#!/usr/bin/env python3
"""M1-LONG Michigan SCA same-person pilot — research diagnostics only.

Links Interview-1 PEXP → Interview-3 PAGO via (IDPREV, DATEPR) chain
(and IDPREV2/DATEPR2 shortcut when populated). REJECTS ID-only linkage.

NO production scoring. Does not touch L1 or score_authorized(M1).

Plan: docs/m1/WAVE2B_M1_LONG_PILOT_PLAN.md
"""

from __future__ import annotations

import hashlib
import json
import math
from collections import Counter, defaultdict
from dataclasses import asdict, dataclass
from datetime import datetime, timezone
from pathlib import Path
from typing import Any, Optional

import pandas as pd

ROOT = Path(__file__).resolve().parents[2]
MICRO_CSV = (
    ROOT
    / "data"
    / "raw"
    / "michigan_sca"
    / "microdata"
    / "sca_micro_2015_2023_pilot.csv"
)
PROVENANCE_MD = (
    ROOT / "data" / "raw" / "michigan_sca" / "microdata" / "PROVENANCE.md"
)
OUT_DIR = ROOT / "data" / "m1"

# SAMPLE codes — pilot plan §3.5
SAMPLE_FRESH = {1, 3, 6}
SAMPLE_REINT2 = {2, 4, 7}
SAMPLE_REINT3 = {5, 8}

# PEXP/PAGO valid response codes (8/9 = missing)
VALID_ATT = {1, 3, 5}
SIGNED = {1: 1, 3: 0, 5: -1}
LABEL = {1: "better", 3: "same", 5: "worse"}

METHOD_LABEL = {
    1: "telephone_cati",
    2: "web",
    3: "telephone_papi",
    4: "in_person",
}


def yyyymm_gap_months(origin: int, dest: int) -> int:
    oy, om = divmod(int(origin), 100)
    dy, dm = divmod(int(dest), 100)
    return (dy - oy) * 12 + (dm - om)


def mode_era(yyyymm: int) -> str:
    """Annotate phone→web 2024 break eras (pilot plan §5)."""
    y = int(yyyymm)
    if y <= 202403:
        return "phone"
    if y <= 202406:
        return "transition"
    return "web"


def sha256_file(path: Path) -> str:
    h = hashlib.sha256()
    with path.open("rb") as f:
        for chunk in iter(lambda: f.read(1 << 20), b""):
            h.update(chunk)
    return h.hexdigest()


def _to_num(series: pd.Series) -> pd.Series:
    return pd.to_numeric(series, errors="coerce")


@dataclass
class PairRow:
    origin_caseid: int
    dest_caseid: int
    mid_caseid: Optional[int]
    origin_id: int
    origin_yyyymm: int
    mid_id: Optional[int]
    mid_yyyymm: Optional[int]
    dest_id: int
    dest_yyyymm: int
    origin_sample: int
    mid_sample: Optional[int]
    dest_sample: int
    origin_method: int
    dest_method: int
    pexp: int
    pago: int
    pexp_signed: int
    pago_signed: int
    error_signed: int
    gap_months: int
    gap_in_11_13: int
    link_method: str  # idprev2_shortcut | idprev_chain
    chain_agrees_idprev2: Optional[bool]
    mode_pair: str
    origin_mode_era: str
    dest_mode_era: str
    wt_dest: float
    wt_origin: float
    age_origin: Optional[float]
    region_origin: Optional[int]


def load_micro(path: Path = MICRO_CSV) -> pd.DataFrame:
    df = pd.read_csv(path)
    for c in ["IDPREV", "DATEPR", "IDPREV2", "DATEPR2", "WT_HH", "AGE", "EDUC"]:
        if c in df.columns:
            df[c] = _to_num(df[c])
    for c in ["CASEID", "YYYYMM", "ID", "SAMPLE", "METHOD", "PEXP", "PAGO", "SEX", "REGION"]:
        if c in df.columns:
            df[c] = _to_num(df[c])
    if "WT" in df.columns:
        df["WT"] = _to_num(df["WT"])
    return df


def index_by_id_yyyymm(df: pd.DataFrame) -> dict[tuple[int, int], Any]:
    """Map (ID, YYYYMM) → row namedtuple. Fail on duplicates."""
    idx: dict[tuple[int, int], Any] = {}
    for row in df.itertuples(index=False):
        key = (int(row.ID), int(row.YYYYMM))
        if key in idx:
            raise ValueError(f"duplicate (ID, YYYYMM) in microdata: {key}")
        idx[key] = row
    return idx


