from __future__ import annotations

import os
import json
import datetime
import pandas as pd
import numpy as np

from .gate_row20_cosmology import gate_cosmo


def _utc_now_z() -> str:
    return datetime.datetime.now(datetime.UTC).isoformat(timespec="seconds").replace("+00:00", "Z")


def _read_vector_csv(path: str) -> np.ndarray:
    vals = []
    with open(path, "r", encoding="utf-8", errors="ignore") as f:
        for line in f:
            s = line.strip()
            if not s:
                continue
            if s.startswith("#"):
                continue
            try:
                vals.append(float(s))
            except ValueError:
                continue
    return np.asarray(vals, dtype=float)


def _read_matrix_csv(path: str) -> np.ndarray:
    rows = []
    with open(path, "r", encoding="utf-8", errors="ignore") as f:
        for line in f:
            s = line.strip()
            if not s:
                continue
            if s.startswith("#"):
                continue
            try:
                row = [float(x) for x in s.replace("\t", ",").split(",") if x.strip() != ""]
                if row:
                    rows.append(row)
            except ValueError:
                continue
    return np.asarray(rows, dtype=float)


def _chi2(data: np.ndarray, theory: np.ndarray, cov: np.ndarray) -> float:
    delta = data - theory
    inv = np.linalg.inv(cov)
    return float(delta.T @ inv @ delta)


def _resolve_ahg(dataset_dir: str) -> tuple[str, str, str]:
    pred_path = os.path.join(dataset_dir, "axion_pred.csv")
    lim_path = os.path.join(dataset_dir, "axion_limit95.csv")

    if not os.path.exists(pred_path) or not os.path.exists(lim_path):
        return "PENDING", "missing input file(s)", ""

    try:
        pred_df = pd.read_csv(pred_path)
        lim_df = pd.read_csv(lim_path)
    except Exception as e:
        return "PENDING", f"csv read failed: {e}", ""

    pred_x_candidates = ["mass", "m", "ma", "mass_uev", "x"]
    pred_y_candidates = ["pred", "prediction", "g_pred", "value", "y"]

    lim_x_candidates = ["mass", "m", "ma", "mass_uev", "x"]
    lim_y_candidates = ["limit95", "limit", "g_limit95", "g95", "value", "y"]

    def pick_col(df: pd.DataFrame, candidates: list[str]):
        cols = list(df.columns)
        lower_map = {str(c).strip().lower(): c for c in cols}
        for name in candidates:
            if name.lower() in lower_map:
                return lower_map[name.lower()]
        return None

    pred_x = pick_col(pred_df, pred_x_candidates)
    pred_y = pick_col(pred_df, pred_y_candidates)
    lim_x = pick_col(lim_df, lim_x_candidates)
    lim_y = pick_col(lim_df, lim_y_candidates)

    if pred_x is None or pred_y is None:
        if pred_df.shape[1] >= 2:
            pred_x, pred_y = pred_df.columns[:2]
        else:
            return "PENDING", "axion_pred.csv lacks usable columns", ""

    if lim_x is None or lim_y is None:
        if lim_df.shape[1] >= 2:
            lim_x, lim_y = lim_df.columns[:2]
        else:
            return "PENDING", "axion_limit95.csv lacks usable columns", ""

    work_pred = pred_df[[pred_x, pred_y]].copy()
    work_lim = lim_df[[lim_x, lim_y]].copy()

    work_pred.columns = ["x", "pred"]
    work_lim.columns = ["x", "limit95"]

    work_pred["x"] = pd.to_numeric(work_pred["x"], errors="coerce")
    work_pred["pred"] = pd.to_numeric(work_pred["pred"], errors="coerce")
    work_lim["x"] = pd.to_numeric(work_lim["x"], errors="coerce")
    work_lim["limit95"] = pd.to_numeric(work_lim["limit95"], errors="coerce")

    work_pred = work_pred.dropna()
    work_lim = work_lim.dropna()

    merged = pd.merge(work_pred, work_lim, on="x", how="inner").sort_values("x")

    if merged.empty:
        return "PENDING", "no overlapping numeric x-values between axion_pred.csv and axion_limit95.csv", ""

    delta = merged["limit95"].to_numpy() - merged["pred"].to_numpy()
    delta_min = float(np.min(delta))

    if delta_min >= 0.0:
        return "PASS", "", f"n={len(merged)}, min(limit95-pred)={delta_min:.6g}"
    return "FAIL", "", f"n={len(merged)}, min(limit95-pred)={delta_min:.6g}"


