"""Reproduce the stale-price essay's descriptive sensitivity experiment.

Inputs are public observations captured on 2026-10-30; see README.md.
Run with Python 3, numpy, pandas and matplotlib. No network calls are made.
"""

from pathlib import Path
import hashlib
import json
import re

import matplotlib
matplotlib.use("Agg")
import matplotlib.pyplot as plt
import numpy as np
import pandas as pd

ROOT = Path(__file__).resolve().parent


def fit_ar1(series):
    pairs = pd.concat([series.rename("current"), series.shift().rename("lag")], axis=1).dropna()
    design = np.column_stack([np.ones(len(pairs)), pairs["lag"].values])
    intercept, slope = np.linalg.lstsq(design, pairs["current"].values, rcond=None)[0]
    return {"intercept": float(intercept), "alpha": float(slope), "n": len(pairs)}


def adjusted_returns(current, lag, alpha):
    # The inverse one-lag filter is a model assumption, not identification of stale marks.
    result = (current - alpha * lag) / (1 - alpha)
    # A finite sample has a boundary effect. Match means exactly to isolate risk inputs.
    return result - result.mean() + current.mean()


cdli = pd.read_csv(ROOT / "cdli-public-chart.tsv", sep="\t")
cdli.index = pd.PeriodIndex(pd.to_datetime(cdli["Date"]), freq="Q")
expected_dates = pd.period_range("2004Q3", "2026Q2", freq="Q")
assert cdli.index.equals(expected_dates), "Missing or repeated CDLI quarters"
levels = cdli["Cliffwater Direct Lending Index"].astype(float)
assert (levels > 0).all()
reported = levels.pct_change().dropna()
fit = fit_ar1(reported)
alpha_hat = fit["alpha"]
assert 0 < alpha_hat < 1

raw_french = (ROOT / "fama-french-factors-source.csv").read_text()
monthly_rows = [line.split(",") for line in raw_french.splitlines() if re.match(r"^\d{6},", line)]
french = pd.DataFrame(monthly_rows, columns=["date", "Mkt-RF", "SMB", "HML", "RF"])
french.index = pd.PeriodIndex(french["date"], freq="M")
french = french[["Mkt-RF", "RF"]].astype(float) / 100
french = french.loc["2004-10":"2026-09"]
assert len(french) == 264 and np.isfinite(french.values).all()
assert (french.values > -0.99).all(), "Missing-data sentinel or invalid return"
monthly = pd.DataFrame({"equity": french["Mkt-RF"] + french["RF"], "rf": french["RF"]})
quarter_labels = monthly.index.asfreq("Q")
assert monthly.groupby(quarter_labels).size().eq(3).all()
quarterly = (1 + monthly).groupby(quarter_labels).prod() - 1

sample = pd.DataFrame({"reported": reported, "lag": reported.shift()}).dropna().join(quarterly)
assert len(sample) == 86 and sample.notna().all().all()
sample["adjusted"] = adjusted_returns(sample["reported"], sample["lag"], alpha_hat)
equity_sigma = sample["equity"].std(ddof=1)
market_excess = sample["equity"] - sample["rf"]


def metrics(series):
    excess = series - sample["rf"]
    annual_vol = 2 * series.std(ddof=1)
    weight_credit = equity_sigma / (equity_sigma + series.std(ddof=1))
    covariance = np.cov(np.stack([series, sample["equity"]]), ddof=1)
    weights = np.array([weight_credit, 1 - weight_credit])
    risk_contributions = weights * (covariance @ weights)
    # In this two-asset example, inverse-volatility weights have equal risk contributions.
    assert np.isclose(risk_contributions[0], risk_contributions[1], rtol=1e-10)
    return {
        "annualised_volatility_sqrt4": float(annual_vol),
        "conventional_annualised_sharpe": float(2 * excess.mean() / excess.std(ddof=1)),
        "equity_beta_excess_returns_with_intercept": float(
            np.cov(excess, market_excess, ddof=1)[0, 1] / market_excess.var(ddof=1)
        ),
        "equity_correlation_total_returns": float(series.corr(sample["equity"])),
        "arithmetic_mean_return_times4": float(4 * series.mean()),
        "quarterly_excess_mean": float(excess.mean()),
        "quarterly_excess_std": float(excess.std(ddof=1)),
        "credit_weight_equal_risk_contribution": float(weight_credit),
        "credit_weight_after_30pct_cap": float(min(weight_credit, 0.30)),
        "vol_target_position_relative_to_reported": float(sample["reported"].std(ddof=1) / series.std(ddof=1)),
        "lag1_autocorrelation": float(series.autocorr()),
    }


