#!/usr/bin/env python3
"""WAVE2B Track B simple candidate rebuild (research only — not a score).

Rebuilds PEXP_R(t-12) vs A229RX(t)/A229RX(t-12)-1 monthly pairs + annual means
from already-ingested Michigan SCA PEXP + BEA A229RX monthly files.

Run from srp-observatory root:
  python3 scripts/wave2b_track_b_rebuild.py

Inputs (must already exist):
  data/m1/mich_pago_pexp_monthly.json
  data/m1/a229rx_monthly.csv

Outputs:
  data/m1/track_b_simple_pexp_a229rx_pairs.json
  data/m1/track_b_simple_pexp_a229rx_monthly.csv
  data/m1/track_b_simple_pexp_a229rx_annual.csv

Gates: no M1 scoring; no formula freeze; descriptive Pearson r only.
"""
from __future__ import annotations

import csv
import json
import math
from collections import defaultdict
from datetime import datetime, timezone
from pathlib import Path
from statistics import mean

ROOT = Path(__file__).resolve().parents[1]
M1 = ROOT / "data" / "m1"
CANDIDATE = "B_simple_pexp_vs_a229rx_yoy"
DEFINITION = "PEXP_R(t-12) vs A229RX(t)/A229RX(t-12)-1"


def pearson(xs, ys):
    n = len(xs)
    if n < 2:
        return float("nan")
    mx, my = mean(xs), mean(ys)
    num = sum((x - mx) * (y - my) for x, y in zip(xs, ys))
    den = math.sqrt(sum((x - mx) ** 2 for x in xs) * sum((y - my) ** 2 for y in ys))
    return num / den if den else float("nan")


def load_pexp():
    raw = json.loads((M1 / "mich_pago_pexp_monthly.json").read_text())
    out = {}
    for o in raw["pexp"]["observations"]:
        out[o["obs_date"][:10]] = float(o["value"])
    return out


def load_a229():
    out = {}
    with (M1 / "a229rx_monthly.csv").open() as f:
        for r in csv.DictReader(f):
            out[r["obs_date"][:10]] = float(r["value"])
    return out


def shift_months(ymd: str, delta: int) -> str:
    y, m, d = map(int, ymd.split("-"))
    m0 = y * 12 + (m - 1) + delta
    y2, m2 = divmod(m0, 12)
    return f"{y2:04d}-{m2 + 1:02d}-{d:02d}"


def main():
    pexp = load_pexp()
    a229 = load_a229()
    pairs = []
    for date, level in sorted(a229.items()):
        lag = shift_months(date, -12)
        pexp_date = lag
        if lag not in a229 or pexp_date not in pexp:
            continue
        yoy = level / a229[lag] - 1.0
        y, m, _ = map(int, date.split("-"))
        pairs.append(
            {
                "realization_date": date,
                "realization_year": y,
                "realization_month": m,
                "pexp_date": pexp_date,
                "pexp_r_prior": pexp[pexp_date],
                "real_dpi_pc_yoy": round(yoy, 6),
                "real_dpi_pc_yoy_pct": round(yoy * 100, 4),
                "a229rx_level": level,
                "a229rx_level_lag12": a229[lag],
                "candidate": CANDIDATE,
                "definition": DEFINITION,
                "status": "research_only_not_frozen",
            }
        )

    by_year = defaultdict(list)
    for p in pairs:
        by_year[p["realization_year"]].append(p)

    annual = []
    for year in sorted(by_year):
        rows = by_year[year]
        if len(rows) != 12:
            continue
        annual.append(
            {
                "year": year,
                "n_months": 12,
                "pexp_r_prior_mean": round(mean(r["pexp_r_prior"] for r in rows), 4),
                "real_dpi_pc_yoy_mean": round(mean(r["real_dpi_pc_yoy"] for r in rows), 6),
                "real_dpi_pc_yoy_pct_mean": round(mean(r["real_dpi_pc_yoy_pct"] for r in rows), 4),
                "status": "research_only_not_frozen",
            }
        )

    r_annual = pearson(
        [a["pexp_r_prior_mean"] for a in annual],
        [a["real_dpi_pc_yoy_mean"] for a in annual],
    )
    r_monthly = pearson(
        [p["pexp_r_prior"] for p in pairs],
        [p["real_dpi_pc_yoy"] for p in pairs],
    )

    payload = {
        "status": "RESEARCH ONLY — no M1 production scoring; no formula freeze",
        "generated_at": datetime.now(timezone.utc).isoformat(),
        "candidate": CANDIDATE,
        "definition": DEFINITION,
        "expectation_leg": "MICH_PEXP_R at t-12",
        "outcome_leg": "A229RX YoY (t vs t-12)",
        "notes": [
            "Research pairing only — do not freeze weights/formula",
            "Descriptive Pearson r on full-12 annual means is not a score",
            "Monthly Pearson reported for transparency; primary reported r is annual-means",
            "score_authorized(M1)=false",
        ],
        "n_monthly_pairs": len(pairs),
        "n_full12_annual_years": len(annual),
        "descriptive_pearson_r_annual_means": round(r_annual, 6),
        "descriptive_pearson_r_monthly": round(r_monthly, 6),
        "sample_window": {
            "monthly_first": pairs[0]["realization_date"] if pairs else None,
            "monthly_last": pairs[-1]["realization_date"] if pairs else None,
            "annual_first": annual[0]["year"] if annual else None,
            "annual_last": annual[-1]["year"] if annual else None,
            "annual_frequency": "calendar-year mean of 12 monthly pairs (full-12 years only)",
            "r_primary": "pearson of annual means (n=full-12 years)",
        },
        "monthly_pairs": pairs,
        "annual_full12": annual,
    }

    (M1 / "track_b_simple_pexp_a229rx_pairs.json").write_text(json.dumps(payload, indent=2) + "\n")
    with (M1 / "track_b_simple_pexp_a229rx_monthly.csv").open("w", newline="") as f:
        w = csv.DictWriter(f, fieldnames=list(pairs[0].keys()))
        w.writeheader()
        w.writerows(pairs)
    with (M1 / "track_b_simple_pexp_a229rx_annual.csv").open("w", newline="") as f:
        w = csv.DictWriter(f, fieldnames=list(annual[0].keys()))
        w.writeheader()
        w.writerows(annual)

    print(json.dumps({
        "n_monthly_pairs": len(pairs),
        "n_full12_annual_years": len(annual),
        "descriptive_pearson_r_annual_means": round(r_annual, 6),
        "descriptive_pearson_r_monthly": round(r_monthly, 6),
        "annual_years": f"{annual[0]['year']}-{annual[-1]['year']}" if annual else None,
        "monthly_window": f"{pairs[0]['realization_date']}..{pairs[-1]['realization_date']}" if pairs else None,
    }, indent=2))


if __name__ == "__main__":
    main()