def _resolve_mhg(dataset_dir: str) -> tuple[str, str, str]:
    data_path = os.path.join(dataset_dir, "gw_data.csv")
    theory_path = os.path.join(dataset_dir, "gw_theory.csv")
    cov_path = os.path.join(dataset_dir, "gw_cov.csv")

    if not all(os.path.exists(p) for p in [data_path, theory_path, cov_path]):
        return "PENDING", "missing input file(s)", ""

    d = _read_vector_csv(data_path)
    t = _read_vector_csv(theory_path)
    C = _read_matrix_csv(cov_path)

    if len(d) == 0 or len(t) == 0 or C.size == 0:
        return "PENDING", "empty numeric arrays after parsing", ""

    if len(t) == 1 and len(d) > 1:
        t = np.full(len(d), float(t[0]), dtype=float)

    n = min(len(d), len(t), C.shape[0], C.shape[1])
    d = d[:n]
    t = t[:n]
    C = C[:n, :n]

    try:
        x2 = _chi2(d, t, C)
        chi2_dof = x2 / max(n, 1)
    except Exception as e:
        return "PENDING", str(e), ""

    if chi2_dof <= 1.35:
        return "PASS", "", f"n={n}, chi2={x2:.6g}, chi2_dof={chi2_dof:.6g}"
    return "FAIL", "", f"n={n}, chi2={x2:.6g}, chi2_dof={chi2_dof:.6g}"


def _resolve_ccg(dataset_dir: str) -> tuple[str, str, str]:
    """
    v21-aligned cosmology consistency check.

    Instead of recomputing an older standalone CMB+BAO-only aggregate, use the
    current Row 20 gate directly so Row 25 stays aligned with the repaired
    cosmology acceptance definition.
    """
    try:
        row20 = gate_cosmo(dataset_dir)
    except Exception as e:
        return "PENDING", f"row20 invocation failed: {e}", ""

    verdict = str(row20.get("verdict", "PENDING_DATA")).upper()
    summary = row20.get("summary", {}) or {}
    chi2_dof = float(summary.get("chi2_per_dof", np.inf))
    threshold = float(summary.get("threshold_chi2_per_dof", 1.25))
    signature = row20.get("signature", "ROW20_UNKNOWN")

    if verdict == "PASS":
        return "PASS", "", f"{signature}: chi2_dof={chi2_dof:.6g} <= {threshold:.6g}"

    if verdict == "PENDING_DATA":
        missing_required = summary.get("missing_required", [])
        return "PENDING", "", f"{signature}: pending row20 data; missing_required={missing_required}"

    return "FAIL", "", f"{signature}: chi2_dof={chi2_dof:.6g} > {threshold:.6g}"


def _resolve_cig(dataset_dir: str) -> tuple[str, str, str]:
    path = os.path.join(dataset_dir, "cig_status.json")
    if not os.path.exists(path):
        return "PENDING", "missing input file(s)", ""

    try:
        with open(path, "r", encoding="utf-8") as f:
            obj = json.load(f)
    except Exception as e:
        return "PENDING", str(e), ""

    spectral_radius_pass = bool(obj.get("spectral_radius_pass", False))
    ci_width = float(obj.get("ci_width", np.inf))
    ci_width_threshold = float(obj.get("ci_width_threshold", 0.15))

    if spectral_radius_pass and ci_width <= ci_width_threshold:
        return "PASS", "", f"rho_guard=PASS, ci_width={ci_width:.6g} <= {ci_width_threshold:.6g}"
    return "FAIL", "", f"rho_guard={'PASS' if spectral_radius_pass else 'FAIL'}, ci_width={ci_width:.6g} > {ci_width_threshold:.6g}"