def link_idprev2(dest, by_key) -> Optional[Any]:
    """Primary shortcut: (IDPREV2, DATEPR2) → (ID, YYYYMM)."""
    if pd.isna(dest.IDPREV2) or pd.isna(dest.DATEPR2):
        return None
    return by_key.get((int(dest.IDPREV2), int(dest.DATEPR2)))


def link_chain_to_origin(dest, by_key) -> tuple[Optional[Any], Optional[Any]]:
    """Chain: dest --(IDPREV,DATEPR)→ mid --(IDPREV,DATEPR)→ origin.

    Returns (origin, mid). REJECTS ID-only matching.
    """
    if pd.isna(dest.IDPREV) or pd.isna(dest.DATEPR):
        return None, None
    mid = by_key.get((int(dest.IDPREV), int(dest.DATEPR)))
    if mid is None:
        return None, None
    if pd.isna(mid.IDPREV) or pd.isna(mid.DATEPR):
        return None, mid
    origin = by_key.get((int(mid.IDPREV), int(mid.DATEPR)))
    return origin, mid


def mode_pair_label(origin_method: int, dest_method: int) -> str:
    om = METHOD_LABEL.get(int(origin_method), f"m{origin_method}")
    dm = METHOD_LABEL.get(int(dest_method), f"m{dest_method}")
    if origin_method == dest_method == 1:
        return "phone→phone"
    if origin_method == dest_method == 2:
        return "web→web"
    if origin_method != dest_method:
        return f"cross-mode:{om}→{dm}"
    return f"{om}→{dm}"


def build_pairs(df: pd.DataFrame) -> tuple[list[PairRow], dict[str, Any]]:
    """Build interview-3 → interview-1 pairs. Research-only."""
    by_key = index_by_id_yyyymm(df)
    pairs: list[PairRow] = []
    seen: set[tuple[int, int]] = set()
    qc = Counter()

    dest_mask = df["SAMPLE"].isin(SAMPLE_REINT3) | (
        df["IDPREV2"].notna() & df["DATEPR2"].notna()
    )
    dests = df.loc[dest_mask]

    for dest in dests.itertuples(index=False):
        if int(dest.PAGO) not in VALID_ATT:
            qc["dest_pago_missing"] += 1
            continue

        origin_shortcut = link_idprev2(dest, by_key)
        origin_chain, mid = link_chain_to_origin(dest, by_key)

        if origin_shortcut is not None and origin_chain is not None:
            if (int(origin_shortcut.ID), int(origin_shortcut.YYYYMM)) != (
                int(origin_chain.ID),
                int(origin_chain.YYYYMM),
            ):
                qc["chain_idprev2_disagree"] += 1
                continue  # drop contradictory pairs (plan §4.2)
            origin = origin_shortcut
            link_method = "idprev2_shortcut"
            chain_agrees = True
        elif origin_shortcut is not None:
            origin = origin_shortcut
            link_method = "idprev2_shortcut"
            chain_agrees = None if origin_chain is None else False
            mid = mid  # may be None
        elif origin_chain is not None:
            origin = origin_chain
            link_method = "idprev_chain"
            chain_agrees = None  # shortcut absent
        else:
            qc["no_link"] += 1
            continue

        if int(origin.SAMPLE) not in SAMPLE_FRESH:
            qc["origin_not_fresh"] += 1
            continue
        if int(origin.PEXP) not in VALID_ATT:
            qc["origin_pexp_missing"] += 1
            continue

        person_key = (int(origin.ID), int(origin.YYYYMM))
        if person_key in seen:
            qc["duplicate_person_key"] += 1
            continue  # keep first; note in QC (do not crash research run)
        seen.add(person_key)

        gap = yyyymm_gap_months(int(origin.YYYYMM), int(dest.YYYYMM))
        pexp_s = SIGNED[int(origin.PEXP)]
        pago_s = SIGNED[int(dest.PAGO)]
        age = float(origin.AGE) if not pd.isna(getattr(origin, "AGE", float("nan"))) else None
        region = (
            int(origin.REGION)
            if not pd.isna(getattr(origin, "REGION", float("nan")))
            else None
        )

        pairs.append(
            PairRow(
                origin_caseid=int(origin.CASEID),
                dest_caseid=int(dest.CASEID),
                mid_caseid=int(mid.CASEID) if mid is not None else None,
                origin_id=int(origin.ID),
                origin_yyyymm=int(origin.YYYYMM),
                mid_id=int(mid.ID) if mid is not None else None,
                mid_yyyymm=int(mid.YYYYMM) if mid is not None else None,
                dest_id=int(dest.ID),
                dest_yyyymm=int(dest.YYYYMM),
                origin_sample=int(origin.SAMPLE),
                mid_sample=int(mid.SAMPLE) if mid is not None else None,
                dest_sample=int(dest.SAMPLE),
                origin_method=int(origin.METHOD),
                dest_method=int(dest.METHOD),
                pexp=int(origin.PEXP),
                pago=int(dest.PAGO),
                pexp_signed=pexp_s,
                pago_signed=pago_s,
                error_signed=pexp_s - pago_s,
                gap_months=gap,
                gap_in_11_13=int(11 <= gap <= 13),
                link_method=link_method,
                chain_agrees_idprev2=chain_agrees,
                mode_pair=mode_pair_label(int(origin.METHOD), int(dest.METHOD)),
                origin_mode_era=mode_era(int(origin.YYYYMM)),
                dest_mode_era=mode_era(int(dest.YYYYMM)),
                wt_dest=float(dest.WT) if not pd.isna(dest.WT) else float("nan"),
                wt_origin=float(origin.WT) if not pd.isna(origin.WT) else float("nan"),
                age_origin=age,
                region_origin=region,
            )
        )
        qc["paired"] += 1

    meta = {"qc_counts": dict(qc), "n_pairs": len(pairs)}
    return pairs, meta


