"""Corporate Actions — Adjusting for splits and dividends in historical price series."""

import matplotlib.pyplot as plt
import numpy as np
import polars as pl
from ml4t.data.adjustments import apply_corporate_actions

from data import load_us_equities
from utils.style import COLORS

# Production defaults — Papermill injects overrides for CI

raw_wiki = load_us_equities()
print(f"Wiki Prices loaded: {len(raw_wiki):,} rows")

dividend_col = "ex_dividend" if "ex_dividend" in raw_wiki.columns else "ex-dividend"
aapl = raw_wiki.filter(pl.col("symbol") == "AAPL").sort("timestamp")

print(f"AAPL records: {len(aapl):,}  ({aapl['timestamp'].min()} to {aapl['timestamp'].max()})")

splits = aapl.filter(pl.col("split_ratio") != 1.0)
splits.select(["timestamp", "close", "split_ratio", "adj_close"])

dividends = aapl.filter(pl.col(dividend_col) > 0)
dividends.select(["timestamp", "close", dividend_col, "adj_close"]).head(10)

dates_arr = aapl["timestamp"].to_numpy()
raw_close = aapl["close"].to_numpy()
adj_close = aapl["adj_close"].to_numpy()
split_dates = splits["timestamp"].to_numpy()
split_ratios = splits["split_ratio"].to_numpy()

raw_returns = np.diff(raw_close) / raw_close[:-1]
adj_returns = np.diff(adj_close) / adj_close[:-1]
raw_cumret = np.cumprod(1 + raw_returns)
adj_cumret = np.cumprod(1 + adj_returns)

fig, ax = plt.subplots(figsize=(12, 6), layout="tight")
ax.plot(
    dates_arr[1:], adj_cumret, label="Adjusted prices (correct)", color=COLORS["blue"], linewidth=2
)
ax.plot(
    dates_arr[1:], raw_cumret, label="Raw prices (wrong)", color=COLORS["negative"], linewidth=2
)

for d, ratio in zip(split_dates, split_ratios, strict=False):
    ax.axvline(d, color="gray", linestyle="--", alpha=0.4, linewidth=1)
    idx = np.searchsorted(dates_arr[1:], d)
    if idx < len(adj_cumret):
        ax.annotate(
            f"{int(ratio)}:1 split",
            xy=(d, adj_cumret[idx]),
            xytext=(10, 20),
            textcoords="offset points",
            fontsize=9,
            color="gray",
            arrowprops=dict(arrowstyle="-", color="gray", alpha=0.5),
        )

ax.annotate(
    f"Raw: {raw_cumret[-1]:.0f}x\nAdjusted: {adj_cumret[-1]:.0f}x\n({adj_cumret[-1] / raw_cumret[-1]:.0f}x difference)",
    xy=(dates_arr[-1], (raw_cumret[-1] + adj_cumret[-1]) / 2),
    xytext=(-120, 0),
    textcoords="offset points",
    fontsize=11,
    fontweight="bold",
    bbox=dict(boxstyle="round,pad=0.3", facecolor="wheat", alpha=0.8),
    ha="right",
)
ax.annotate(
    "Each split looks like a crash\nin raw prices, causing\ncumulative returns to diverge",
    xy=(split_dates[2], raw_cumret[np.searchsorted(dates_arr[1:], split_dates[2])]),
    xytext=(50, -30),
    textcoords="offset points",
    fontsize=9,
    style="italic",
    bbox=dict(boxstyle="round,pad=0.2", facecolor="white", alpha=0.8, edgecolor="gray"),
    arrowprops=dict(arrowstyle="->", color="gray"),
)

ax.set_ylabel("Cumulative return (start = 1)")
ax.set_xlabel("Date")
ax.set_title("Raw vs adjusted cumulative return — AAPL 1980–2018")
ax.legend(loc="upper left", fontsize=10)
ax.set_yscale("log")
ax.set_ylim(0.5, adj_cumret[-1] * 1.5)
ax.yaxis.set_major_formatter(plt.FuncFormatter(lambda x, _: f"{x:.0f}x" if x >= 1 else f"{x:.1f}x"))
ax.grid(True, alpha=0.3)
plt.show()

