import csv
import json
import math
import os
from pathlib import Path

ROOT = Path(".")
DATASETS = ROOT / "datasets"
HARNESS = ROOT / "harness"

INPUT_JSON = ROOT / "results" / "row20_h0_attribution" / "row20_thin_slice_attribution.json"
OUTPUT_JSON = ROOT / "results" / "row20_h0_attribution" / "row20_lA_origin_trace.json"
OUTPUT_REPORT = ROOT / "results" / "row20_h0_attribution" / "row20_lA_origin_trace_REPORT.md"


SEARCH_TERMS = ["ell_A", "theta", "theta_star", "theta*", "acoustic", "lA", "l_A"]


def read_text_safe(path):
    try:
        return path.read_text(encoding="utf-8")
    except UnicodeDecodeError:
        return path.read_text(encoding="cp1252", errors="replace")
    except Exception as e:
        return f"<<READ ERROR: {e}>>"


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


def find_block(data, tag):
    for block in data.get("blocks", []):
        if block.get("tag") == tag:
            return block
    raise ValueError(f"Block not found: {tag}")


def find_node(block, slice_name):
    for node in block.get("nodes", []):
        if node.get("slice") == slice_name:
            return node
    raise ValueError(f"Slice not found: {slice_name}")


def search_files():
    matches = []

    search_roots = [DATASETS, HARNESS]
    extensions = {".py", ".csv", ".json", ".txt", ".md"}

    for base in search_roots:
        if not base.exists():
            continue

        for path in base.rglob("*"):
            if not path.is_file():
                continue
            if path.suffix.lower() not in extensions:
                continue

            text = read_text_safe(path)
            lines = text.splitlines()

            for idx, line in enumerate(lines, start=1):
                lower = line.lower()
                if any(term.lower() in lower for term in SEARCH_TERMS):
                    matches.append({
                        "file": str(path),
                        "line": idx,
                        "text": line.strip()
                    })

    return matches


def inspect_cmb_csvs():
    cmb_files = []
    for path in DATASETS.glob("*cmb*"):
        if path.is_file():
            cmb_files.append(path)

    inspected = []

    for path in cmb_files:
        item = {
            "file": str(path),
            "rows": []
        }

        try:
            with open(path, "r", encoding="utf-8") as f:
                sample = f.read().splitlines()
        except UnicodeDecodeError:
            with open(path, "r", encoding="cp1252", errors="replace") as f:
                sample = f.read().splitlines()

        for i, line in enumerate(sample[:20], start=1):
            item["rows"].append({"line": i, "text": line})

        inspected.append(item)

    return inspected


def maybe_derive_lA_from_theta(data):
    """
    Tries to detect theta_* or theta from CMB nodes if it exists.
    If theta_* is stored as radians, l_A ≈ pi / theta_*.
    If 100*theta_* is stored, l_A ≈ 100*pi / value.
    This is diagnostic only.
    """
    derived = []

    for block in data.get("blocks", []):
        if block.get("tag") != "CMB":
            continue

        for node in block.get("nodes", []):
            name = node.get("slice", "")
            value_data = node.get("data", None)
            value_theory = node.get("theory", None)

            if "theta" in name.lower():
                for label, val in [("data", value_data), ("theory", value_theory)]:
                    if val is None:
                        continue

                    try:
                        val = float(val)
                    except Exception:
                        continue

                    candidates = []

                    if val != 0:
                        candidates.append({
                            "assumption": "theta_* stored in radians",
                            "source": label,
                            "theta_value": val,
                            "derived_lA": math.pi / val
                        })

                        candidates.append({
                            "assumption": "100*theta_* stored",
                            "source": label,
                            "theta_value": val,
                            "derived_lA": 100.0 * math.pi / val
                        })

                    derived.append({
                        "node": name,
                        "candidates": candidates
                    })

    return derived