def attrition_funnel(df: pd.DataFrame, pairs: list[PairRow]) -> dict[str, Any]:
    fresh = df[df["SAMPLE"].isin(SAMPLE_FRESH)]
    fresh_pexp = fresh[fresh["PEXP"].isin(VALID_ATT)]
    # Origins that appear as prior of a reinterview-2 row
    mid_rows = df[df["SAMPLE"].isin(SAMPLE_REINT2)]
    mid_priors = set()
    for r in mid_rows.itertuples(index=False):
        if pd.isna(r.IDPREV) or pd.isna(r.DATEPR):
            continue
        mid_priors.add((int(r.IDPREV), int(r.DATEPR)))
    n1 = int(len(fresh_pexp))
    n2 = int(
        sum(
            1
            for r in fresh_pexp.itertuples(index=False)
            if (int(r.ID), int(r.YYYYMM)) in mid_priors
        )
    )
    n3 = len(pairs)
    n_pair_primary = sum(1 for p in pairs if p.gap_in_11_13)
    return {
        "N1_fresh_valid_pexp": n1,
        "N2_appear_as_reint2_prior": n2,
        "N2_over_N1": round(n2 / n1, 4) if n1 else None,
        "N3_linked_pairs_any_gap": n3,
        "N3_over_N1": round(n3 / n1, 4) if n1 else None,
        "N_pair_gap_11_13": n_pair_primary,
        "N_pair_over_N1": round(n_pair_primary / n1, 4) if n1 else None,
        "notes": [
            "N2 is presence of fresh (ID,YYYYMM) as IDPREV/DATEPR of SAMPLE∈{2,4,7}.",
            "N3 requires successful chain/shortcut link + valid PEXP/PAGO + fresh origin.",
            "Selective attrition into interview 3 is expected; no IPW in v1.",
        ],
    }


def miss_tables(pairs: list[PairRow]) -> dict[str, Any]:
    primary = [p for p in pairs if p.gap_in_11_13]
    contig_u: dict[str, int] = defaultdict(int)
    contig_w: dict[str, float] = defaultdict(float)
    err_u: Counter = Counter()
    err_w: dict[int, float] = defaultdict(float)
    disappointment = 0
    pleasant = 0
    w_sum = 0.0
    for p in primary:
        key = f"pexp_{LABEL[p.pexp]}__pago_{LABEL[p.pago]}"
        contig_u[key] += 1
        w = p.wt_dest if not math.isnan(p.wt_dest) else 1.0
        contig_w[key] += w
        err_u[p.error_signed] += 1
        err_w[p.error_signed] += w
        w_sum += w
        # expected better & realized worse
        if p.pexp == 1 and p.pago == 5:
            disappointment += 1
        if p.pexp == 5 and p.pago == 1:
            pleasant += 1

    # matrix form
    matrix = {}
    for e in (1, 3, 5):
        for a in (1, 3, 5):
            matrix[f"{LABEL[e]}→{LABEL[a]}"] = contig_u.get(
                f"pexp_{LABEL[e]}__pago_{LABEL[a]}", 0
            )

    return {
        "n_primary_gap_11_13": len(primary),
        "contingency_unweighted": dict(sorted(contig_u.items())),
        "contingency_weighted_wt_dest": {k: round(v, 3) for k, v in sorted(contig_w.items())},
        "contingency_matrix_counts": matrix,
        "error_signed_unweighted": {str(k): int(v) for k, v in sorted(err_u.items())},
        "error_signed_weighted_wt_dest": {
            str(k): round(v, 3) for k, v in sorted(err_w.items())
        },
        "share_disappointed_exp_better_got_worse": round(disappointment / len(primary), 4)
        if primary
        else None,
        "share_pleasant_exp_worse_got_better": round(pleasant / len(primary), 4)
        if primary
        else None,
        "weight_sum_dest": round(w_sum, 3),
        "coding_map_research_only": {"1": "+1", "3": "0", "5": "-1", "8/9": "missing"},
        "error_definition": "error_i = s(PEXP_origin) - s(PAGO_dest); NOT production",
    }


