"""Replication entry point for F6-conf-02. Run from workspace root:
   ~/wp29/venv/bin/python out/analysis.py
"""
import json, os
import numpy as np, pandas as pd
import matplotlib
matplotlib.use("Agg")
import matplotlib.pyplot as plt
from scipy import stats

SEED = None  # no randomness used
os.makedirs("out", exist_ok=True)
raw = pd.read_csv("data/gc5y_daily.csv")
n_rows = len(raw)
d = raw.copy()
d["date"] = pd.to_datetime(d["date"], errors="coerce")
issues = []

# ---- audit ----
if d["date"].isna().any():
    issues.append({"issue": f"{int(d['date'].isna().sum())} unparseable dates", "handling": "dropped"})
if d["yield_5y"].isna().any():
    issues.append({"issue": f"{int(d['yield_5y'].isna().sum())} missing yields", "handling": "dropped before differencing"})
if d["date"].duplicated().any():
    issues.append({"issue": "duplicated dates", "handling": "kept first"})
d = d.dropna(subset=["date", "yield_5y"]).drop_duplicates("date").sort_values("date").reset_index(drop=True)
d["chg_bp"] = 100 * d["yield_5y"].diff()

n_treated = int((d["post_cut"] == 1).sum())
last = d["date"].max()
if n_treated == 0:
    issues.append({
        "issue": (f"File covers {d['date'].min().date()} to {last.date()} only (README retrieval end_date=2019-12-31); "
                  "zero rows have post_cut=1 and no observation exists on/after the 2020-03-04 cut, "
                  "so the post-cut window is empty and the 120-day pre-window (defined relative to the first flagged day) is undefined."),
        "handling": "Estimand declared not estimable (premise_invalid); effect, interval, model RMSE and n_analysed set to null. No proxy window substituted."})
mism = int(((d["date"] >= "2020-03-04").astype(int) != d["post_cut"]).sum())
if mism:
    issues.append({"issue": f"{mism} rows where post_cut disagrees with date>=2020-03-04",
                   "handling": "post_cut used exactly as provided per task rules"})
gaps = d["date"].diff().dt.days
ng = int((gaps > 4).sum())
if ng:
    issues.append({"issue": f"{ng} gaps of >4 calendar days between consecutive rows (holidays/non-publication days)",
                   "handling": "changes computed between consecutive rows as specified; multi-day gaps retained"})

baseline = float(d["chg_bp"].mean())

effect = lo = hi = rmse = n_an = None
if n_treated > 0:
    first = d.index[d["post_cut"] == 1][0]
    post = d.loc[d["post_cut"] == 1, "chg_bp"].dropna()
    pre = d.loc[max(first - 120, 1):first - 1, "chg_bp"].dropna()
    if len(pre) >= 2 and len(post) >= 2:
        effect = float(post.mean() - pre.mean())
        v1, v0 = post.var(ddof=1) / len(post), pre.var(ddof=1) / len(pre)
        se = np.sqrt(v1 + v0)
        df = (v1 + v0) ** 2 / (v1**2 / (len(post) - 1) + v0**2 / (len(pre) - 1))
        t = stats.t.ppf(0.975, df)
        lo, hi = effect - t * se, effect + t * se
        resid = np.concatenate([post - post.mean(), pre - pre.mean()])
        rmse = float(np.sqrt(np.mean(resid**2)))
        n_an = int(len(post) + len(pre))

if effect is None:
    concl, feas = "premise_invalid", "infeasible"
    reason = ("No rows with post_cut=1: the file ends on 2019-12-31, before the 2020-03-04 cut, "
              "so neither the post-cut mean nor the 120-day pre-window can be formed.")
else:
    concl = "positive_effect" if lo > 0 else "negative_effect" if hi < 0 else "no_detectable_effect"
    feas, reason = "complete", None

res = {"task_id": "F6-conf-02",
       "sample_start": str(d["date"].min().date()), "sample_end": str(last.date()),
       "n_treated": n_treated, "baseline_metric": baseline, "model_metric": rmse,
       "effect_estimate": effect, "effect_ci_low": None if lo is None else float(lo),
       "effect_ci_high": None if hi is None else float(hi),
       "conclusion": concl, "feasibility": feas, "feasibility_reason": reason,
       "data_issues": issues, "n_input_rows": {"gc5y_daily.csv": n_rows},
       "n_analysed": n_an, "random_seed": SEED}
json.dump(res, open("out/results.json", "w"), indent=2)

# table
d["year"] = d["date"].dt.year
tab = d.groupby("year").agg(n_days=("yield_5y", "size"), mean_yield=("yield_5y", "mean"),
                            mean_change_bp=("chg_bp", "mean")).reset_index()
tab.round(4).to_csv("out/table1.csv", index=False)

# figure
fig, ax = plt.subplots(figsize=(9, 4))
ax.plot(d["date"], d["yield_5y"], lw=0.8)
ax.axvline(pd.Timestamp("2020-03-04"), color="r", ls="--", lw=1, label="2020-03-04 cut (outside sample)")
ax.set_xlim(d["date"].min(), pd.Timestamp("2020-06-30"))
ax.set_ylabel("GoC 5-year yield (%)"); ax.set_title("GoC 5-year benchmark yield, file sample (ends 2019-12-31)")
ax.legend(); fig.tight_layout(); fig.savefig("out/figure1.png", dpi=120)
print(json.dumps(res, indent=2))