def classify_origin(ell_node, matches, derived_lA):
    theory = ell_node["theory"]
    data_val = ell_node["data"]
    residual = ell_node["residual_after_gate_handling"]
    frac_shift = (data_val - theory) / theory

    notes = []

    if matches:
        notes.append("ell_A/acoustic/theta references were found in the harness or datasets.")
    else:
        notes.append("No direct ell_A/acoustic/theta references were found in searched harness/dataset files.")

    if derived_lA:
        notes.append("theta-derived l_A candidates exist and should be compared against stored ell_A theory.")
    else:
        notes.append("No theta_* node was found in the CMB diagnostic block; ell_A appears to be directly stored or computed elsewhere.")

    if abs(frac_shift) < 0.002:
        notes.append("The ell_A offset is small in fractional terms (<0.2%) but high significance due to tight covariance.")

    return {
        "ell_A_data": data_val,
        "ell_A_theory": theory,
        "ell_A_residual": residual,
        "fractional_shift_to_match_data": frac_shift,
        "origin_status": "TRACE_ONLY_NO_CORRECTION",
        "notes": notes
    }


def main():
    if not INPUT_JSON.exists():
        raise FileNotFoundError(f"Missing input JSON: {INPUT_JSON}")

    data = load_json(INPUT_JSON)

    cmb = find_block(data, "CMB")
    ell_node = find_node(cmb, "ell_A")

    matches = search_files()
    cmb_csvs = inspect_cmb_csvs()
    derived_lA = maybe_derive_lA_from_theta(data)
    origin = classify_origin(ell_node, matches, derived_lA)

    result = {
        "signature": "ROW20_LA_ORIGIN_TRACE_v1_NATIVE_DIAGNOSTIC",
        "context": "DIAGNOSTIC_NON_GATING",
        "rules": data.get("rules", {}),
        "ell_A_origin_trace": origin,
        "search_terms": SEARCH_TERMS,
        "matches": matches,
        "cmb_file_samples": cmb_csvs,
        "theta_to_lA_candidates": derived_lA,
        "verdict": (
            "origin trace complete; 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 v21 — Row 20 ell_A Origin Trace\n\n")
        f.write("**Context:** DIAGNOSTIC_NON_GATING. No physics, thresholds, datasets, parameters, or gate verdicts changed.\n\n")

        f.write("## ell_A summary\n\n")
        f.write("| Quantity | Value |\n")
        f.write("|---|---:|\n")
        f.write(f"| data | {origin['ell_A_data']:.9f} |\n")
        f.write(f"| theory | {origin['ell_A_theory']:.9f} |\n")
        f.write(f"| residual | {origin['ell_A_residual']:.9f} |\n")
        f.write(f"| fractional shift to match data | {origin['fractional_shift_to_match_data']:.9e} |\n\n")

        f.write("## Origin notes\n\n")
        for note in origin["notes"]:
            f.write(f"- {note}\n")

        f.write("\n## File references found\n\n")
        if matches:
            f.write("| File | Line | Text |\n")
            f.write("|---|---:|---|\n")
            for m in matches[:50]:
                safe_text = m["text"].replace("|", "\\|")
                f.write(f"| `{m['file']}` | {m['line']} | `{safe_text}` |\n")
        else:
            f.write("No direct references found.\n")

        f.write("\n## CMB file samples\n\n")
        for item in cmb_csvs:
            f.write(f"### `{item['file']}`\n\n")
            f.write("```text\n")
            for row in item["rows"]:
                f.write(f"{row['line']}: {row['text']}\n")
            f.write("```\n\n")

        f.write("## theta-derived l_A candidates\n\n")
        if derived_lA:
            for item in derived_lA:
                f.write(f"### {item['node']}\n\n")
                f.write("| Assumption | Source | theta value | derived l_A |\n")
                f.write("|---|---|---:|---:|\n")
                for c in item["candidates"]:
                    f.write(
                        f"| {c['assumption']} | {c['source']} | "
                        f"{c['theta_value']:.12g} | {c['derived_lA']:.9f} |\n"
                    )
                f.write("\n")
        else:
            f.write("No theta-derived candidates found.\n\n")

        f.write("## Verdict\n\n")
        f.write(result["verdict"] + "\n")

    print("Wrote:")
    print(os.path.abspath(OUTPUT_JSON))
    print(os.path.abspath(OUTPUT_REPORT))
    print()
    print("ell_A origin trace:")
    print(f"data={origin['ell_A_data']:.9f}")
    print(f"theory={origin['ell_A_theory']:.9f}")
    print(f"residual={origin['ell_A_residual']:.9f}")
    print(f"fractional_shift={origin['fractional_shift_to_match_data']:.9e}")
    print(f"matches_found={len(matches)}")
    print(f"theta_candidates_found={len(derived_lA)}")


if __name__ == "__main__":
    main()