from __future__ import annotations

import os
import json
import math
import datetime
import numpy as np
from scipy.optimize import minimize

from .common import load_vector_csv, load_matrix_csv, cov_diagnostics, chi2, sha256_file
from .cosmology_theory import generate_theory_vectors, write_vector_csv


def gate_cosmo(dataset_dir: str) -> dict:
    prov_path = os.path.join(dataset_dir, "cosmology_provenance.json")
    provenance = {}
    if os.path.exists(prov_path):
        try:
            with open(prov_path, "r", encoding="utf-8") as f:
                provenance = json.load(f)
        except Exception:
            provenance = {}

    parts: list[dict] = []
    total_chi2 = 0.0
    total_n = 0

    # Row 20 is a fit, not a score at one hand-picked background point.
    # The protected manuscript declares constrained minimization and charges
    # N_eff fitted degrees of freedom.  The active v23 DE gate already selects
    # flat LambdaCDM (Omega_k=0, w0=-1, wa=0), leaving the three background
    # coordinates (h, Omega_m, Omega_b) to be minimized here.
    base_params_path = os.path.join(dataset_dir, "cosmo_params.json")
    try:
        with open(base_params_path, "r", encoding="utf-8") as f:
            base_params = json.load(f)
    except Exception:
        base_params = {}

    fit_info = {"status": "NOT_RUN"}
    cosmo_meta = None
    try:
        cmb_data = load_vector_csv(os.path.join(dataset_dir, "cmb_data.csv"))
        cmb_cov = load_matrix_csv(os.path.join(dataset_dir, "cmb_cov.csv"))
        bao_data = load_vector_csv(os.path.join(dataset_dir, "bao_data_DESI_DR2_FULL.csv"))
        bao_cov = load_matrix_csv(os.path.join(dataset_dir, "bao_cov_DESI_DR2_FULL.csv"))
        inv_cmb = np.linalg.inv(cmb_cov)
        inv_bao = np.linalg.inv(bao_cov)

        h0 = float(base_params.get("H0", 67.4)) / 100.0
        obh2_0 = float(base_params.get("ombh2", 0.02235976))
        och2_0 = float(base_params.get("omch2", 0.1200))
        om0 = (obh2_0 + och2_0) / (h0 * h0)
        ob0 = obh2_0 / (h0 * h0)
        x0 = np.array([h0, om0, ob0], dtype=float)

        def params_from_x(x):
            h, om, ob = [float(v) for v in x]
            pfit = dict(base_params)
            pfit.update({
                "H0": 100.0 * h,
                "ombh2": ob * h * h,
                "omch2": (om - ob) * h * h,
                "ok": 0.0,
                "w0": -1.0,
                "wa": 0.0,
            })
            return pfit

        def objective(x):
            h, om, ob = [float(v) for v in x]
            # Protected-manuscript hard domain, specialized to the active flat
            # LambdaCDM slice.  Omega_b <= Omega_m is additionally physical.
            if not (0.4 < h < 1.0 and 0.0 <= ob <= om <= 1.0):
                return 1.0e12
            try:
                trial = generate_theory_vectors(params_from_x(x), dataset_dir=dataset_dir)
                rc = cmb_data - trial["cmb"]
                rb = bao_data - trial["bao_desi_dr2_full"]
                return float(rc @ inv_cmb @ rc + rb @ inv_bao @ rb)
            except Exception:
                return 1.0e12

        opt = minimize(
            objective,
            x0,
            method="Nelder-Mead",
            options={"maxiter": 1200, "xatol": 1.0e-10, "fatol": 1.0e-10},
        )
        if not opt.success or not np.isfinite(opt.fun):
            raise RuntimeError(f"Row20 constrained minimization failed: {opt.message}")

        best_params = params_from_x(opt.x)
        gen = generate_theory_vectors(best_params, dataset_dir=dataset_dir)
        write_vector_csv(os.path.join(dataset_dir, "cmb_theory.csv"), gen["cmb"])
        write_vector_csv(os.path.join(dataset_dir, "bao_theory_DV.csv"), gen["bao_dv"])
        write_vector_csv(os.path.join(dataset_dir, "bao_theory_DM_H.csv"), gen["bao_dm_h"])
        write_vector_csv(os.path.join(dataset_dir, "bao_theory_DESI_DR2_FULL.csv"), gen["bao_desi_dr2_full"])
        write_vector_csv(os.path.join(dataset_dir, "sn_theory.csv"), gen["sn"])
        write_vector_csv(os.path.join(dataset_dir, "gw_theory.csv"), gen["gw"])
        cosmo_meta = gen.get("meta", {})
        fit_info = {
            "status": "PASS_CONSTRAINED_MINIMIZATION",
            "method": "Nelder-Mead",
            "fitted_parameters": ["h", "Omega_m", "Omega_b"],
            "n_fitted": 3,
            "fixed_active_slice": {"Omega_k": 0.0, "w0": -1.0, "wa": 0.0},
            "hard_domain": {
                "h": "0.4 < h < 1.0",
                "Omega_m": "0 <= Omega_m <= 1",
                "Omega_b": "0 <= Omega_b <= Omega_m",
            },
            "initial": {"h": float(x0[0]), "Omega_m": float(x0[1]), "Omega_b": float(x0[2])},
            "best_fit": {
                "h": float(opt.x[0]),
                "H0": float(best_params["H0"]),
                "Omega_m": float(opt.x[1]),
                "Omega_b": float(opt.x[2]),
                "ombh2": float(best_params["ombh2"]),
                "omch2": float(best_params["omch2"]),
            },
            "critical_chi2_min_from_optimizer": float(opt.fun),
            "iterations": int(getattr(opt, "nit", 0)),
        }
    except Exception as e:
        fit_info = {"status": "FIT_ERROR", "error": str(e), "n_fitted": 0}
        cosmo_meta = {"generator_error": str(e)}

    def looks_like_template(pth: str) -> bool:
        if not os.path.exists(pth):
            return False
        try:
            with open(pth, "r", encoding="utf-8", errors="ignore") as fh:
                head = fh.read(1024)
            markers = (
                "remove this line after editing",
                "expects numeric rows only",
                "placeholder",
                "template",
                "demo",
                "example",
                "fill me",
                "edit me",
            )
            stripped = head.strip().lower()
            return (stripped == "") or any(m in stripped for m in markers)
        except Exception:
            return False

    def preview(vec: np.ndarray, n: int = 8) -> list[float]:
        arr = np.atleast_1d(vec).reshape(-1)
        return [float(x) for x in arr[: min(n, arr.shape[0])]]

    def diag_sigma(C: np.ndarray) -> np.ndarray:
        d = np.diag(C).astype(float)
        d = np.maximum(d, 0.0)
        return np.sqrt(d)

    def safe_inv(C: np.ndarray) -> tuple[np.ndarray, str]:
        try:
            return np.linalg.inv(C), "inv"
        except np.linalg.LinAlgError:
            return np.linalg.pinv(C), "pinv"

    def whitened(residual: np.ndarray, C: np.ndarray) -> np.ndarray:
        try:
            L = np.linalg.cholesky(C)
            return np.linalg.solve(L, residual)
        except np.linalg.LinAlgError:
            invC, _ = safe_inv(C)
            return invC @ residual

    def bao_dm_h_covariance_bridge(d: np.ndarray, t: np.ndarray, C: np.ndarray) -> dict:
        """
        v23 BAO covariance audit.

        The official score uses the unmodified correlated D_M/H residual. The
        older NZ/CFGL residual correction is recomputed only as a preserved
        historical diagnostic so it cannot improve the Critical Row 20 score.
        """

        d = np.asarray(d, dtype=float).reshape(-1)
        t = np.asarray(t, dtype=float).reshape(-1)
        C = np.asarray(C, dtype=float)

        if d.shape != t.shape:
            raise ValueError(f"BAO_DM_H shape mismatch: data {d.shape} vs theory {t.shape}.")

        n = int(d.shape[0])
        if n != 4:
            raise ValueError(
                "BAO_DM_H v22 bridge expects four entries ordered as "
                "[D_M(0.510), H(0.510), D_M(0.706), H(0.706)]."
            )

        if C.shape != (n, n):
            raise ValueError(f"BAO_DM_H covariance shape mismatch: expected {(n, n)} got {C.shape}.")

        residual = d - t
        invC, inverse_method = safe_inv(C)

        original_chi2 = float(residual @ invC @ residual)
        original_white = whitened(residual, C)

        idx_dm706 = 2
        idx_h706 = 3

        cov_dm_h = float(C[idx_dm706, idx_h706])
        var_h = float(C[idx_h706, idx_h706])
        var_dm = float(C[idx_dm706, idx_dm706])

        alpha = cov_dm_h / var_h if abs(var_h) > 1.0e-15 else 0.0
        h_sigma = math.sqrt(max(var_h, 0.0))

        # Native bridge source:
        # This is the same operator value used by the diagnostic wall test:
        # transformed r_H_BAO after the first NZ/CFGL bridge operator.
        #
        # It is anchored to the existing BAO residual geometry:
        # b_star = |r_D_M_BAO(z=0.706)|
        # r_H_BAO_prime = r_H_BAO - b_star
        #
        # No new parameter is introduced.
        sigma = diag_sigma(C)
        r_h_bao = residual[idx_h706] / sigma[idx_h706] if sigma[idx_h706] > 0 else 0.0
        r_dm_bao = residual[idx_dm706] / sigma[idx_dm706] if sigma[idx_dm706] > 0 else 0.0
        b_star = abs(r_dm_bao)
        bao_h_after_first = r_h_bao - b_star

        h_delta_data_units = bao_h_after_first * h_sigma

        correction = np.zeros_like(residual)
        correction[idx_h706] = -h_delta_data_units
        correction[idx_dm706] = -alpha * h_delta_data_units

        bridge_residual = residual + correction
        bridge_chi2 = float(bridge_residual @ invC @ bridge_residual)
        bridge_white = whitened(bridge_residual, C)

        chi2_delta = float(bridge_chi2 - original_chi2)

        pair_cov = C[np.ix_([idx_dm706, idx_h706], [idx_dm706, idx_h706])]
        pair_inv, pair_inverse_method = safe_inv(pair_cov)

        original_pair = residual[[idx_dm706, idx_h706]]
        bridge_pair = bridge_residual[[idx_dm706, idx_h706]]

        original_pair_chi2 = float(original_pair @ pair_inv @ original_pair)
        bridge_pair_chi2 = float(bridge_pair @ pair_inv @ bridge_pair)
        pair_chi2_delta = float(bridge_pair_chi2 - original_pair_chi2)

        denom = math.sqrt(max(var_dm, 0.0) * max(var_h, 0.0))
        cov_corr = cov_dm_h / denom if denom > 0.0 else float("nan")

        tests = {
            "covariance_loaded": bool(C.shape == (4, 4)),
            "z0706_pair_correlation_nonzero": bool(abs(cov_corr) > 1.0e-6),
            "bridge_reduces_H_white_component": bool(abs(bridge_white[idx_h706]) < abs(original_white[idx_h706])),
            "full_BAO_chi2_not_worse": bool(bridge_chi2 <= original_chi2),
            "pair_chi2_not_worse": bool(bridge_pair_chi2 <= original_pair_chi2),
        }

        score = sum(1 for v in tests.values() if v)
        max_score = len(tests)

        if tests["full_BAO_chi2_not_worse"] and tests["pair_chi2_not_worse"]:
            bridge_status = "BAO_H_wall_covariance_artifact_candidate_promoted_to_bridge"
        elif tests["bridge_reduces_H_white_component"] and pair_chi2_delta < 1.0:
            bridge_status = "BAO_H_wall_covariance_softened_promoted_to_bridge"
        else:
            bridge_status = "BAO_H_wall_covariance_real_no_bridge_improvement"

        pair_details = []
        for i in range(n // 2):
            j = 2 * i
            d_pair = d[j:j + 2]
            t_pair = t[j:j + 2]
            r_pair_original = residual[j:j + 2]
            r_pair_bridge = bridge_residual[j:j + 2]
            C_pair = C[j:j + 2, j:j + 2]
            inv_pair, inv_pair_method = safe_inv(C_pair)

            sig_dm = float(np.sqrt(max(C_pair[0, 0], 0.0)))
            sig_h = float(np.sqrt(max(C_pair[1, 1], 0.0)))
            denom_pair = sig_dm * sig_h
            rho = float(C_pair[0, 1] / denom_pair) if denom_pair > 0.0 else float("nan")

            pair_details.append({
                "pair_index": i,
                "redshift": 0.510 if i == 0 else 0.706,
                "vector_order": ["D_M", "H"],
                "data": preview(d_pair, 2),
                "theory": preview(t_pair, 2),
                "original_residual": preview(r_pair_original, 2),
                "bridge_residual": preview(r_pair_bridge, 2),
                "delta_residual": preview(correction[j:j + 2], 2),
                "sigma": [sig_dm, sig_h],
                "rho_DM_H": rho,
                "original_pair_chi2": float(r_pair_original @ inv_pair @ r_pair_original),
                "bridge_pair_chi2": float(r_pair_bridge @ inv_pair @ r_pair_bridge),
                "inverse_method": inv_pair_method,
            })

        # Explicit local map to the official DESI-like observable space
        # (D_M/r_d^fid, D_H/r_d^fid) with D_H=c/H. This derives the nonlinear
        # observable map and the first-order covariance Jacobian instead of
        # treating H and D_H interchangeably.
        rd_fid = 147.78
        c_light = 299792.458
        official_data = np.array([
            d[0] / rd_fid, c_light / (d[1] * rd_fid),
            d[2] / rd_fid, c_light / (d[3] * rd_fid),
        ])
        official_theory = np.array([
            t[0] / rd_fid, c_light / (t[1] * rd_fid),
            t[2] / rd_fid, c_light / (t[3] * rd_fid),
        ])
        J = np.diag([
            1.0 / rd_fid, -c_light / (d[1] ** 2 * rd_fid),
            1.0 / rd_fid, -c_light / (d[3] ** 2 * rd_fid),
        ])
        official_cov = J @ C @ J.T
        J_inv = np.linalg.inv(J)
        covariance_roundtrip = J_inv @ official_cov @ J_inv.T
        covariance_roundtrip_residual = float(np.max(np.abs(covariance_roundtrip - C)))

        return {
            "chi2": original_chi2,
            "n": n,
            "n_eff": n,
            "residual": residual,
            "residual_preview": preview(residual),
            "pull_diag_preview": preview(
                np.divide(
                    residual,
                    sigma,
                    out=np.full_like(residual, np.nan),
                    where=sigma > 0,
                )
            ),
            "bridge_diagnostic": {
                "status": "HISTORICAL_DIAGNOSTIC_ONLY__" + bridge_status,
                "official_chi2_uses_covariance_bridge": False,
                "original_chi2": original_chi2,
                "covariance_bridge_chi2": bridge_chi2,
                "delta_chi2": chi2_delta,
                "original_pair_z0706_chi2": original_pair_chi2,
                "covariance_bridge_pair_z0706_chi2": bridge_pair_chi2,
                "pair_z0706_delta_chi2": pair_chi2_delta,
                "covariance_correlation_DM_H_z0706": float(cov_corr),
                "alpha_DM_compensation": float(alpha),
                "r_H_BAO_original_sigma": float(r_h_bao),
                "r_D_M_BAO_original_sigma": float(r_dm_bao),
                "b_star": float(b_star),
                "r_H_BAO_after_NZ_CFGL_bridge_sigma": float(bao_h_after_first),
                "delta_residual": correction.tolist(),
                "original_residual": residual.tolist(),
                "covariance_bridge_residual": bridge_residual.tolist(),
                "original_whitened_residual": original_white.tolist(),
                "covariance_bridge_whitened_residual": bridge_white.tolist(),
                "inverse_method": inverse_method,
                "pair_inverse_method": pair_inverse_method,
                "tests": tests,
                "score": int(score),
                "max_score": int(max_score),
                "pair_details": pair_details,
                "official_observable_space": {
                    "order": ["DM_over_rd_z0p510", "DH_over_rd_z0p510", "DM_over_rd_z0p706", "DH_over_rd_z0p706"],
                    "rd_fid_Mpc": rd_fid,
                    "data": official_data.tolist(),
                    "theory": official_theory.tolist(),
                    "covariance": official_cov.tolist(),
                    "jacobian": J.tolist(),
                    "covariance_roundtrip_max_abs_residual": covariance_roundtrip_residual,
                    "nonlinear_relation": "D_H/r_d = c/(H*r_d)",
                    "status": "DERIVED_LOCAL_JACOBIAN_MAP",
                },
                "notes": (
                    "The b_star=|r_D_M| NZ/CFGL residual correction is preserved for historical "
                    "diagnostic comparison only. It is not used in the v23 Critical Row 20 score."
                ),
            },
        }

    def add_part(
        tag: str,
        data_fn: str,
        theory_fn: str,
        cov_fn: str,
        required: bool = True,
        include_in_total: bool = True,
    ) -> None:
        nonlocal total_chi2, total_n

        data_path = os.path.join(dataset_dir, data_fn)
        theory_path = os.path.join(dataset_dir, theory_fn)
        cov_path = os.path.join(dataset_dir, cov_fn) if cov_fn else None

        missing_files = [
            name for name, path in [
                ("data", data_path),
                ("theory", theory_path),
                ("cov", cov_path),
            ]
            if path and not os.path.exists(path)
        ]

        if missing_files:
            parts.append({
                "tag": tag,
                "required": required,
                "include_in_total": include_in_total,
                "status": "PENDING_DATA",
                "files": {"data": data_fn, "theory": theory_fn, "cov": cov_fn},
                "notes": f"Missing file(s): {', '.join(missing_files)}.",
            })
            return

        pmeta = provenance.get(tag, {}) or {}
        if pmeta.get("verified", False) is not True:
            parts.append({
                "tag": tag,
                "required": required,
                "include_in_total": include_in_total,
                "status": "PENDING_DATA",
                "files": {"data": data_fn, "theory": theory_fn, "cov": cov_fn},
                "notes": "Dataset provenance not verified for LIVE use.",
            })
            return

        if looks_like_template(data_path) or looks_like_template(theory_path) or (cov_path and looks_like_template(cov_path)):
            parts.append({
                "tag": tag,
                "required": required,
                "include_in_total": include_in_total,
                "status": "PENDING_DATA",
                "files": {"data": data_fn, "theory": theory_fn, "cov": cov_fn},
                "notes": "Template/placeholder detected.",
            })
            return

        try:
            d = load_vector_csv(data_path)
            t = load_vector_csv(theory_path)
            C = load_matrix_csv(cov_path)
        except Exception as e:
            parts.append({
                "tag": tag,
                "required": required,
                "include_in_total": include_in_total,
                "status": "PENDING_DATA",
                "files": {"data": data_fn, "theory": theory_fn, "cov": cov_fn},
                "notes": f"Load error: {e}",
            })
            return

        if d.shape != t.shape:
            parts.append({
                "tag": tag,
                "required": required,
                "include_in_total": include_in_total,
                "status": "PENDING_DATA",
                "files": {"data": data_fn, "theory": theory_fn, "cov": cov_fn},
                "notes": f"Shape mismatch: data {d.shape} vs theory {t.shape}.",
            })
            return

        n = int(d.shape[0])

        if C.shape != (n, n):
            parts.append({
                "tag": tag,
                "required": required,
                "include_in_total": include_in_total,
                "status": "PENDING_DATA",
                "files": {"data": data_fn, "theory": theory_fn, "cov": cov_fn},
                "notes": f"Covariance shape mismatch: expected {(n, n)} got {C.shape}.",
            })
            return

        if (not np.all(np.isfinite(d))) or (not np.all(np.isfinite(t))) or (not np.all(np.isfinite(C))):
            parts.append({
                "tag": tag,
                "required": required,
                "include_in_total": include_in_total,
                "status": "PENDING_DATA",
                "files": {"data": data_fn, "theory": theory_fn, "cov": cov_fn},
                "notes": "Non-finite values detected.",
            })
            return

        if np.allclose(d, t, rtol=0.0, atol=0.0):
            parts.append({
                "tag": tag,
                "required": required,
                "include_in_total": include_in_total,
                "status": "PENDING_DATA",
                "files": {"data": data_fn, "theory": theory_fn, "cov": cov_fn},
                "notes": "Theory identical to data.",
            })
            return

        try:
            diag = cov_diagnostics(C)

            part = {
                "tag": tag,
                "required": required,
                "include_in_total": include_in_total,
                "status": "RUN",
                "n": n,
                "cov": diag,
                "files": {"data": data_fn, "theory": theory_fn, "cov": cov_fn},
                "data_preview": preview(d),
                "theory_preview": preview(t),
            }

            if tag == "SN":
                invC = np.linalg.inv(C)
                raw_resid = d - t
                z_sn = np.asarray((cosmo_meta or {}).get("sn_z", []), dtype=float).reshape(-1)
                if z_sn.shape != d.shape:
                    raise ValueError(f"SN redshift vector shape mismatch: z {z_sn.shape} vs data {d.shape}.")
                pivot = float((cosmo_meta or {}).get("sn_mu_pivot", np.mean(z_sn)))
                X = np.column_stack([np.ones_like(z_sn), z_sn - pivot])
                normal = X.T @ invC @ X
                rhs = X.T @ invC @ raw_resid
                try:
                    nuisance = np.linalg.solve(normal, rhs)
                    nuisance_method = "solve"
                except np.linalg.LinAlgError:
                    nuisance = np.linalg.pinv(normal) @ rhs
                    nuisance_method = "pinv"
                nuisance_rank = int(np.linalg.matrix_rank(X))
                shifted_t = t + X @ nuisance
                post_resid = d - shifted_t
                x2 = float(post_resid @ invC @ post_resid)
                n_eff = max(n - nuisance_rank, 1)

                sig = diag_sigma(C)
                raw_pull = np.divide(raw_resid, sig, out=np.full_like(raw_resid, np.nan), where=sig > 0)
                post_pull = np.divide(post_resid, sig, out=np.full_like(post_resid, np.nan), where=sig > 0)

                part.update({
                    "chi2": x2,
                    "chi2_per_dof": x2 / max(n_eff, 1),
                    "n_eff": n_eff,
                    "notes": f"SN GLS nuisance profile: intercept={nuisance[0]:.6f}, slope={nuisance[1]:.6f}; rank={nuisance_rank}",
                    "sn_diagnostic": {
                        "nuisance_model": "Delta_mu(z)=M+s*(z-z_pivot)",
                        "nuisance_fit_method": nuisance_method,
                        "nuisance_rank_charged_to_dof": nuisance_rank,
                        "M_fit": float(nuisance[0]),
                        "slope_fit": float(nuisance[1]),
                        "pivot_z": pivot,
                        "legacy_sweep_selected_slope_ignored": (cosmo_meta or {}).get("legacy_sn_mu_slope_ignored"),
                        "raw_residual_preview": preview(raw_resid),
                        "post_nuisance_residual_preview": preview(post_resid),
                        "raw_pull_diag_preview": preview(raw_pull),
                        "post_nuisance_pull_diag_preview": preview(post_pull),
                        "theory_shifted_preview": preview(shifted_t),
                        "mean_raw_abs_residual": float(np.mean(np.abs(raw_resid))),
                        "mean_post_abs_residual": float(np.mean(np.abs(post_resid))),
                    },
                })

            elif tag == "BAO_DM_H":
                bridge = bao_dm_h_covariance_bridge(d, t, C)
                x2 = float(bridge["chi2"])
                n_eff = int(bridge["n_eff"])

                part.update({
                    "chi2": x2,
                    "chi2_per_dof": x2 / max(n_eff, 1),
                    "n_eff": n_eff,
                    "residual_preview": bridge["residual_preview"],
                    "pull_diag_preview": bridge["pull_diag_preview"],
                    "bao_dm_h_bridge": bridge["bridge_diagnostic"],
                    "notes": "v23 direct correlated BAO score; legacy NZ/CFGL residual correction retained only as diagnostic.",
                })

            else:
                x2 = float(chi2(d, t, C))
                n_eff = n
                resid = d - t
                sig = diag_sigma(C)
                pull = np.divide(resid, sig, out=np.full_like(resid, np.nan), where=sig > 0)

                part.update({
                    "chi2": x2,
                    "chi2_per_dof": x2 / max(n_eff, 1),
                    "n_eff": n_eff,
                    "residual_preview": preview(resid),
                    "pull_diag_preview": preview(pull),
                })

        except Exception as e:
            parts.append({
                "tag": tag,
                "required": required,
                "include_in_total": include_in_total,
                "status": "PENDING_DATA",
                "files": {"data": data_fn, "theory": theory_fn, "cov": cov_fn},
                "notes": f"chi2/covariance error: {e}",
            })
            return

        if include_in_total:
            total_chi2 += x2
            total_n += n_eff

        parts.append(part)

    add_part("CMB", "cmb_data.csv", "cmb_theory.csv", "cmb_cov.csv", required=True, include_in_total=True)
    # v23 completion: score the full official DESI DR2 Gaussian BAO likelihood
    # once, in its native DV/rd, DM/rd, DH/rd basis.  The historical low-z
    # transformed slices remain executable below as diagnostics only so no
    # measurement is double-counted in the Critical score.
    add_part("BAO_DESI_DR2_FULL", "bao_data_DESI_DR2_FULL.csv", "bao_theory_DESI_DR2_FULL.csv", "bao_cov_DESI_DR2_FULL.csv", required=True, include_in_total=True)
    add_part("BAO_DV", "bao_data_DV.csv", "bao_theory_DV.csv", "bao_cov_DV.csv", required=False, include_in_total=False)
    add_part("BAO_DM_H", "bao_data_DM_H.csv", "bao_theory_DM_H.csv", "bao_cov_DM_H.csv", required=False, include_in_total=False)
    # SN and GW files in this bundle are harness-compatible stress/proxy datasets,
    # not independent observational likelihoods.  Preserve and execute them as
    # diagnostics, but do not allow their dof to dilute the Critical CMB+BAO score.
    add_part("SN", "sn_data.csv", "sn_theory.csv", "sn_cov.csv", required=False, include_in_total=False)
    add_part("GW", "gw_data.csv", "gw_theory.csv", "gw_cov.csv", required=False, include_in_total=False)

    for p in parts:
        if p.get("tag") in {"SN", "GW"}:
            p["evidence_role"] = "HARNESS_STRESS_OR_PROXY_DIAGNOSTIC__EXCLUDED_FROM_CRITICAL_OBSERVATIONAL_SCORE"
        elif p.get("tag") in {"BAO_DV", "BAO_DM_H"}:
            p["evidence_role"] = "LEGACY_LOW_Z_BAO_DIAGNOSTIC__EXCLUDED_TO_AVOID_DOUBLE_COUNTING_FULL_DR2"
        else:
            p["evidence_role"] = "CRITICAL_OBSERVATIONAL_COMPRESSION"

    n_fitted = int(fit_info.get("n_fitted", 0)) if fit_info.get("status") == "PASS_CONSTRAINED_MINIMIZATION" else 0
    effective_dof = total_n - n_fitted
    chi2_per_dof = total_chi2 / max(effective_dof, 1)

    cov_ok = all(
        (p.get("status") != "RUN") or (p.get("cov", {}).get("min_eig", 0.0) > 0.0)
        for p in parts
        if p.get("include_in_total", True)
    )

    threshold_chi2_per_dof = 1.25
    pass_rule = (chi2_per_dof <= threshold_chi2_per_dof) and cov_ok

    missing_required = [
        p["tag"]
        for p in parts
        if p.get("status") == "PENDING_DATA" and p.get("required")
    ]

    verdict = "PENDING_DATA" if missing_required else ("PASS" if pass_rule else "FAIL")

    timestamp = datetime.datetime.now(datetime.UTC).isoformat(timespec="seconds").replace("+00:00", "Z")

    to_hash = [
        "cosmo_params.json",
        "cosmology_provenance.json",
        "row20_background_provenance.json",
        "cmb_data.csv", "cmb_theory.csv", "cmb_cov.csv",
        "bao_data_DV.csv", "bao_theory_DV.csv", "bao_cov_DV.csv",
        "bao_data_DM_H.csv", "bao_theory_DM_H.csv", "bao_cov_DM_H.csv",
        "bao_data_DESI_DR2_FULL.csv", "bao_theory_DESI_DR2_FULL.csv", "bao_cov_DESI_DR2_FULL.csv",
        "desi_dr2_bao_provenance.json",
        "sn_data.csv", "sn_theory.csv", "sn_cov.csv",
        "gw_data.csv", "gw_theory.csv", "gw_cov.csv",
    ]

    bao_audit = {}
    for p in parts:
        if p.get("tag") == "BAO_DM_H" and "bao_dm_h_bridge" in p:
            bridge = p["bao_dm_h_bridge"]
            bao_audit = {
                "original_BAO_DM_H_chi2": bridge.get("original_chi2"),
                "covariance_bridge_BAO_DM_H_chi2": bridge.get("covariance_bridge_chi2"),
                "BAO_DM_H_delta_chi2": bridge.get("delta_chi2"),
                "bridge_status": bridge.get("status"),
                "bridge_score": bridge.get("score"),
                "bridge_max_score": bridge.get("max_score"),
            }
            break

    direct_total_chi2 = total_chi2
    diagnostic_all_blocks_chi2 = float(sum(float(p.get("chi2", 0.0)) for p in parts if p.get("status") == "RUN"))
    diagnostic_all_blocks_n = int(sum(int(p.get("n_eff", 0)) for p in parts if p.get("status") == "RUN"))
    diagnostic_all_blocks_chi2_per_dof = diagnostic_all_blocks_chi2 / max(diagnostic_all_blocks_n, 1)
    legacy_bridge_total_chi2 = total_chi2
    if bao_audit:
        legacy_bridge_total_chi2 = (
            total_chi2
            - float(bao_audit["original_BAO_DM_H_chi2"])
            + float(bao_audit["covariance_bridge_BAO_DM_H_chi2"])
        )

    return {
        "row": 20,
        "signature": "ROW20_COSMOLOGY_GATE_v7_CONSTRAINED_LCDM_FULL_DESI_DR2_PRECISION_RD",
        "gate": "compressed_cosmology_fit",
        "context": "LIVE",
        "tier": "Critical",
        "timestamp_utc": timestamp,
        "summary": {
            "chi2_total": total_chi2,
            "n_total": total_n,
            "n_fitted": n_fitted,
            "effective_dof": effective_dof,
            "fit": fit_info,
            "chi2_per_dof": chi2_per_dof,
            "cosmo_meta": cosmo_meta,
            "covariance_all_pd": cov_ok,
            "threshold_chi2_per_dof": threshold_chi2_per_dof,
            "missing_required": missing_required,
            "mode": "Constrained flat-LambdaCDM fit to CMB + full official DESI DR2 native Gaussian BAO likelihood; three fitted background dof charged; legacy low-z transformed BAO and SN/GW blocks diagnostic-only",
            "critical_evidence_blocks": ["CMB", "BAO_DESI_DR2_FULL"],
            "diagnostic_only_blocks": ["BAO_DV", "BAO_DM_H", "SN", "GW"],
            "diagnostic_all_blocks_chi2": diagnostic_all_blocks_chi2,
            "diagnostic_all_blocks_n_eff": diagnostic_all_blocks_n,
            "diagnostic_all_blocks_chi2_per_dof": diagnostic_all_blocks_chi2_per_dof,
            "v22_bridge": {
                "BAO_DM_H": "historical residual-derived NZ/CFGL correction retained only as diagnostic",
                "official_v23_total_chi2": float(direct_total_chi2),
                "historical_bridge_total_chi2": float(legacy_bridge_total_chi2),
                "historical_bridge_delta_chi2": float(legacy_bridge_total_chi2 - direct_total_chi2),
                "physics_changed": False,
                "thresholds_changed": False,
                "datasets_changed": False,
                "new_free_parameters": False,
                "audit": bao_audit,
            },
        },
        "parts": parts,
        "verdict": verdict,
        "blocking": True,
        "inputs": {
            "dataset_dir": os.path.basename(dataset_dir),
            "hashes": {
                fn: sha256_file(os.path.join(dataset_dir, fn))
                for fn in to_hash
                if os.path.exists(os.path.join(dataset_dir, fn))
            },
        },
    }


if __name__ == "__main__":
    result = gate_cosmo("datasets")

    out_dir = os.path.join("results", "row20_cosmology")
    os.makedirs(out_dir, exist_ok=True)

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

    print("\n==============================")
    print(f"ROW 20 - {result['verdict']}")
    print(f"chi2/dof: {result['summary']['chi2_per_dof']:.6f}")
    print("==============================\n")

    print(json.dumps(result, indent=2))
    print(f"\nWrote {out_path}")