from __future__ import annotations

import json
import math
import sys
from pathlib import Path

ROOT = Path(__file__).resolve().parents[1]
if str(ROOT) not in sys.path:
    sys.path.insert(0, str(ROOT))

from harness.gate_row20_cosmology import gate_cosmo
from tools.row20_h0_composite_release_test import composite_h0_release_test


def find_part(row20: dict, tag: str) -> dict:
    for part in row20.get("parts", []):
        if part.get("tag") == tag:
            return part
    return {}


def h0_sn_geometry_test(dataset_dir: str = "datasets") -> dict:
    row20 = gate_cosmo(dataset_dir)
    composite = composite_h0_release_test(dataset_dir)

    h0_early = float(composite["summary"]["h0_early_anchor"])
    h0_local = float(composite["summary"]["h0_local_target"])

    raw_composite_h0 = float(composite["summary"]["raw_composite_h0"])
    raw_remaining_gap = float(composite["summary"]["raw_remaining_gap"])

    sn = find_part(row20, "SN")
    sn_diag = sn.get("sn_diagnostic", {})

    m_fit = float(sn_diag.get("M_fit", 0.0))
    mean_raw_abs = float(sn_diag.get("mean_raw_abs_residual", 0.0))
    mean_post_abs = float(sn_diag.get("mean_post_abs_residual", 0.0))

    if mean_raw_abs > 1.0e-12:
        geometry_efficiency = 1.0 - (mean_post_abs / mean_raw_abs)
    else:
        geometry_efficiency = 0.0

    # Geometry-only SN contribution:
    # We do NOT treat M_fit as H0.
    # We only allow the fractional residual compression to act on the remaining
    # distributed H0 gap as a projection efficiency.
    sn_geometry_shift = raw_remaining_gap * geometry_efficiency

    sn_geometry_h0 = raw_composite_h0 + sn_geometry_shift
    sn_geometry_remaining_gap = h0_local - sn_geometry_h0

    within_local_tolerance = abs(sn_geometry_h0 - h0_local) <= 1.5

    if within_local_tolerance:
        verdict = "CANDIDATE"
        status = "SN_GEOMETRY_H0_RELEASE_CANDIDATE"
    elif sn_geometry_shift > 0:
        verdict = "PARTIAL"
        status = "SN_GEOMETRY_PRESENT_BUT_INSUFFICIENT"
    else:
        verdict = "CANNOT_SOLVE"
        status = "NO_SN_GEOMETRY_RELEASE_FOUND"

    return {
        "test": "UTFANSWF_v22_H0_SN_GEOMETRY_TEST",
        "row20_signature": row20.get("signature"),
        "source_test": composite.get("test"),
        "summary": {
            "h0_early_anchor": h0_early,
            "h0_local_target": h0_local,
            "raw_composite_h0_before_sn_geometry": raw_composite_h0,
            "raw_remaining_gap_before_sn_geometry": raw_remaining_gap,
            "sn_M_fit_logged_not_used_as_H0": m_fit,
            "sn_mean_raw_abs_residual": mean_raw_abs,
            "sn_mean_post_abs_residual": mean_post_abs,
            "sn_geometry_efficiency": geometry_efficiency,
            "sn_geometry_shift": sn_geometry_shift,
            "sn_geometry_h0": sn_geometry_h0,
            "sn_geometry_remaining_gap": sn_geometry_remaining_gap,
            "within_local_tolerance": within_local_tolerance,
            "status": status,
            "rules": {
                "M_fit_used_directly_as_H0": False,
                "no_fitted_weights": True,
                "physics_changed": False,
                "thresholds_changed": False,
                "datasets_changed": False,
                "new_free_parameters": False,
            },
        },
        "interpretation": {
            "main_point": (
                "This tests whether SN contributes through projection geometry by measuring "
                "how much the nuisance-offset step compresses residual structure, then applying "
                "that compression only to the remaining distributed H0 gap."
            ),
            "warning": (
                "A candidate result still needs derivation in the manuscript. This is a probe, "
                "not yet a final H0 proof."
            ),
        },
        "verdict": verdict,
    }


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

    print("\n================================")
    print(f"H0 SN GEOMETRY TEST — {result['verdict']}")
    print(f"status: {result['summary']['status']}")
    print(f"SN geometry H0: {result['summary']['sn_geometry_h0']:.6f}")
    print(f"remaining gap: {result['summary']['sn_geometry_remaining_gap']:.6f}")
    print("================================\n")

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