import json
import os

INPUT_FILE = "results/row20_h0_attribution/row20_thin_slice_attribution.json"
OUTPUT_FILE = "results/row20_h0_attribution/row20_node_trace.json"

def load_data():
    if not os.path.exists(INPUT_FILE):
        raise FileNotFoundError(f"Missing input file: {INPUT_FILE}")
    with open(INPUT_FILE, "r") as f:
        return json.load(f)

def classify_source(slice_name):
    if "ell_A" in slice_name:
        return "CMB compression / acoustic scale mapping"
    elif "SN mu" in slice_name:
        return "SN calibration / distance-modulus mapping"
    elif "H z=" in slice_name or "D_M z=" in slice_name:
        return "BAO projection / DM-H conversion"
    else:
        return "Unknown / needs deeper inspection"

def analyze_node(node):
    return {
        "block": node["block"],
        "slice": node["slice"],
        "delta_chi2_removed": node["delta_chi2_removed"],
        "chi2_per_dof_without_slice": node["chi2_per_dof_without_slice"],
        "likely_source": classify_source(node["slice"])
    }

def main():
    data = load_data()

    nodes = data.get("thin_slice_targets", [])

    if not nodes:
        print("No thin_slice_targets found in JSON. Check structure.")
        return

    results = []
    for node in nodes:
        results.append(analyze_node(node))

    os.makedirs(os.path.dirname(OUTPUT_FILE), exist_ok=True)

    with open(OUTPUT_FILE, "w") as f:
        json.dump(results, f, indent=2)

    print("\nNode attribution written to:")
    print(OUTPUT_FILE)

    print("\nTop interpreted sources:")
    for i, r in enumerate(results[:5], start=1):
        print(f"{i}. {r['slice']} → {r['likely_source']}")

if __name__ == "__main__":
    main()