import json
import math
import os
from pathlib import Path

import numpy as np

ROOT = Path(".")

DATA_FILE = ROOT / "datasets" / "bao_data_DM_H.csv"
THEORY_FILE = ROOT / "datasets" / "bao_theory_DM_H.csv"
COV_FILE = ROOT / "datasets" / "bao_cov_DM_H.csv"

INPUT_FIRST = ROOT / "results" / "row20_h0_attribution" / "row20_NZ_CFGL_bridge_operator.json"

OUTPUT_JSON = ROOT / "results" / "row20_h0_attribution" / "row20_BAO_covariance_wall_test.json"
OUTPUT_REPORT = ROOT / "results" / "row20_h0_attribution" / "row20_BAO_covariance_wall_test.md"


def load_json(path):
    with open(path, "r", encoding="utf-8") as f:
        return json.load(f)


def load_vector(path):
    return np.loadtxt(path, delimiter=",")


def load_matrix(path):
    return np.loadtxt(path, delimiter=",")


def chi2(residual, cov):
    inv = np.linalg.inv(cov)
    return float(residual.T @ inv @ residual)


def whitened(residual, cov):
    L = np.linalg.cholesky(cov)
    return np.linalg.solve(L, residual)


def main():
    data_vec = load_vector(DATA_FILE)
    theory_vec = load_vector(THEORY_FILE)
    cov = load_matrix(COV_FILE)

    residual = data_vec - theory_vec
    original_chi2 = chi2(residual, cov)
    original_white = whitened(residual, cov)

    labels = [
        "D_M z=0.510",
        "H z=0.510",
        "D_M z=0.706",
        "H z=0.706",
    ]

    first = load_json(INPUT_FIRST)
    first_resid = first["transformed_residual_vector"]

    bao_h_after_first = first_resid["r_H_BAO"]

    d = np.zeros_like(residual)

    idx_dm706 = 2
    idx_h706 = 3

    cov_dm_h = cov[idx_dm706, idx_h706]
    var_h = cov[idx_h706, idx_h706]

    alpha = cov_dm_h / var_h if abs(var_h) > 1e-15 else 0.0

    h_sigma = math.sqrt(var_h)
    h_delta_data_units = bao_h_after_first * h_sigma

    d[idx_h706] = -h_delta_data_units
    d[idx_dm706] = -alpha * h_delta_data_units

    candidate_residual = residual + d

    candidate_chi2 = chi2(candidate_residual, cov)
    candidate_white = whitened(candidate_residual, cov)

    chi2_delta = candidate_chi2 - original_chi2

    pair_cov = cov[np.ix_([idx_dm706, idx_h706], [idx_dm706, idx_h706])]
    original_pair = residual[[idx_dm706, idx_h706]]
    candidate_pair = candidate_residual[[idx_dm706, idx_h706]]

    original_pair_chi2 = chi2(original_pair, pair_cov)
    candidate_pair_chi2 = chi2(candidate_pair, pair_cov)

    pair_chi2_delta = candidate_pair_chi2 - original_pair_chi2

    cov_corr = cov_dm_h / math.sqrt(cov[idx_dm706, idx_dm706] * var_h)

    tests = {
        "covariance_loaded": bool(cov.shape == (4, 4)),
        "z0706_pair_correlation_nonzero": bool(abs(cov_corr) > 1e-6),
        "candidate_reduces_H_white_component": bool(abs(candidate_white[idx_h706]) < abs(original_white[idx_h706])),
        "full_BAO_chi2_not_worse": bool(candidate_chi2 <= original_chi2),
        "pair_chi2_not_worse": bool(candidate_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"]:
        status = "BAO_H_wall_covariance_artifact_candidate"
    elif tests["candidate_reduces_H_white_component"] and pair_chi2_delta < 1.0:
        status = "BAO_H_wall_covariance_softened"
    else:
        status = "BAO_H_wall_covariance_real"

    result = {
        "signature": "ROW20_BAO_COVARIANCE_WALL_TEST_v1",
        "context": "DIAGNOSTIC_NON_GATING",
        "labels": labels,
        "covariance_correlation_DM_H_z0706": float(cov_corr),
        "alpha_DM_compensation": float(alpha),
        "bao_h_after_first_operator_sigma_proxy": float(bao_h_after_first),
        "original": {
            "residual": residual.tolist(),
            "whitened_residual": original_white.tolist(),
            "chi2": float(original_chi2),
            "pair_z0706_chi2": float(original_pair_chi2),
        },
        "candidate": {
            "delta_residual": d.tolist(),
            "residual": candidate_residual.tolist(),
            "whitened_residual": candidate_white.tolist(),
            "chi2": float(candidate_chi2),
            "pair_z0706_chi2": float(candidate_pair_chi2),
        },
        "deltas": {
            "full_chi2_delta": float(chi2_delta),
            "pair_chi2_delta": float(pair_chi2_delta),
        },
        "tests": tests,
        "score": int(score),
        "max_score": int(max_score),
        "verdict": f"{status}; diagnostic only; no physics, thresholds, datasets, parameters, or gate verdicts changed.",
    }

    OUTPUT_JSON.parent.mkdir(parents=True, exist_ok=True)

    with open(OUTPUT_JSON, "w", encoding="utf-8") as f:
        json.dump(result, f, indent=2)

    with open(OUTPUT_REPORT, "w", encoding="utf-8") as f:
        f.write("# UTFANSWF v22 — BAO Covariance Wall Test\n\n")

        f.write("## z=0.706 covariance geometry\n\n")
        f.write(f"corr(D_M,H) = {cov_corr:.6f}\n")
        f.write(f"alpha compensation = {alpha:.6f}\n\n")

        f.write("## chi2\n\n")
        f.write(f"original = {original_chi2:.6f}\n")
        f.write(f"candidate = {candidate_chi2:.6f}\n")
        f.write(f"delta = {chi2_delta:.6f}\n\n")

        f.write("## pair chi2\n\n")
        f.write(f"original = {original_pair_chi2:.6f}\n")
        f.write(f"candidate = {candidate_pair_chi2:.6f}\n")
        f.write(f"delta = {pair_chi2_delta:.6f}\n\n")

        f.write("## verdict\n\n")
        f.write(status + "\n")

    print("BAO covariance wall test:")
    print(f"corr={cov_corr:.6f}")
    print(f"alpha={alpha:.6f}")
    print(f"original chi2={original_chi2:.6f}")
    print(f"candidate chi2={candidate_chi2:.6f}")
    print(f"delta chi2={chi2_delta:.6f}")
    print(f"pair delta chi2={pair_chi2_delta:.6f}")
    print(f"score={score}/{max_score}")
    print(f"status={status}")


if __name__ == "__main__":
    main()