"""Shared tools for the From Papers to Portfolios notebooks.

Quarto's freeze and the Jupyter cache do not track this file. After
editing it, render the notebooks that import it with --cache-refresh.
"""

import os
from pathlib import Path

import matplotlib as mpl
import numpy as np
import pandas as pd
from scipy.stats import t as t_dist

# Make SVG figures byte-identical across reruns, so the frozen outputs
# change in git only when a figure does
os.environ.setdefault("SOURCE_DATE_EPOCH", "0")
mpl.rcParams["svg.hashsalt"] = "from-papers-to-portfolios"

DATA_DIR = Path(__file__).resolve().parent.parent / "_data"


def load_wrds(name, sql, data_dir=DATA_DIR):
    file = data_dir / f"{name}.parquet"
    if not file.exists():
        import wrds
        conn = wrds.Connection()
        df = conn.raw_sql(sql, date_cols=["date"])
        conn.close()
        data_dir.mkdir(exist_ok=True)
        df.to_parquet(file)
    return pd.read_parquet(file)


def load_jkp(name, url, data_dir=DATA_DIR):
    file = data_dir / f"{name}.parquet"
    if not file.exists():
        df = pd.read_csv(url, usecols=["name", "date", "ret"], parse_dates=["date"])
        data_dir.mkdir(exist_ok=True)
        df.to_parquet(file)
    return pd.read_parquet(file)


def sharpe(r):
    return np.sqrt(12) * r.mean() / r.std()


def proj_simplex(v):
    u = np.sort(v)[::-1]
    cssv = np.cumsum(u)
    rho = np.nonzero(u * np.arange(1, len(v) + 1) > (cssv - 1))[0][-1]
    theta = (cssv[rho] - 1) / (rho + 1)
    w = np.maximum(v - theta, 0.0)
    return w / w.sum()


def minvar_longonly_pgd(S, tol=1e-7, max_iters=100000, ridge=1e-8):
    S = 0.5 * (S + S.T) + ridge * np.eye(S.shape[0])
    n = S.shape[0]
    w = np.full(n, 1 / n)
    eta = 1.0 / np.linalg.eigvalsh(S).max()

    # Accelerated projected gradient (FISTA), restarted when the objective rises
    y, t, f = w, 1.0, w @ S @ w
    for _ in range(max_iters):
        w_new = proj_simplex(y - eta * (S @ y))
        if np.linalg.norm(w_new - w) < tol:
            return w_new
        f_new = w_new @ S @ w_new
        if f_new > f:
            y, t = w_new, 1.0
        else:
            t_new = (1 + np.sqrt(1 + 4 * t * t)) / 2
            y = w_new + (t - 1) / t_new * (w_new - w)
            t = t_new
        w, f = w_new, f_new
    return w


def cov_ledoit_wolf(X):
    X = X.to_numpy(dtype=float)
    X = X - X.mean(axis=0)
    T, N = X.shape
    S = X.T @ X / T
    var = np.diag(S)
    s = np.sqrt(var)

    # Constant-correlation target
    r_bar = ((S / np.outer(s, s)).sum() - N) / (N * (N - 1))
    F = r_bar * np.outer(s, s)
    np.fill_diagonal(F, var)

    # Estimated optimal shrinkage intensity
    pi_mat = (X ** 2).T @ (X ** 2) / T - S ** 2
    theta = (X ** 3).T @ X / T - var[:, None] * S
    np.fill_diagonal(theta, 0)
    rho = np.trace(pi_mat) + r_bar * ((s[None, :] / s[:, None]) * theta).sum()
    gamma = ((S - F) ** 2).sum()
    delta = min(1.0, max(0.0, (pi_mat.sum() - rho) / gamma / T))
    return delta * F + (1 - delta) * S, delta


def summarize(r, m, rf):
    ex, mex = r - rf, m - rf
    X = np.column_stack([np.ones(len(mex)), mex])
    coef, *_ = np.linalg.lstsq(X, ex, rcond=None)
    resid = ex - X @ coef
    se = np.sqrt(resid @ resid / (len(ex) - 2) * np.linalg.inv(X.T @ X)[0, 0])
    wealth = (1 + r).cumprod()
    return pd.Series({
        "Mean (%)": 1200 * r.mean(),
        "Volatility (%)": 100 * np.sqrt(12) * r.std(),
        "Sharpe ratio": np.sqrt(12) * ex.mean() / ex.std(),
        "Beta": coef[1],
        "Alpha (%)": 1200 * coef[0],
        "Alpha p-value": 2 * t_dist.sf(abs(coef[0] / se), len(ex) - 2) if se > 1e-12 else np.nan,
        "Max drawdown (%)": 100 * (wealth / wealth.cummax() - 1).min(),
    })


def factor_regression(r, rf, F):
    ex = (r - rf).to_numpy()
    X = np.column_stack([np.ones(len(ex)), F.to_numpy()])
    coef, *_ = np.linalg.lstsq(X, ex, rcond=None)
    resid = ex - X @ coef
    se = np.sqrt(resid @ resid / (len(ex) - X.shape[1]) * np.diag(np.linalg.inv(X.T @ X)))
    r2 = 1 - resid @ resid / ((ex - ex.mean()) @ (ex - ex.mean()))
    p = 2 * t_dist.sf(np.abs(coef / se), len(ex) - X.shape[1])
    return coef, p, r2