"""0001: SPY vs IWM monthly returns scatter with OLS regression line (live data).

Snapshot values cited in lesson 0001 (as of 2026-08-20): 119 monthly obs,
corr 0.862, beta 1.173 (IWM regressed on SPY), cov 0.002314 (decimal returns).
Statistics are computed from DECIMAL returns (matching fetch_data.py and the
lesson's cited numbers); returns are converted to percentage points only for
plotting.
"""
import sys

import matplotlib.pyplot as plt
import numpy as np

sys.path.insert(0, ".")
import course_style as cs

cs.apply_style()

import yfinance as yf

s_dec = yf.download("SPY", start="2016-08-20", end="2026-08-20", interval="1mo", auto_adjust=True)["Close"].iloc[:, 0].dropna()
i_dec = yf.download("IWM", start="2016-08-20", end="2026-08-20", interval="1mo", auto_adjust=True)["Close"].iloc[:, 0].dropna()
common = s_dec.index.intersection(i_dec.index)
s_dec, i_dec = s_dec.loc[common].pct_change().dropna(), i_dec.loc[common].pct_change().dropna()

beta = np.cov(s_dec, i_dec, ddof=1)[0, 1] / np.var(s_dec, ddof=1)
alpha = i_dec.mean() - beta * s_dec.mean()
corr = np.corrcoef(s_dec, i_dec)[0, 1]
cov = np.cov(s_dec, i_dec, ddof=1)[0, 1]

s, i = s_dec * 100, i_dec * 100  # percentage points for plotting only

fig, ax = plt.subplots(figsize=(5, 3.3))
ax.scatter(s, i, s=14, color=cs.SECONDARY, edgecolor="none", alpha=0.85, label="Monthly returns")

xs = np.array([s.min(), s.max()])
ax.plot(xs, alpha + beta * xs, color=cs.PRIMARY, linewidth=2, label="OLS fit")

ax.set_title("SPY vs IWM: monthly returns, 10 years")
ax.set_xlabel("SPY monthly return (percent)")
ax.set_ylabel("IWM monthly return (percent)")
ax.axhline(0, color=cs.TICK, linewidth=0.8)
ax.axvline(0, color=cs.TICK, linewidth=0.8)
ax.text(
    0.03,
    0.97,
    f"beta {beta:.2f}  corr {corr:.2f}  n = {len(s)}",
    transform=ax.transAxes,
    ha="left",
    va="top",
    fontfamily=cs.MONO[0],
    fontsize=8,
    color=cs.INK,
)
ax.legend(loc="upper right")
ax.grid(True, axis="y")
cs.save_fig(fig, "0001-scatter-ols.png")