def method_stratum(pairs: list[PairRow]) -> dict[str, Any]:
    by_mode = Counter(p.mode_pair for p in pairs)
    by_method = Counter((p.origin_method, p.dest_method) for p in pairs)
    by_era = Counter((p.origin_mode_era, p.dest_mode_era) for p in pairs)
    return {
        "mode_pair_counts": dict(by_mode),
        "method_code_pair_counts": {
            f"{a}→{b}": c for (a, b), c in by_method.items()
        },
        "mode_era_pair_counts": {
            f"{a}→{b}": c for (a, b), c in by_era.items()
        },
        "note_2024_break": (
            "This extract ends YYYYMM=202312. Official phone→web mix begins 202404 "
            "and web-only ≥202407, so the 2024 mode break is OUT OF SAMPLE here. "
            "All pairs in this pilot are METHOD=1 telephone (phone→phone)."
        ),
    }


def horizon_stats(pairs: list[PairRow]) -> dict[str, Any]:
    gaps = [p.gap_months for p in pairs]
    if not gaps:
        return {"n": 0}
    s = pd.Series(gaps)
    return {
        "n": len(gaps),
        "min": int(s.min()),
        "median": float(s.median()),
        "mean": float(s.mean()),
        "max": int(s.max()),
        "distribution": {str(k): int(v) for k, v in sorted(Counter(gaps).items())},
        "share_in_11_13": round(sum(1 for g in gaps if 11 <= g <= 13) / len(gaps), 4),
    }