# Persist Figure 2.3 artifact for book script (Hard Rule 15).
import polars as _pl

from utils.paths import get_chapter_dir as _get_chapter_dir

_ARTIFACTS_DIR = _get_chapter_dir(2) / "output" / "book_figure_artifacts"
_ARTIFACTS_DIR.mkdir(parents=True, exist_ok=True)
_pl.DataFrame(
    {
        "timestamp": dates_arr[1:],
        "adj_cumret": adj_cumret,
        "raw_cumret": raw_cumret,
    }
).write_parquet(_ARTIFACTS_DIR / "figure_2_3_corporate_actions_curves.parquet")
_pl.DataFrame(
    {
        "split_date": split_dates,
        "split_ratio": split_ratios,
    }
).write_parquet(_ARTIFACTS_DIR / "figure_2_3_corporate_actions_splits.parquet")
print(
    f"Persisted figure_2_3_corporate_actions: {len(adj_cumret)} dates, "
    f"{len(split_dates)} splits, final raw={raw_cumret[-1]:.1f}x adj={adj_cumret[-1]:.1f}x"
)

print(f"Final cumulative return — raw prices:      {raw_cumret[-1]:.1f}x")
print(f"Final cumulative return — adjusted prices: {adj_cumret[-1]:.1f}x")
print(f"Ratio: {adj_cumret[-1] / raw_cumret[-1]:.0f}x difference")
print(
    f"Raw prices understate the cumulative return by {(adj_cumret[-1] / raw_cumret[-1] - 1) * 100:.0f}%"
)

help(apply_corporate_actions)

adjusted = apply_corporate_actions(
    aapl.rename({"timestamp": "date"}),
    split_col="split_ratio",
    dividend_col=dividend_col,
    price_cols=["open", "high", "low", "close"],
    volume_col="volume",
).rename({"date": "timestamp"})

print("Adjusted columns:", [c for c in adjusted.columns if c.startswith("adj_")])

from datetime import date as _date

example_dates = [
    _date(1980, 12, 12),
    _date(1987, 6, 16),
    _date(2000, 6, 21),
    _date(2014, 6, 9),
    _date(2018, 3, 27),
]

comparison = (
    adjusted.filter(pl.col("timestamp").is_in(example_dates))
    .select(
        pl.col("timestamp"),
        pl.col("close").alias("raw_close"),
        pl.col("adj_close").alias("our_adj_close"),
        pl.col("split_ratio"),
    )
    .join(
        aapl.select(
            pl.col("timestamp"),
            pl.col("adj_close").alias("quandl_adj_close"),
        ),
        on="timestamp",
    )
    .sort("timestamp")
)

comparison

our_adj = adjusted["adj_close"].to_numpy()
quandl_adj = aapl["adj_close"].to_numpy()
dates = adjusted["timestamp"].to_numpy()

absolute_diff = np.abs(our_adj - quandl_adj)
relative_diff = absolute_diff / quandl_adj * 100
tolerance_pct = 0.05  # 0.05 %

print(f"Total comparisons:        {len(our_adj):,}")
print(f"Max absolute difference:  ${absolute_diff.max():.6f}")
print(f"Max relative difference:  {relative_diff.max():.4f}%")
print(f"Mean relative difference: {relative_diff.mean():.6f}%")
print(f"Median relative difference: {np.median(relative_diff):.6f}%")
print()
if np.allclose(our_adj, quandl_adj, rtol=tolerance_pct / 100):
    print(f"VALIDATION PASSED — within {tolerance_pct:.4f}% relative tolerance.")
else:
    n_out = int((relative_diff > tolerance_pct).sum())
    print(f"VALIDATION FAILED — {n_out} rows outside {tolerance_pct:.4f}% tolerance.")

fig, axes = plt.subplots(2, 1, figsize=(14, 8), layout="tight")

