from __future__ import annotations

import os
import json
import datetime
import numpy as np

from .gate_row20_cosmology import gate_cosmo


def _find_part(row20: dict, tag: str) -> dict | None:
    for part in row20.get("parts", []):
        if part.get("tag") == tag:
            return part
    return None


def gate_h0_resolution(dataset_dir: str = "datasets") -> dict:
    """
    UTFANSWF v22 — H0 Resolution Gate

    Purpose:
    Determine whether the existing v22 machinery actually resolves the
    H0 tension, or only improves the BAO covariance bridge.

    Rules:
    - No new physics.
    - No new free parameters.
    - No threshold changes.
    - No dataset changes.
    - Uses Row 20 output as source of truth.
    """

    row20 = gate_cosmo(dataset_dir)

    h0_early_anchor = float(
        row20.get("summary", {})
        .get("cosmo_meta", {})
        .get("H0", 67.4)
    )

    h0_local_target = 73.0
    h0_local_tolerance = 1.5

    bao = _find_part(row20, "BAO_DM_H")
    sn = _find_part(row20, "SN")

    if bao is None or bao.get("status") != "RUN":
        verdict = "PENDING_DATA"
        status = "BAO_DM_H unavailable"
        h0_bridge_prediction = h0_early_anchor
        delta_h0_from_bao = 0.0
        remaining_gap_to_local = h0_local_target - h0_bridge_prediction
        evidence = {}
    else:
        bridge = bao.get("bao_dm_h_bridge", {})

        theory = np.array(bao.get("theory_preview", []), dtype=float)
        delta_residual = np.array(
            bridge.get("delta_residual", [0.0, 0.0, 0.0, 0.0]),
            dtype=float,
        )

        # Expected ordering:
        # [D_M(0.510), H(0.510), D_M(0.706), H(0.706)]
        if theory.shape[0] >= 4 and delta_residual.shape[0] >= 4:
            h_theory_z0706 = float(theory[3])
            h_correction_z0706 = float(delta_residual[3])

            # A lower model H(z) residual correction corresponds to an
            # upward inferred H0 movement needed to match local expansion.
            delta_h0_from_bao = float(
                h0_early_anchor * abs(h_correction_z0706) / max(abs(h_theory_z0706), 1.0e-12)
            )
        else:
            h_theory_z0706 = None
            h_correction_z0706 = None
            delta_h0_from_bao = 0.0

        h0_bridge_prediction = float(h0_early_anchor + delta_h0_from_bao)
        remaining_gap_to_local = float(h0_local_target - h0_bridge_prediction)

        sn_diag = sn.get("sn_diagnostic", {}) if sn else {}
        sn_m_fit = sn_diag.get("M_fit")

        evidence = {
            "row20_verdict": row20.get("verdict"),
            "row20_chi2_per_dof": row20.get("summary", {}).get("chi2_per_dof"),
            "bao_bridge_status": bridge.get("status"),
            "bao_original_chi2": bridge.get("original_chi2"),
            "bao_covariance_bridge_chi2": bridge.get("covariance_bridge_chi2"),
            "bao_delta_chi2": bridge.get("delta_chi2"),
            "bao_bridge_score": bridge.get("score"),
            "bao_bridge_max_score": bridge.get("max_score"),
            "h_theory_z0706": h_theory_z0706,
            "h_correction_z0706": h_correction_z0706,
            "sn_M_fit": sn_m_fit,
            "sn_note": (
                "SN nuisance offset is logged as calibration structure only; "
                "it is not treated as a direct H0 solution."
            ),
        }

        row20_pass = row20.get("verdict") == "PASS"
        bridge_valid = (
            bridge.get("score") == bridge.get("max_score")
            and bridge.get("delta_chi2", 0.0) < 0.0
        )

        reaches_local = abs(h0_bridge_prediction - h0_local_target) <= h0_local_tolerance
        stays_near_early = abs(h0_bridge_prediction - h0_early_anchor) <= 1.0

        if row20_pass and bridge_valid and reaches_local:
            verdict = "PASS"
            status = "H0_RESOLUTION_CANDIDATE"
        elif row20_pass and bridge_valid and stays_near_early:
            verdict = "CANNOT_SOLVE"
            status = "H0_NOT_SOLVED_BY_V22_CURRENT_FORM"
        elif row20_pass and bridge_valid:
            verdict = "FAIL"
            status = "H0_PARTIAL_SHIFT_INSUFFICIENT"
        else:
            verdict = "FAIL"
            status = "H0_BRIDGE_NOT_VALIDATED"

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

    return {
        "row": "20-H0",
        "signature": "ROW20_H0_RESOLUTION_GATE_v1",
        "gate": "h0_resolution_test",
        "context": "LIVE",
        "tier": "Critical",
        "timestamp_utc": timestamp,
        "summary": {
            "h0_early_anchor": h0_early_anchor,
            "h0_local_target": h0_local_target,
            "h0_local_tolerance": h0_local_tolerance,
            "h0_bridge_prediction": h0_bridge_prediction,
            "delta_h0_from_bao_bridge": delta_h0_from_bao,
            "remaining_gap_to_local": remaining_gap_to_local,
            "status": status,
            "rules": {
                "physics_changed": False,
                "thresholds_changed": False,
                "datasets_changed": False,
                "new_free_parameters": False,
            },
        },
        "evidence": evidence,
        "verdict": verdict,
        "blocking": False,
    }


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

    print("\n================================")
    print(f"ROW 20-H0 — {result['verdict']}")
    print(f"status: {result['summary']['status']}")
    print(f"H0 bridge prediction: {result['summary']['h0_bridge_prediction']:.6f}")
    print(f"remaining gap to local: {result['summary']['remaining_gap_to_local']:.6f}")
    print("================================\n")

    print(json.dumps(result, indent=2))