"""Remove graph-proposed placements while retaining native population aliases.

Uses the already rescored evidence and identical inference code, fallback and
gates. This isolates graph retrieval within the existing reference configuration;
it does not rebase that configuration on full-genomic IPD.
"""
import argparse
import gzip
import hashlib
import json
import os
import sys
from pathlib import Path


def sha(path):
    h = hashlib.sha256()
    with path.open("rb") as stream:
        for block in iter(lambda: stream.read(8 * 1024**2), b""):
            h.update(block)
    return h.hexdigest()


def native_only(records):
    for r in records:
        yield dict(r, candidates={k: v for k, v in r["candidates"].items()
                                  if k.startswith("IPD:")})


def records(path):
    with (path.open() if path.exists() else gzip.open(str(path) + ".gz", "rt")) as stream:
        yield from map(json.loads, stream)


def evidence_sha(path):
    """Receipts bind uncompressed contents even when storage was later gzipped."""
    if path.exists():
        return sha(path)
    h = hashlib.sha256()
    with gzip.open(str(path) + ".gz", "rb") as stream:
        for block in iter(lambda: stream.read(8 * 1024**2), b""):
            h.update(block)
    return h.hexdigest()


def main():
    p = argparse.ArgumentParser(description=__doc__)
    for name in ("source", "code", "native-table", "aliases", "output"):
        p.add_argument("--" + name, type=Path, required=True)
    a = p.parse_args()
    if not os.environ.get("SLURM_JOB_ID"):
        raise RuntimeError("Run inference under Slurm")
    if a.output.exists():
        raise FileExistsError(a.output)
    sys.path.insert(0, str(a.code))
    from context_inference import allele_fields, emission_rows, gene_evidence_rows, infer
    from run_context_inference import MODES
    from sequence_identity import core_slot
    from score_native import decode_slots

    def load(path):
        return json.loads(path.read_text())

    evidence = a.source / "evidence"
    complete = load(a.source / "COMPLETE.json")
    assert complete["status"] == "CV_inference_complete"
    assert sha(evidence / "COMPLETE.json") == complete["evidence_manifest_sha256"]
    manifest = load(evidence / "COMPLETE.json")
    assert manifest["status"] == "complete"
    for name in ("fragments.jsonl", "candidates.json"):
        assert evidence_sha(evidence / name) == manifest["output_sha256"][name]
    # Verify native fallback identity, not just the continued presence of a file.
    baseline = load(a.source / "baseline_calls.json")
    for mode in ("E1", "E2", "E3"):
        original = load(a.source / mode / "COMPLETE.json")
        assert sha(a.source / "baseline_calls.json") == original["baseline_sha256"]
        for name, digest in original["code_sha256"].items():
            assert sha(a.code / name) == digest
    meta = load(evidence / "candidates.json")
    aliases = load(a.aliases)
    native_cores = {}
    for line in a.native_table.read_text().splitlines():
        if not line or line.startswith("#"):
            continue
        gene = line.split("\t")[0].removeprefix("HLA-")
        slots = decode_slots(line, aliases)
        native_cores[gene] = [{k: s[k] for k in ("cores", "core_unknown")} for s in slots]
        expected = baseline[gene]["resolutions"]["4"]
        assert [sorted({allele_fields(x, 4) for x in s["labels"]}) for s in slots] == [
            sorted(s["alleles"]) for s in expected]
        assert [s["label_unknown"] for s in slots] == [s["unresolved"] for s in expected]
    native_sequence_hashes = {m["sequence_sha256"] for k, m in meta.items() if k.startswith("IPD:")}
    graph = {k: m for k, m in meta.items() if not k.startswith("IPD:")}
    missing = [k for k, m in graph.items() if m["sequence_sha256"] not in native_sequence_hashes]
    if missing:
        raise ValueError("Graph sequence absent from native catalogue: " + repr(missing))
    a.output.mkdir(parents=True)
    provenance = {
        "source": str(a.source), "source_complete_sha256": sha(a.source / "COMPLETE.json"),
        "source_evidence_sha256": sha(evidence / "COMPLETE.json"),
        "code_sha256": {x.name: sha(x) for x in sorted(a.code.glob("*.py"))},
        "script_sha256": sha(Path(__file__)), "native_table_sha256": sha(a.native_table),
        "aliases_sha256": sha(a.aliases), "job_id": os.environ["SLURM_JOB_ID"],
        "graph_sequences_absent_from_native_catalogue": len(missing),
        "native_panel_aliases_retained": sum(k.startswith("IPD:") and m["source"] == "graph"
                                             for k, m in meta.items()),
        "truth_loaded": False,
        "intervention": "Drop non-IPD-prefixed candidate placements before normalization; retain native panel aliases",
        "limitations": [
            "Existing original.fa/flanked reference configuration, not a full-genomic-IPD rebase.",
            "Tests incremental graph retrieval over native retrieval, not all possible graph algorithms.",
            "Candidate availability is held equal; retrieved candidates may differ by design.",
            "Same scorer, fallback and decision thresholds as original union replay.",
        ],
    }
    counts = {"fragments": 0, "native_supported_fragments": 0, "graph_supported_fragments": 0,
              "graph_only_fragments": 0}
    for r in records(evidence / "fragments.jsonl"):
        n = any(k.startswith("IPD:") and v for k, v in r["candidates"].items())
        g = any(not k.startswith("IPD:") and v for k, v in r["candidates"].items())
        counts["fragments"] += 1
        counts["native_supported_fragments"] += bool(n)
        counts["graph_supported_fragments"] += bool(g)
        counts["graph_only_fragments"] += bool(g and not n)
    provenance["retrieval_counts"] = counts
    outputs = {}
    for mode in ("E1", "E2", "E3"):
        settings = MODES[mode]
        identifiers, rows = emission_rows(native_only(records(evidence / "fragments.jsonl")),
                                          meta, context=settings["context"], include_graph=True)
        core = (gene_evidence_rows(native_only(records(evidence / "fragments.jsonl")), identifiers)
                if settings["context"] else None)
        calls, diagnostics = infer(rows, meta, baseline,
                                   sequence_first=settings["sequence_first"], core_rows=core)
        sequence = {}
        for gene in baseline:
            d = diagnostics["decisions"].get(gene, {})
            selected = d.get("reason") == "sequence_pair_selected"
            slots = ([core_slot(paths, meta) for paths in d["selected_paths"]]
                     if selected else native_cores.get(gene, [{"cores": [], "core_unknown": True}] * 2))
            sequence[gene] = {"slots": slots, "origin": "native_common_score" if selected else "native_fallback"}
        target = a.output / mode
        target.mkdir()
        for name, payload in (("calls.json", calls), ("decisions.json", diagnostics),
                              ("sequence_calls.json", sequence)):
            path = target / name
            path.write_text(json.dumps(payload, indent=2) + "\n")
            outputs[str(path.relative_to(a.output))] = sha(path)
        print(json.dumps({"mode": mode, "fragments": len(identifiers), "complete": True}), flush=True)
    provenance.update(status="native_proposal_ablation_complete", output_sha256=outputs)
    (a.output / "COMPLETE.json").write_text(json.dumps(provenance, indent=2) + "\n")


if __name__ == "__main__":
    main()