ax1 = axes[0]
ax1.plot(dates, our_adj, label="Our adjustment", alpha=0.85)
ax1.plot(dates, quandl_adj, label="Quandl adjustment", alpha=0.85, linestyle="--")
ax1.set_ylabel("Adjusted price ($)")
ax1.set_title("Local vs Quandl pre-calculated adjusted close")
ax1.set_yscale("log")
ax1.legend()

ax2 = axes[1]
ax2.plot(dates, relative_diff, color=COLORS["blue"], linewidth=1.2, alpha=0.9)
ax2.axhline(
    tolerance_pct,
    color="red",
    linestyle="--",
    linewidth=2,
    label=f"Tolerance ({tolerance_pct:.4f}%)",
)
ax2.set_ylabel("Relative difference (%)")
ax2.set_xlabel("Date")
ax2.set_title("Relative difference between local and Quandl series")
ax2.set_ylim(0, max(relative_diff.max() * 1.1, tolerance_pct * 2))
ax2.legend()

plt.show()

def demonstrate_split_adjustment(
    price_before: float = 270.0, shares_before: int = 100, ratio: float = 3.0
) -> None:
    """Print a worked example of backward split adjustment."""
    value_before = price_before * shares_before
    price_after = price_before / ratio
    shares_after = shares_before * ratio
    value_after = price_after * shares_after

    print(
        f"Before {int(ratio)}:1 split — price ${price_before:.2f}, shares {shares_before:d}, value ${value_before:,.2f}"
    )
    print(
        f"After  {int(ratio)}:1 split — price ${price_after:.2f}, shares {int(shares_after):d}, value ${value_after:,.2f}"
    )
    print(
        f"Backward-adjusted historical price: ${price_before:.2f} / {ratio:.0f} = ${price_before / ratio:.2f}"
    )
    print(f"Adjusted series is now continuous: ${price_before / ratio:.2f} -> ${price_after:.2f}")

demonstrate_split_adjustment()

def demonstrate_dividend_adjustment(price_before_ex: float = 100.0, dividend: float = 5.0) -> None:
    """Print a worked example of backward dividend adjustment."""
    price_on_ex = price_before_ex - dividend
    raw_return = (price_on_ex - price_before_ex) / price_before_ex
    multiplier = price_on_ex / price_before_ex
    adjusted_before = price_before_ex * multiplier
    adj_return = (price_on_ex - adjusted_before) / adjusted_before

    print(
        f"Day before ex: ${price_before_ex:.2f}    Dividend: ${dividend:.2f}    Ex-date: ${price_on_ex:.2f}"
    )
    print(f"Raw return on ex-date:      {raw_return:.1%}")
    print(f"Backward multiplier: {price_on_ex:.2f} / {price_before_ex:.2f} = {multiplier:.4f}")
    print(
        f"Adjusted pre-ex price: ${price_before_ex:.2f} × {multiplier:.4f} = ${adjusted_before:.2f}"
    )
    print(f"Adjusted return on ex-date: {adj_return:.1%}")

demonstrate_dividend_adjustment()

universe = pl.DataFrame(
    {
        "symbol": ["SPY", "QQQ", "TLT", "GLD", "EFA"],
        "price_return": [0.10, 0.15, -0.02, 0.08, 0.07],
        "dividend_yield": [0.015, 0.005, 0.025, 0.00, 0.025],
    }
).with_columns(total_return=pl.col("price_return") + pl.col("dividend_yield"))

ranked = universe.with_columns(
    rank_price=pl.col("price_return").rank(descending=True, method="ordinal"),
    rank_total=pl.col("total_return").rank(descending=True, method="ordinal"),
).with_columns(
    # Cast to a signed dtype before subtracting: ranks are u32, so a drop in
    # rank (e.g. GLD 3 -> 4) would underflow to a huge number otherwise.
    rank_change=pl.col("rank_price").cast(pl.Int32) - pl.col("rank_total").cast(pl.Int32)
)

ranked.select(
    [
        "symbol",
        "price_return",
        "dividend_yield",
        "total_return",
        "rank_price",
        "rank_total",
        "rank_change",
    ]
)