"""Locate the largest observed drawdown, attribute the loss to each holding and compare the period with VIX.

Research example. Live observations can change.
"""

from datetime import UTC, datetime
from decimal import Decimal
from textwrap import fill

import matplotlib.pyplot as plt
import numpy as np
import pandas as pd

AS_OF = datetime.now(UTC).date()
NAVY, BLUE, PINK = "#132244", "#8fb9ff", "#e7a2cc"


def fetch(qj):
    return {
        "prices": qj.eod.get_historical_prices(
            symbols=["SPY", "TLT", "GLD", "DBC", "UUP"], start_date="2020-01-01", end_date="2025-12-31", frequency="1d"
        ),
        "vix_history": qj.fred.get_fred_data_series_by_id(search_id="VIXCLS", start="2020-01-01", end="2025-12-31"),
    }


def analyze(raw):
    weights = notebook_weights()
    prices = pd.DataFrame({symbol: history(raw["prices"], symbol)["price"] for symbol in weights}).dropna()
    if len(prices) < 252:
        raise ValueError("Need at least one year of common prices")
    units = {symbol: weights[symbol] / decimal(prices[symbol].iloc[0]) for symbol in weights}
    values = prices.apply(lambda column: column * units[column.name])
    nav = values.sum(axis=1)
    drawdown = pd.Series(
        [ratio(value, peak) - 1 for value, peak in zip(nav, nav.cummax(), strict=True)], index=nav.index
    )
    trough = drawdown.idxmin()
    peak = nav.loc[:trough].idxmax()
    if peak == trough:
        raise ValueError("No drawdown in the observed sample")
    contribution = values.loc[trough] - values.loc[peak]
    result = pd.DataFrame(
        {"Symbol": contribution.index, "Contribution pp": [100 * ratio(value, nav.loc[peak]) for value in contribution]}
    )
    gross_losses = -result.loc[result["Contribution pp"] < 0, "Contribution pp"].sum()
    if gross_losses <= 0:
        raise ValueError("No negative contribution")
    result = result.sort_values("Contribution pp")
    shares = (-result["Contribution pp"].clip(upper=0) / gross_losses).cumsum()
    result["In 80% loss group"] = shares.shift(fill_value=0) < 0.8
    result["Gross loss share %"] = 100 * (-result["Contribution pp"].clip(upper=0) / gross_losses)
    if not np.isclose(result["Contribution pp"].sum(), 100 * drawdown.loc[trough]):
        raise ValueError("Attribution does not reconcile")
    vix = macro(raw["vix_history"], "VIXCLS").reindex(drawdown.index)
    if vix.loc[peak:trough].dropna().empty:
        raise ValueError("No VIX observations in the drawdown window")
    fig, (top, bottom) = canvas(
        "Portfolio concentration · explain the drawdown", "Contribution to drawdown · pp", panels=2
    )
    top.barh(
        result["Symbol"],
        result["Contribution pp"],
        color=[PINK if flag else BLUE for flag in result["In 80% loss group"]],
    )
    top.axvline(0, color="#afc1dc", lw=0.8)
    top.set_xlabel(f"Peak {peak:%Y-%m-%d} → trough {trough:%Y-%m-%d}")
    bottom.plot(drawdown.index, 100 * drawdown, label="Portfolio drawdown %", color=BLUE)
    bottom.axvspan(peak, trough, color=PINK, alpha=0.12)
    right = bottom.twinx()
    right.plot(vix.index, vix, alpha=0.4, color=PINK)
    right.set_ylabel("VIX level", color=PINK)
    bottom.set_ylabel("Drawdown %")
    result["Peak"] = peak
    result["Trough"] = trough
    return result, fig


def payload(response):
    if hasattr(response, "to_dict"):
        response = response.to_dict()
    while isinstance(response, dict):
        if response.get("error") or response.get("error_code") or response.get("success") is False:
            raise ValueError("Provider request failed")
        if "data" in response:
            response = response["data"]
        elif "value" in response:
            response = response["value"]
        else:
            break
    return response


def frame(response, date="date"):
    rows = payload(response)
    if not isinstance(rows, list) or not rows:
        raise ValueError("No provider observations returned")
    result = pd.DataFrame(rows)
    if date:
        result[date] = pd.to_datetime(result[date], utc=True)
        result = result.sort_values(date)
    return result


def decimal(value):
    if pd.isna(value):
        raise ValueError("Required monetary value or quantity is missing")
    return Decimal(str(value))


def ratio(numerator, denominator):
    denominator = decimal(denominator)
    if denominator == 0:
        raise ValueError("Required denominator is zero")
    return float(decimal(numerator) / denominator)


def history(response, symbol, adjusted=True):
    values = payload(response)
    if not isinstance(values, dict) or symbol not in values:
        raise ValueError(f"No price coverage for {symbol}")
    result = frame(values[symbol]).set_index("date")
    field = "adjusted_close" if adjusted else "close"
    result["price"] = result[field].map(decimal)
    result["volume"] = result["volume"].map(decimal)
    if result.index.has_duplicates:
        raise ValueError("Duplicate trading dates")
    return result


def macro(response, column):
    values = frame(response).set_index("date")[column]
    values = pd.to_numeric(values, errors="coerce").dropna()
    if values.empty:
        raise ValueError(f"No numeric observations for {column}")
    return values


def canvas(title, ylabel, panels=1):
    plt.rcParams.update(
        {
            "figure.facecolor": NAVY,
            "axes.facecolor": NAVY,
            "savefig.facecolor": NAVY,
            "text.color": "#e9eef8",
            "axes.labelcolor": "#afc1dc",
            "xtick.color": "#afc1dc",
            "ytick.color": "#afc1dc",
            "axes.edgecolor": "#34486b",
            "grid.color": "#34486b",
            "font.size": 10,
            "axes.prop_cycle": plt.cycler(color=[BLUE, PINK, "#ba9dea"]),
        }
    )
    fig, axes = plt.subplots(panels, 1, figsize=(9, 6.1), squeeze=False, layout="constrained")
    fig.suptitle(fill(title, width=58), fontsize=17, x=0.04, ha="left", weight="medium")
    for ax in axes.flat:
        ax.grid(axis="y", alpha=0.45)
        ax.set_axisbelow(True)
        ax.spines[["top", "right"]].set_visible(False)
    axes.flat[0].set_ylabel(ylabel)
    return fig, list(axes.flat)


def run():
    from quantjourney.sdk import QuantJourneyAPI

    qj = QuantJourneyAPI.from_env()
    raw = fetch(qj)
    table, fig = analyze(raw)
    print(table.to_string(index=False))
    fig.savefig(OUTPUT, dpi=180)
    plt.show()


def notebook_weights():
    return {
        "SPY": Decimal("0.42"),
        "TLT": Decimal("0.22"),
        "GLD": Decimal("0.14"),
        "DBC": Decimal("0.12"),
        "UUP": Decimal("0.10"),
    }


OUTPUT = "sdk_drawdown_attribution.png"

if __name__ == "__main__":
    run()