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 tools.row20_h0_composite_release_test import composite_h0_release_test


def h0_cross_coupling_test(dataset_dir: str = "datasets") -> dict:
    base = composite_h0_release_test(dataset_dir)
    s = base["summary"]

    h0_early = float(s["h0_early_anchor"])
    h0_local = float(s["h0_local_target"])

    bao = float(s["bao_bridge_shift"])
    cmb = float(s["cmb_calibration_shift"])
    gw = float(s["gw_proxy_shift"])

    raw_shift = float(s["raw_composite_shift"])
    needed_shift = h0_local - h0_early
    missing_after_raw = h0_local - (h0_early + raw_shift)

    # Conservative no-double-counting norm
    orthogonal_shift = math.sqrt(bao**2 + cmb**2 + gw**2)
    orthogonal_h0 = h0_early + orthogonal_shift

    # Cross-coupling probe:
    # Uses geometric mean interaction terms, not fitted weights.
    bao_cmb = math.sqrt(abs(bao * cmb))
    bao_gw = math.sqrt(abs(bao * gw))
    cmb_gw = math.sqrt(abs(cmb * gw))

    total_cross = bao_cmb + bao_gw + cmb_gw

    # Conservative coupled candidate:
    # orthogonal base + normalized cross support.
    # Normalization by number of channels avoids raw double-counting.
    normalized_cross = total_cross / 3.0

    coupled_shift = orthogonal_shift + normalized_cross
    coupled_h0 = h0_early + coupled_shift
    coupled_remaining_gap = h0_local - coupled_h0

    within_local_tolerance = abs(coupled_h0 - h0_local) <= 1.5

    if within_local_tolerance:
        verdict = "CANDIDATE"
        status = "H0_CROSS_COUPLING_CANDIDATE"
    elif coupled_h0 > orthogonal_h0:
        verdict = "PARTIAL"
        status = "CROSS_COUPLING_PRESENT_BUT_INSUFFICIENT"
    else:
        verdict = "CANNOT_SOLVE"
        status = "NO_CROSS_COUPLING_GAIN"

    return {
        "test": "UTFANSWF_v22_H0_CROSS_COUPLING_TEST",
        "source_test": base["test"],
        "summary": {
            "h0_early_anchor": h0_early,
            "h0_local_target": h0_local,
            "needed_shift": needed_shift,

            "bao_shift": bao,
            "cmb_shift": cmb,
            "gw_shift": gw,

            "raw_additive_shift": raw_shift,
            "raw_additive_h0": h0_early + raw_shift,
            "raw_remaining_gap": missing_after_raw,

            "orthogonal_no_double_counting_shift": orthogonal_shift,
            "orthogonal_no_double_counting_h0": orthogonal_h0,

            "bao_cmb_cross": bao_cmb,
            "bao_gw_cross": bao_gw,
            "cmb_gw_cross": cmb_gw,
            "total_cross_support": total_cross,
            "normalized_cross_support": normalized_cross,

            "coupled_shift": coupled_shift,
            "coupled_h0": coupled_h0,
            "coupled_remaining_gap": coupled_remaining_gap,

            "within_local_tolerance": within_local_tolerance,
            "status": status,

            "rules": {
                "no_fitted_weights": True,
                "physics_changed": False,
                "thresholds_changed": False,
                "datasets_changed": False,
                "new_free_parameters": False,
            },
        },
        "interpretation": {
            "main_point": (
                "This probes whether the missing H0 shift appears through BAO-CMB-GW "
                "cross-coupling geometry rather than raw additive stacking."
            ),
            "warning": (
                "This is still a probe. A CANDIDATE result must be promoted only if the "
                "cross terms can be derived from UTFANSWF operators rather than treated "
                "as numerical convenience."
            ),
        },
        "verdict": verdict,
    }


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

    print("\n================================")
    print(f"H0 CROSS-COUPLING TEST — {result['verdict']}")
    print(f"status: {result['summary']['status']}")
    print(f"coupled H0: {result['summary']['coupled_h0']:.6f}")
    print(f"remaining gap: {result['summary']['coupled_remaining_gap']:.6f}")
    print("================================\n")

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