"""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
import platform
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 wrds_username():
    """The WRDS username, read from the pgpass file."""
    if platform.system() == "Windows":
        pgpass = Path(os.environ["APPDATA"]) / "postgresql" / "pgpass.conf"
    else:
        pgpass = Path.home() / ".pgpass"
    for line in pgpass.read_text().splitlines():
        if line.startswith("wrds-pgdata"):
            return line.split(":")[3]


def wrds_connect():
    """Connect to the WRDS Postgres server with psycopg."""
    import psycopg
    from psycopg.types.numeric import FloatLoader
    conn = psycopg.connect(host="wrds-pgdata.wharton.upenn.edu", port=9737,
                           dbname="wrds", user=wrds_username(), sslmode="require")
    # Return SQL numbers as floats rather than Python Decimals
    conn.adapters.register_loader("numeric", FloatLoader)
    return conn


def wrds_sql(conn, sql, date_cols=("date",)):
    """Run one query and return a DataFrame with date_cols as datetimes."""
    cur = conn.execute(sql)
    df = pd.DataFrame(cur.fetchall(), columns=[c.name for c in cur.description])
    for c in date_cols:
        if c in df:
            df[c] = pd.to_datetime(df[c])
    return df


def load_wrds(name, sql, data_dir=DATA_DIR):
    """Run a query on WRDS and cache the result in data_dir/<name>.parquet."""
    file = data_dir / f"{name}.parquet"
    if not file.exists():
        with wrds_connect() as conn:
            df = wrds_sql(conn, sql)
        data_dir.mkdir(exist_ok=True)
        df.to_parquet(file)
    return pd.read_parquet(file)


def load_wrds_by_year(name, sql, years, data_dir=DATA_DIR):
    """Like load_wrds, with one query per year for tables too large for
    one. `sql` takes the year as {year}."""
    file = data_dir / f"{name}.parquet"
    if not file.exists():
        with wrds_connect() as conn:
            parts = [wrds_sql(conn, sql.format(year=y)) for y in years]
        data_dir.mkdir(exist_ok=True)
        pd.concat(parts, ignore_index=True).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