baseline = metrics(sample["reported"])
adjusted = metrics(sample["adjusted"])
assert np.isclose(sample["reported"].mean(), sample["adjusted"].mean(), atol=1e-14)
assert np.isclose(
    baseline["quarterly_excess_mean"], adjusted["quarterly_excess_mean"], atol=1e-14
)

grid = []
for alpha in np.linspace(0, 0.45, 181):
    row = metrics(adjusted_returns(sample["reported"], sample["lag"], float(alpha)))
    grid.append({"alpha": float(alpha), **row})

results = {
    "captured_on": "2026-10-30",
    "sample_start": str(sample.index[0]),
    "sample_end": str(sample.index[-1]),
    "quarters": len(sample),
    "ar1_fit_on_87_returns": fit,
    "ar1_fit_through_index_launch": fit_ar1(reported.loc[:"2015Q3"]),
    "ar1_fit_after_index_launch": fit_ar1(reported.loc["2015Q4":]),
    "reported": baseline,
    "adjusted": adjusted,
    "sensitivity": [
        {"alpha": a, **metrics(adjusted_returns(sample["reported"], sample["lag"], a))}
        for a in [0.0, 0.15, alpha_hat, 0.45]
    ],
    "source_sha256": {
        name: hashlib.sha256((ROOT / name).read_bytes()).hexdigest()
        for name in ["cdli-public-chart.tsv", "fama-french-factors-source.csv"]
    },
}
post_alpha = results["ar1_fit_after_index_launch"]["alpha"]
results["post_launch_alpha_applied_to_full_matched_sample"] = {
    "alpha": post_alpha,
    **metrics(adjusted_returns(sample["reported"], sample["lag"], post_alpha)),
}


def distributed_lag_beta(lags=1, include_lead=False):
    market = quarterly["equity"] - quarterly["rf"]
    columns = {"current": market}
    columns.update({f"lag{k}": market.shift(k) for k in range(1, lags + 1)})
    if include_lead:
        columns["lead1"] = market.shift(-1)
    x = pd.DataFrame(columns).reindex(sample.index)
    data = x.join((sample["reported"] - sample["rf"]).rename("response")).dropna()
    design = np.column_stack([np.ones(len(data)), data[list(columns)].values])
    response = data["response"].values
    coefficients = np.linalg.lstsq(design, response, rcond=None)[0]
    residuals = response - design @ coefficients
    # Newey-West covariance, two quarterly lags, with n/(n-k) finite-sample factor.
    score = design * residuals[:, None]
    meat = score.T @ score
    for lag in [1, 2]:
        cross = score[lag:].T @ score[:-lag]
        meat += (1 - lag / 3) * (cross + cross.T)
    bread = np.linalg.inv(design.T @ design)
    covariance = bread @ meat @ bread * len(data) / (len(data) - design.shape[1])
    contrast = np.r_[0.0, np.ones(len(columns))]
    return {
        "n": len(data),
        "start": str(data.index[0]),
        "end": str(data.index[-1]),
        "coefficients": dict(zip(["intercept", *columns], map(float, coefficients))),
        "sum_beta": float(coefficients[1:].sum()),
        "sum_beta_hac2_standard_error": float(np.sqrt(contrast @ covariance @ contrast)),
    }


results["dimson_style_checks"] = {
    "current_plus_one_lag": distributed_lag_beta(),
    "current_plus_two_lags": distributed_lag_beta(lags=2),
    "one_lead_current_one_lag": distributed_lag_beta(include_lead=True),
}
assert results["dimson_style_checks"]["current_plus_one_lag"]["n"] == len(sample)
assert results["dimson_style_checks"]["one_lead_current_one_lag"]["n"] == len(sample)
reported_excess = sample["reported"] - sample["rf"]
results["illustrative_flat_fee_drag"] = {
    "annualised_mean_excess_return_before_drag": float(4 * reported_excess.mean()),
    "annual_drag_1_5pct_sharpe": float(2 * (reported_excess - 0.015 / 4).mean() / reported_excess.std(ddof=1)),
    "annual_drag_2pct_sharpe": float(2 * (reported_excess - 0.02 / 4).mean() / reported_excess.std(ddof=1)),
}
(ROOT / "results.json").write_text(json.dumps(results, indent=2) + "\n")
sample.to_csv(ROOT / "matched-quarterly-observations.csv", index_label="quarter", float_format="%.12f")