def write_outputs(pairs: list[PairRow], df: pd.DataFrame, link_meta: dict) -> dict[str, Any]:
    OUT_DIR.mkdir(parents=True, exist_ok=True)
    pairs_path = OUT_DIR / "m1long_pairs_pilot_stratumA.csv"
    summary_path = OUT_DIR / "m1long_pilot_summary.json"
    protocol_path = OUT_DIR / "m1long_pilot_protocol.json"
    miss_path = OUT_DIR / "m1long_pilot_miss_table.csv"
    horizon_path = OUT_DIR / "m1long_pilot_horizon_months.csv"
    attrition_path = OUT_DIR / "m1long_pilot_attrition.json"

    pdf = pd.DataFrame([asdict(p) for p in pairs])
    pdf.to_csv(pairs_path, index=False)

    # miss table CSV (matrix)
    primary = [p for p in pairs if p.gap_in_11_13]
    rows = []
    for e in (1, 3, 5):
        for a in (1, 3, 5):
            n = sum(1 for p in primary if p.pexp == e and p.pago == a)
            w = sum(
                (p.wt_dest if not math.isnan(p.wt_dest) else 1.0)
                for p in primary
                if p.pexp == e and p.pago == a
            )
            rows.append(
                {
                    "pexp_code": e,
                    "pexp_label": LABEL[e],
                    "pago_code": a,
                    "pago_label": LABEL[a],
                    "n": n,
                    "wt_dest_sum": round(w, 3),
                    "share_unweighted": round(n / len(primary), 4) if primary else None,
                }
            )
    pd.DataFrame(rows).to_csv(miss_path, index=False)

    gap_dist = Counter(p.gap_months for p in pairs)
    pd.DataFrame(
        [{"gap_months": g, "n": gap_dist[g]} for g in sorted(gap_dist)]
    ).to_csv(horizon_path, index=False)

    attrition = attrition_funnel(df, pairs)
    attrition_path.write_text(json.dumps(attrition, indent=2) + "\n")

    miss = miss_tables(pairs)
    horizon = horizon_stats(pairs)
    method = method_stratum(pairs)

    idprev2_nonnull = int(df["IDPREV2"].notna().sum())
    datepr2_nonnull = int(df["DATEPR2"].notna().sum())

    summary = {
        "title": "M1-LONG Michigan SCA same-person pilot diagnostics",
        "status": "RESEARCH_ONLY",
        "score_authorized_M1": False,
        "L1_rule_id": "L1-US-v0.1",
        "formula_freeze": False,
        "generated_utc": datetime.now(timezone.utc).strftime("%Y-%m-%dT%H:%M:%SZ"),
        "microdata": {
            "path": str(MICRO_CSV.relative_to(ROOT)),
            "rows": int(len(df)),
            "yyyymm_min": int(df["YYYYMM"].min()),
            "yyyymm_max": int(df["YYYYMM"].max()),
            "sha256": sha256_file(MICRO_CSV),
            "provenance": str(PROVENANCE_MD.relative_to(ROOT)),
            "IDPREV2_nonnull": idprev2_nonnull,
            "DATEPR2_nonnull": datepr2_nonnull,
            "IDPREV2_note": (
                "IDPREV2/DATEPR2 are empty in this SDA extract; primary linkage used "
                "IDPREV/DATEPR two-step chain (interview 3→2→1). REJECT ID-only."
            ),
        },
        "n_pairs": len(pairs),
        "n_pairs_gap_11_13": sum(1 for p in pairs if p.gap_in_11_13),
        "horizon": horizon,
        "attrition": attrition,
        "miss": miss,
        "method_stratum": method,
        "link_qc": link_meta,
        "link_method_counts": dict(Counter(p.link_method for p in pairs)),
        "outputs": {
            "pairs": str(pairs_path.relative_to(ROOT)),
            "miss_table": str(miss_path.relative_to(ROOT)),
            "horizon": str(horizon_path.relative_to(ROOT)),
            "attrition": str(attrition_path.relative_to(ROOT)),
            "summary": str(summary_path.relative_to(ROOT)),
            "protocol": str(protocol_path.relative_to(ROOT)),
        },
        "gates": {
            "research_only": True,
            "no_production_scoring": True,
            "no_formula_freeze": True,
            "A2_M1D_remain_retired": True,
            "aggregate_A2_not_substitute": True,
        },
    }
    summary_path.write_text(json.dumps(summary, indent=2) + "\n")

    protocol = {
        "name": "m1long_pilot_stratumA_v0",
        "research_only": True,
        "score_authorized": False,
        "linkage": {
            "reject_id_only": True,
            "primary_when_available": "(IDPREV2, DATEPR2) → (ID, YYYYMM)",
            "used_in_this_extract": "(IDPREV, DATEPR) chain 3→2→1",
            "destination_sample": sorted(SAMPLE_REINT3),
            "origin_sample_fresh": sorted(SAMPLE_FRESH),
            "valid_pexp_pago": sorted(VALID_ATT),
            "missing_codes": [8, 9],
            "gap_primary_months": [11, 12, 13],
        },
        "window": {
            "extract_yyyymm": "201501–202312",
            "stratum": "A_phone_core",
            "mode_break_2024": "out_of_sample",
        },
        "signed_map_predeclared_not_frozen": SIGNED,
        "inputs": {
            "micro_csv": str(MICRO_CSV.relative_to(ROOT)),
            "sha256": summary["microdata"]["sha256"],
        },
    }
    protocol_path.write_text(json.dumps(protocol, indent=2) + "\n")

    # SHA256SUMS for pilot artifacts
    sums_path = OUT_DIR / "m1long_pilot_SHA256SUMS.txt"
    lines = []
    for p in [
        pairs_path,
        miss_path,
        horizon_path,
        attrition_path,
        summary_path,
        protocol_path,
    ]:
        lines.append(f"{sha256_file(p)}  {p.name}")
    sums_path.write_text("\n".join(lines) + "\n")

    return summary


def main() -> None:
    if not MICRO_CSV.exists():
        raise SystemExit(f"microdata missing: {MICRO_CSV}")
    df = load_micro()
    pairs, link_meta = build_pairs(df)
    summary = write_outputs(pairs, df, link_meta)
    print(
        json.dumps(
            {
                "n_pairs": summary["n_pairs"],
                "horizon_median": summary["horizon"].get("median"),
                "link_methods": summary["link_method_counts"],
                "summary": summary["outputs"]["summary"],
            },
            indent=2,
        )
    )


if __name__ == "__main__":
    main()