"""Fit market, sector and style exposures before an earnings event, then reconcile their contributions to the observed stock return.

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=["AAPL", "SPY", "XLK"], start_date="2023-01-01", end_date="2025-12-31", frequency="1d"
        ),
        "daily_factors": qj.ff.get_factors(
            region="america", factor="5_factors", frequency="daily", start_date="2023-01-01", end_date="2025-12-31"
        ),
        "earnings": qj.fmp.get_earnings_surprises(symbol="AAPL"),
    }


def analyze(raw):
    returns = pd.DataFrame(
        {symbol: change(history(raw["prices"], symbol)["price"]) for symbol in ["AAPL", "SPY", "XLK"]}
    )
    data = returns.join(factor_daily(raw["daily_factors"]), how="inner").dropna()
    data["Market"] = data["SPY"] - data["RF"]
    data["Sector"] = data["XLK"] - data["SPY"]
    names = ["Market", "Sector", "SMB", "HML", "RMW", "CMA"]
    events = frame(raw["earnings"])
    chosen = None
    for event in events["date"].sort_values(ascending=False):
        train = data.loc[data.index < event].tail(120)
        after = data.loc[data.index > event].head(5)
        if len(train) == 120 and len(after) == 5:
            chosen = event
            break
    if chosen is None:
        raise ValueError("No event with a complete pre-event and post-event sample")
    design = np.column_stack([np.ones(len(train)), train[names]])
    coefficients, _, rank, _ = np.linalg.lstsq(design, train["AAPL"] - train["RF"], rcond=None)
    if rank != design.shape[1]:
        raise ValueError("Attribution design is rank deficient")
    contribution = after[names].mul(coefficients[1:]).sum()
    rows = [{"Component": name, "Contribution pp": 100 * contribution[name]} for name in names]
    rows.append(
        {"Component": "Cash + intercept", "Contribution pp": 100 * (after["RF"].sum() + len(after) * coefficients[0])}
    )
    residual = after["AAPL"].sum() - sum(row["Contribution pp"] for row in rows) / 100
    rows.append({"Component": "Residual", "Contribution pp": 100 * residual})
    result = pd.DataFrame(rows)
    if not np.isclose(result["Contribution pp"].sum(), 100 * after["AAPL"].sum()):
        raise ValueError("Attribution does not reconcile")
    result["Event date"] = chosen
    result["Training end"] = train.index[-1]
    result["Response end"] = after.index[-1]
    fig, (ax,) = canvas("Post-event attribution · market to residual", "Return contribution · pp")
    start = 0.0
    levels = [start]
    for index, row in result.iterrows():
        value = row["Contribution pp"]
        ax.bar(index, value, bottom=start, color=PINK if row["Component"] == "Residual" else BLUE)
        ax.text(index, start + value, f"{value:+.2f}", ha="center", va="bottom", fontsize=9)
        start += value
        levels.append(start)
    span = max(max(levels) - min(levels), 0.05)
    ax.set_ylim(min(levels) - 0.12 * span, max(levels) + 0.12 * span)
    ax.axhline(0, color="#afc1dc", lw=0.8)
    ax.bar(len(result), start, color="#afc1dc")
    ax.set_xticks(range(len(result) + 1), list(result["Component"]) + ["Observed"], rotation=35, ha="right")
    ax.set_title(f"Earnings {chosen:%Y-%m-%d} · next five sessions", loc="left", color="#e9eef8", fontsize=11)
    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 change(values, periods=1):
    previous = values.shift(periods)
    return pd.Series(
        [ratio(x, y) - 1 if pd.notna(x) and pd.notna(y) else np.nan for x, y in zip(values, previous, strict=True)],
        index=values.index,
    )


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 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 factor_daily(response):
    factors = frame(response, date=None)
    factors.index = pd.to_datetime(factors["Date"], utc=True)
    if factors.index.has_duplicates:
        raise ValueError("Duplicate factor dates")
    return factors[["Mkt-RF", "SMB", "HML", "RMW", "CMA", "RF"]].apply(pd.to_numeric) / 100


OUTPUT = "sdk_event_attribution.png"

if __name__ == "__main__":
    run()