plt.rcParams.update({"font.family": "DejaVu Sans", "font.size": 10, "axes.spines.top": False,
                     "axes.spines.right": False, "svg.fonttype": "path"})
fig, axes = plt.subplots(1, 2, figsize=(11.8, 5.8))
fig.subplots_adjust(top=0.74, bottom=0.23, left=0.075, right=0.97, wspace=0.30)
navy, orange, teal, muted = "#203b53", "#c05e27", "#197b78", "#596875"
fig.text(0.075, 0.94, "Position size depends on the volatility you choose", fontsize=19, weight="bold", color=navy)
fig.text(0.075, 0.885, "One-lag unsmoothing sensitivity · CDLI · 2005 Q1–2026 Q2 · 86 matched quarters", fontsize=11, color=muted)
fig.text(0.075, 0.835, "Teal: post-launch coefficient (0.107). Orange: full-sample coefficient (0.295).", fontsize=10, color=muted)

alphas = np.array([row["alpha"] for row in grid])
vols = 100 * np.array([row["annualised_volatility_sqrt4"] for row in grid])
positions = 100 * np.array([row["vol_target_position_relative_to_reported"] for row in grid])
post_metrics = results["post_launch_alpha_applied_to_full_matched_sample"]
specs = [
    (axes[0], vols, "Annualised volatility (%)", "annualised_volatility_sqrt4"),
    (axes[1], positions, "Position size (reported σ = 100)", "vol_target_position_relative_to_reported"),
]
for ax, values, title, metric in specs:
    start, fitted, post_value = [100 * row[metric] for row in [baseline, adjusted, post_metrics]]
    ax.plot(alphas, values, color=navy, linewidth=2.2)
    ax.axvline(alpha_hat, color=orange, linestyle="--", linewidth=1.2, alpha=0.8)
    ax.axvline(post_alpha, color=teal, linestyle=":", linewidth=1.2, alpha=0.8)
    ax.scatter([0, post_alpha, alpha_hat], [start, post_value, fitted], color=[navy, teal, orange], s=42, zorder=4)
    ax.set_title(title, loc="left", fontsize=12, weight="bold", pad=12)
    ax.set_xlabel("Assumed smoothing coefficient α")
    ax.set_xlim(-0.015, 0.465)
    ax.set_xticks([0, 0.15, 0.30, 0.45])
    ax.grid(axis="y", alpha=0.16)
    ax.tick_params(axis="both", colors=muted)
    for spine in ax.spines.values():
        spine.set_color("#d4dbe0")
    for coefficient, value, color, offset in [(0, start, navy, (8, 11)), (post_alpha, post_value, teal, (8, 11)), (alpha_hat, fitted, orange, (8, 9))]:
        label = f"{value:.2f}%" if ax is axes[0] else f"{value:.1f}"
        ax.annotate(label, (coefficient, value), xytext=offset, textcoords="offset points", color=color, weight="bold")

axes[1].set_ylim(50, 107)
axes[1].set_yticks([50, 60, 70, 80, 90, 100])
fig.text(0.075, 0.12, "At fixed standalone volatility budget: full-sample adjustment cuts size by 26.2%; post-launch adjustment by 8.3%.", fontsize=9.5, color=navy)
fig.text(0.075, 0.075, "Sources: Cliffwater public CDLI chart; Kenneth French Data Library. Calculations: Rohan Rathod.", fontsize=8.5, color=muted)
fig.text(0.075, 0.04, "Model sensitivity, not recovered market prices. CDLI is unlevered and gross of fees; pre-launch history is reconstructed.", fontsize=8.5, color=muted)
fig.savefig(ROOT / "stale-price-sensitivity.png", dpi=200, facecolor="white")
fig.savefig(ROOT / "stale-price-sensitivity.svg", facecolor="white")
plt.close(fig)
print(json.dumps(results, indent=2))