def _resolve_pmg(dataset_dir: str) -> tuple[str, str, str]:
    path = os.path.join(dataset_dir, "pmg_status.json")
    if not os.path.exists(path):
        return "PENDING", "missing input file(s)", ""

    try:
        with open(path, "r", encoding="utf-8") as f:
            obj = json.load(f)
    except Exception as e:
        return "PENDING", str(e), ""

    frg_fixed_point_found = bool(obj.get("frg_fixed_point_found", False))
    beta_norm = float(obj.get("beta_norm", np.inf))
    beta_norm_threshold = float(obj.get("beta_norm_threshold", 1e-12))
    positivity_pass = bool(obj.get("positivity_pass", False))

    if frg_fixed_point_found and positivity_pass and beta_norm <= beta_norm_threshold:
        return "PASS", "", f"fixed_point=PASS, positivity=PASS, beta_norm={beta_norm:.6g}"
    return "FAIL", "", f"fixed_point={frg_fixed_point_found}, positivity={positivity_pass}, beta_norm={beta_norm:.6g}"


def _resolve_cpg(dataset_dir: str) -> tuple[str, str, str]:
    path = os.path.join(dataset_dir, "cpg_status.json")
    if not os.path.exists(path):
        return "PENDING", "missing input file(s)", ""

    try:
        with open(path, "r", encoding="utf-8") as f:
            obj = json.load(f)
    except Exception as e:
        return "PENDING", str(e), ""

    continuity_pass = bool(obj.get("continuity_pass", False))
    max_jump = float(obj.get("max_jump", np.inf))
    jump_threshold = float(obj.get("jump_threshold", 0.05))

    if continuity_pass and max_jump <= jump_threshold:
        return "PASS", "", f"continuity=PASS, max_jump={max_jump:.6g} <= {jump_threshold:.6g}"
    return "FAIL", "", f"continuity={continuity_pass}, max_jump={max_jump:.6g} > {jump_threshold:.6g}"


