from pathlib import Path
import json
import math
import os

IN_JSON = Path("results/row20_cosmology/result.json")
OUT_TEX = Path("results/row20_results.tex")

def _is_num(x):
    return isinstance(x, (int, float)) and math.isfinite(float(x))

def fmt(x, digits=6):
    # LaTeX-safe display: show -- for missing/null/NaN
    if x is None:
        return "--"
    if isinstance(x, bool):
        return "true" if x else "false"
    if _is_num(x):
        # compact but stable formatting
        return f"{float(x):.{digits}g}"
    # strings / other
    s = str(x)
    return s.replace("_", "\\_")

def find_part(parts, tag):
    for p in parts:
        if p.get("tag") == tag:
            return p
    return None

def main():
    if not IN_JSON.exists():
        raise FileNotFoundError(f"Missing {IN_JSON}")

    data = json.loads(IN_JSON.read_text(encoding="utf-8"))

    verdict = data.get("verdict", "--")
    row = data.get("row", 20)
    signature = data.get("signature", "")
    context = data.get("context", "LIVE")

    summary = data.get("summary", {}) or {}
    chi2_total = summary.get("chi2_total", None)
    n_total = summary.get("n_total", None)
    chi2_per = summary.get("chi2_per_dof", None)
    threshold = summary.get("threshold_chi2_per_dof", None)
    missing_required = summary.get("missing_required", []) or []

    # IMPORTANT: meta lives under summary.cosmo_meta in your gate output
    meta = summary.get("cosmo_meta", {}) or {}

    parts = data.get("parts", []) or []

    # Common tags in your code/logs
    cmb = find_part(parts, "CMB")
    bao_dv = find_part(parts, "BAO_DV")
    bao_dm_h = find_part(parts, "BAO_DM_H")
    sn = find_part(parts, "SN")
    gw = find_part(parts, "GW")

    def part_chi2(p):
        return None if not p or p.get("status") != "RUN" else p.get("chi2", None)

    def part_n(p):
        return None if not p or p.get("status") != "RUN" else p.get("n", None)

    def part_chi2_per(p):
        return None if not p or p.get("status") != "RUN" else p.get("chi2_per_dof", None)

    # Build TeX
    out = []
    out.append("% Auto-generated. DO NOT EDIT BY HAND.\n")

    out.append(f"\\newcommand{{\\RowXXRow}}{{{fmt(row)}}}\n")
    out.append(f"\\newcommand{{\\RowXXContext}}{{{fmt(context)}}}\n")
    out.append(f"\\newcommand{{\\RowXXVerdict}}{{{fmt(verdict)}}}\n")
    out.append(f"\\newcommand{{\\RowXXRunTag}}{{{fmt(signature)}}}\n")

    # Summary totals
    out.append(f"\\newcommand{{\\RowXXChiSqTot}}{{{fmt(chi2_total)}}}\n")
    out.append(f"\\newcommand{{\\RowXXNtot}}{{{fmt(n_total)}}}\n")
    out.append(f"\\newcommand{{\\RowXXChiSqPer}}{{{fmt(chi2_per)}}}\n")
    out.append(f"\\newcommand{{\\RowXXThreshold}}{{{fmt(threshold)}}}\n")
    out.append(f"\\newcommand{{\\RowXXMissingRequired}}{{{fmt(missing_required)}}}\n")

    # Cosmology meta (these were coming out as -- in your PDF)
    out.append(f"\\newcommand{{\\RowXXh}}{{{fmt(meta.get('h'))}}}\n")
    out.append(f"\\newcommand{{\\RowXXHzero}}{{{fmt(meta.get('H0'))}}}\n")
    out.append(f"\\newcommand{{\\RowXXombhTwo}}{{{fmt(meta.get('ombh2'))}}}\n")
    out.append(f"\\newcommand{{\\RowXXomchTwo}}{{{fmt(meta.get('omch2'))}}}\n")

    # Optional baseline params (if present in meta; otherwise read cosmo_params.json)

    params_path = os.path.join("datasets", "cosmo_params.json")
    params = {}

    if os.path.exists(params_path):
        with open(params_path, "r", encoding="utf-8") as f:
            params = json.load(f)

    out.append(f"\\newcommand{{\\RowXXOk}}{{{fmt(meta.get('Omega_k', params.get('ok')))}}}\n")
    out.append(f"\\newcommand{{\\RowXXwZero}}{{{fmt(meta.get('w0', params.get('w0')))}}}\n")
    out.append(f"\\newcommand{{\\RowXXNeff}}{{{fmt(meta.get('Neff', params.get('N_eff')))}}}\n")
    out.append(f"\\newcommand{{\\RowXXTcmb}}{{{fmt(meta.get('T_CMB', params.get('T_CMB')))}}}\n")

    # Per-part numbers (show -- if not RUN)
    out.append(f"\\newcommand{{\\RowXXChiSqCMB}}{{{fmt(part_chi2(cmb))}}}\n")
    out.append(f"\\newcommand{{\\RowXXNCMB}}{{{fmt(part_n(cmb))}}}\n")
    out.append(f"\\newcommand{{\\RowXXChiSqPerCMB}}{{{fmt(part_chi2_per(cmb))}}}\n")

    out.append(f"\\newcommand{{\\RowXXChiSqBAO}}{{{fmt(part_chi2(bao_dv))}}}\n")
    out.append(f"\\newcommand{{\\RowXXNBAO}}{{{fmt(part_n(bao_dv))}}}\n")
    out.append(f"\\newcommand{{\\RowXXChiSqPerBAO}}{{{fmt(part_chi2_per(bao_dv))}}}\n")

    out.append(f"\\newcommand{{\\RowXXChiSqDMH}}{{{fmt(part_chi2(bao_dm_h))}}}\n")
    out.append(f"\\newcommand{{\\RowXXNDMH}}{{{fmt(part_n(bao_dm_h))}}}\n")
    out.append(f"\\newcommand{{\\RowXXChiSqPerDMH}}{{{fmt(part_chi2_per(bao_dm_h))}}}\n")

    out.append(f"\\newcommand{{\\RowXXChiSqSN}}{{{fmt(part_chi2(sn))}}}\n")
    out.append(f"\\newcommand{{\\RowXXNSN}}{{{fmt(part_n(sn))}}}\n")

    out.append(f"\\newcommand{{\\RowXXChiSqGW}}{{{fmt(part_chi2(gw))}}}\n")
    out.append(f"\\newcommand{{\\RowXXNGW}}{{{fmt(part_n(gw))}}}\n")

    OUT_TEX.parent.mkdir(parents=True, exist_ok=True)
    OUT_TEX.write_text("".join(out), encoding="utf-8")
    print(f"Wrote {OUT_TEX}")

if __name__ == "__main__":
    main()