def gate_row25_stress(dataset_dir: str) -> dict:
    report_path = os.path.join(dataset_dir, "stress_report.csv")
    registry_path = os.path.join(dataset_dir, "stress_registry.yaml")

    if not os.path.exists(report_path):
        return {
            "row": 25,
            "gate": "break_utfanswf_stress_registry",
            "context": "LIVE",
            "tier": "Diagnostic",
            "timestamp_utc": _utc_now_z(),
            "verdict": "PENDING_DATA",
            "blocking": False,
            "notes": "Missing datasets/stress_report.csv",
        }

    df = pd.read_csv(report_path)
    df["outcome"] = df["outcome"].astype("object")
    df["error"] = df["error"].astype("object")
    df["reason"] = df["reason"].astype("object")

    required_cols = ["test", "outcome", "error", "reason"]
    missing = [c for c in required_cols if c not in df.columns]
    if missing:
        return {
            "row": 25,
            "gate": "break_utfanswf_stress_registry",
            "context": "LIVE",
            "tier": "Diagnostic",
            "timestamp_utc": _utc_now_z(),
            "verdict": "PENDING_DATA",
            "blocking": False,
            "notes": f"stress_report.csv missing required columns: {missing}",
        }

    resolvers = {
        "AHG": _resolve_ahg,
        "MHG": _resolve_mhg,
        "CCG": _resolve_ccg,
        "CPG": _resolve_cpg,
        "PMG": _resolve_pmg,
        "CIG": _resolve_cig,
    }

    for i, row in df.iterrows():
        test = str(row["test"]).strip().upper()
        if test in resolvers:
            outcome, error, reason = resolvers[test](dataset_dir)
            df.at[i, "outcome"] = outcome
            df.at[i, "error"] = error
            df.at[i, "reason"] = reason

    df.to_csv(report_path, index=False)

    df["outcome_norm"] = df["outcome"].astype(str).str.strip().str.upper()

    n_tests = int(len(df))
    n_pass = int((df["outcome_norm"] == "PASS").sum())
    n_fail = int((df["outcome_norm"] == "FAIL").sum())
    n_pending = int((df["outcome_norm"] == "PENDING").sum())

    critical_triggered = n_fail > 0

    if critical_triggered:
        verdict = "FAIL"
    elif n_pending > 0:
        verdict = "PENDING_DATA"
    else:
        verdict = "PASS"

    notes = (
        f"Registry file present={os.path.exists(registry_path)}; "
        f"tests={n_tests}, pass={n_pass}, pending={n_pending}, fail={n_fail}."
    )

    return {
        "row": 25,
        "gate": "break_utfanswf_stress_registry",
        "context": "LIVE",
        "tier": "Diagnostic",
        "timestamp_utc": _utc_now_z(),
        "verdict": verdict,
        "blocking": False,
        "n_tests": n_tests,
        "n_pass": n_pass,
        "n_pending": n_pending,
        "n_fail": n_fail,
        "critical_triggered": critical_triggered,
        "registry_present": os.path.exists(registry_path),
        "notes": notes,
    }


def _write_row25_tex(res: dict, tex_path: str) -> None:
    def tex_escape(s: object) -> str:
        text = str(s)
        return (
            text.replace("\\", r"\textbackslash{}")
            .replace("{", r"\{")
            .replace("}", r"\}")
            .replace("_", r"\_")
            .replace("&", r"\&")
            .replace("%", r"\%")
            .replace("#", r"\#")
        )

    lines = [
        rf"\def\RowXXVVerdict{{{tex_escape(res.get('verdict', 'PENDING_DATA'))}}}",
        rf"\def\RowXXVRunTag{{{tex_escape(res.get('gate', 'break_utfanswf_stress_registry'))}}}",
        rf"\def\RowXXVTier{{{tex_escape(res.get('tier', 'Diagnostic'))}}}",
        rf"\def\RowXXVContext{{{tex_escape(res.get('context', 'LIVE'))}}}",
        rf"\def\RowXXVNTests{{{tex_escape(res.get('n_tests', 0))}}}",
        rf"\def\RowXXVNPass{{{tex_escape(res.get('n_pass', 0))}}}",
        rf"\def\RowXXVNPending{{{tex_escape(res.get('n_pending', 0))}}}",
        rf"\def\RowXXVNFail{{{tex_escape(res.get('n_fail', 0))}}}",
        rf"\def\RowXXVCriticalTriggered{{{tex_escape(str(res.get('critical_triggered', False)).lower())}}}",
        rf"\def\RowXXVNotes{{{tex_escape(res.get('notes', ''))}}}",
    ]

    with open(tex_path, "w", encoding="utf-8") as f:
        f.write("\n".join(lines) + "\n")


if __name__ == "__main__":
    here = os.path.dirname(__file__)
    repo_root = os.path.abspath(os.path.join(here, ".."))
    dataset_dir = os.path.join(repo_root, "datasets")
    out_dir = os.path.join(repo_root, "results", "row25_stress")
    tex_path = os.path.join(repo_root, "results", "row25_results.tex")

    os.makedirs(out_dir, exist_ok=True)

    res = gate_row25_stress(dataset_dir)

    with open(os.path.join(out_dir, "result.json"), "w", encoding="utf-8") as f:
        json.dump(res, f, indent=2, sort_keys=True)

    _write_row25_tex(res, tex_path)

    print(res["verdict"])