#!/usr/bin/env python3
"""Phase 4H stage 1: read-only MaleCNS mushroom-body learning-circuit audit."""
from __future__ import annotations

import argparse
import hashlib
import json
import re
import shutil
import subprocess
import sys
from collections import Counter
from datetime import datetime, timezone
from pathlib import Path
from typing import Any

import numpy as np
import pandas as pd
import pyarrow.parquet as pq
import yaml


def sha256(path: Path, chunk: int = 8 << 20) -> str:
    h = hashlib.sha256()
    with path.open("rb") as handle:
        while data := handle.read(chunk):
            h.update(data)
    return h.hexdigest()


def write_json(path: Path, value: Any) -> None:
    path.parent.mkdir(parents=True, exist_ok=True)
    path.write_text(json.dumps(value, indent=2, sort_keys=True, default=json_default) + "\n", encoding="utf-8")


def json_default(value: Any) -> Any:
    if isinstance(value, np.generic):
        return value.item()
    if isinstance(value, Path):
        return str(value)
    raise TypeError(type(value).__name__)


def parse_compartment(instance: Any) -> tuple[str | None, str]:
    """Accept only one explicit, non-empty parenthesized token; preserve it verbatim."""
    if instance is None or pd.isna(instance):
        return None, "instance_null"
    text = str(instance)
    matches = re.findall(r"\(([^()]*)\)", text)
    nonempty = [item.strip() for item in matches if item.strip()]
    if len(matches) == 1 and len(nonempty) == 1:
        return nonempty[0], "explicit_single_parenthesized_token"
    if not matches:
        return None, "no_parenthesized_token"
    return None, "ambiguous_parenthesized_tokens"


def verify_checksum(root: Path, control: Path) -> dict[str, Any]:
    checked = 0
    failures: list[dict[str, str]] = []
    for line in control.read_text(encoding="utf-8").splitlines():
        if not line.strip():
            continue
        expected, rel = line.split(None, 1)
        path = root / rel.strip().lstrip("*")
        checked += 1
        if not path.is_file():
            failures.append({"path": str(path), "reason": "missing"})
        else:
            actual = sha256(path)
            if actual != expected:
                failures.append({"path": str(path), "reason": "sha256_mismatch", "actual": actual, "expected": expected})
    return {"control": str(control), "root": str(root), "checked": checked, "failures": failures, "passed": not failures}


def frozen_controls(cfg: dict[str, Any]) -> list[tuple[Path, Path]]:
    p = cfg["paths"]
    roots = [
        (Path(p["reference_root"]), Path(p["reference_root"]) / "checksums.sha256"),
        (Path(p["runtime_root"]), Path(p["runtime_root"]) / "checksums.sha256"),
        (Path(p["io_root"]), Path(p["io_root"]) / "checksums.sha256"),
        (Path(p["phase4c_root"]), Path(p["phase4c_root"]) / "manifests/checksums.sha256"),
        (Path(p["phase4d_root"]), Path(p["phase4d_root"]) / "manifests/checksums.sha256"),
        (Path(p["phase4e_root"]), Path(p["phase4e_root"]) / "manifests/checksums.sha256"),
        (Path(p["phase4f_root"]), Path(p["phase4f_root"]) / "manifests/checksums.sha256"),
        (Path(p["phase4f_model_root"]), Path(p["phase4f_model_root"]) / "checksums.sha256"),
        (Path(p["phase4g_root"]), Path(p["phase4g_root"]) / "manifests/checksums.sha256"),
        (Path(p["phase4g_model_root"]), Path(p["phase4g_model_root"]) / "checksums.sha256"),
        (Path(p["phase4g_learned_root"]), Path(p["phase4g_learned_root"]) / "checksums.sha256"),
    ]
    return roots


def run_frozen_audit(cfg: dict[str, Any]) -> dict[str, Any]:
    controls = frozen_controls(cfg)
    missing = [str(control) for _, control in controls if not control.is_file()]
    if missing:
        raise FileNotFoundError(f"frozen checksum controls missing: {missing}")
    checks = [verify_checksum(root, control) for root, control in controls]
    result = {"all_passed": all(check["passed"] for check in checks), "controls": checks}
    if not result["all_passed"]:
        raise RuntimeError("frozen artifact checksum verification failed")
    return result


def graph_paths(cfg: dict[str, Any]) -> dict[str, Path]:
    root = Path(cfg["paths"]["connectome_root"]).resolve()
    allowed = Path("/data/digital-drosophila/datasets/male-cns-v1.0/processed/neuron-only").resolve()
    if root != allowed:
        raise RuntimeError(f"only allowed connectome input is {allowed}")
    paths = {
        "metadata": root / "neuron-index/unified-metadata.parquet",
        "indptr": root / "gpu/csr-indptr.int64.npy",
        "indices": root / "gpu/csr-indices.int32.npy",
        "weights": root / "gpu/csr-weights.int64.npy",
    }
    missing = [str(path) for path in paths.values() if not path.is_file()]
    if missing:
        raise FileNotFoundError(missing)
    return paths


def edge_frame(edge_ids: np.ndarray, sources: np.ndarray, indices: np.ndarray, weights: np.ndarray,
               metadata: pd.DataFrame, prefix: str = "") -> pd.DataFrame:
    src = sources[edge_ids]
    dst = indices[edge_ids]
    out = pd.DataFrame({
        "edge_id": edge_ids.astype(np.int64),
        "source_neuron_idx": src.astype(np.int32),
        "source_body_id": metadata.loc[src, "body_id"].to_numpy(dtype=np.uint64),
        "target_neuron_idx": dst.astype(np.int32),
        "target_body_id": metadata.loc[dst, "body_id"].to_numpy(dtype=np.uint64),
        "raw_synapse_count": weights[edge_ids].astype(np.int64),
    })
    if prefix:
        out["mapping"] = prefix
    return out


def tree_manifest(root: Path, excludes: set[Path]) -> dict[str, Any]:
    records = []
    for path in sorted(item for item in root.rglob("*") if item.is_file() and item not in excludes):
        records.append({"path": str(path.relative_to(root)), "bytes": path.stat().st_size, "sha256": sha256(path)})
    return {"root": str(root), "generated_utc": datetime.now(timezone.utc).isoformat(), "files": records,
            "file_count": len(records), "total_bytes": sum(item["bytes"] for item in records)}


def main() -> None:
    parser = argparse.ArgumentParser()
    parser.add_argument("--config", type=Path, required=True)
    args = parser.parse_args()
    cfg = yaml.safe_load(args.config.read_text(encoding="utf-8"))
    exp = Path(cfg["paths"]["experiment_root"])
    model = Path(cfg["paths"]["model_root"])
    for sub in ("audit", "mappings", "plastic-mask", "manifests", "regression"):
        (exp / sub).mkdir(parents=True, exist_ok=True)
    model.mkdir(parents=True, exist_ok=True)

    # This is deliberately completed before reading circuit data or writing audit products.
    frozen_before = run_frozen_audit(cfg)
    write_json(exp / "regression/frozen-checksums-before-stage1.json", frozen_before)

    paths = graph_paths(cfg)
    metadata = pq.read_table(paths["metadata"]).to_pandas()
    required = {"neuron_idx", "body_id", "annotation_class", "annotation_type", "annotation_instance",
                "annotation_superclass", "annotation_subclass", "side", "neurotransmitter", "nt_consensus_nt"}
    missing_columns = sorted(required - set(metadata.columns))
    if missing_columns:
        raise RuntimeError(f"metadata columns missing: {missing_columns}")
    metadata = metadata.sort_values("neuron_idx", kind="stable").reset_index(drop=True)
    if not np.array_equal(metadata["neuron_idx"].to_numpy(), np.arange(len(metadata), dtype=np.int32)):
        raise RuntimeError("metadata neuron_idx is not contiguous and row-aligned")

    indptr = np.load(paths["indptr"], mmap_mode="r")
    indices = np.load(paths["indices"], mmap_mode="r")
    weights = np.load(paths["weights"], mmap_mode="r")
    expected = cfg["expected_graph"]
    graph_actual = {"neurons": int(len(indptr) - 1), "directed_edges": int(len(indices)),
                    "synaptic_contact_weight_sum": int(weights.sum(dtype=np.int64))}
    if graph_actual != expected:
        raise RuntimeError(f"graph invariant mismatch: expected={expected}, actual={graph_actual}")

    classes = metadata["annotation_class"].fillna("").astype(str).to_numpy()
    types = metadata["annotation_type"].fillna("").astype(str).to_numpy()
    is_kc = classes == cfg["audit"]["exact_class_values"]["KC"]
    is_mbon = classes == cfg["audit"]["exact_class_values"]["MBON"]
    is_dan = classes == cfg["audit"]["exact_class_values"]["DAN"]
    is_pam = is_dan & np.char.startswith(types.astype(str), cfg["audit"]["pam_type_prefix"])
    is_ppl1 = is_dan & np.char.startswith(types.astype(str), cfg["audit"]["ppl1_type_prefix"])
    is_other_dan = is_dan & ~is_pam & ~is_ppl1

    populations = np.full(len(metadata), "other", dtype=object)
    populations[is_kc] = "KC"
    populations[is_mbon] = "MBON"
    populations[is_pam] = "PAM"
    populations[is_ppl1] = "PPL1"
    populations[is_other_dan] = "other_DAN"
    compartment_values: list[str | None] = []
    compartment_status: list[str] = []
    for value in metadata["annotation_instance"]:
        compartment, status = parse_compartment(value)
        compartment_values.append(compartment)
        compartment_status.append(status)
    metadata["phase4h_population"] = populations
    metadata["mb_compartment_verbatim"] = compartment_values
    metadata["mb_compartment_mapping_status"] = compartment_status
    selected = is_kc | is_mbon | is_dan
    population_columns = [
        "neuron_idx", "body_id", "phase4h_population", "annotation_superclass", "annotation_class",
        "annotation_subclass", "annotation_type", "annotation_instance", "side", "neurotransmitter",
        "nt_consensus_nt", "mb_compartment_verbatim", "mb_compartment_mapping_status",
    ]
    metadata.loc[selected, population_columns].to_parquet(exp / "audit/mb-populations.parquet", index=False)

    # Vectorized CSR edge source materialization is audit-only and never replaces the frozen CSR.
    sources = np.repeat(np.arange(graph_actual["neurons"], dtype=np.int32), np.diff(indptr).astype(np.int64))
    targets = np.asarray(indices)
    kc_mbon_ids = np.flatnonzero(is_kc[sources] & is_mbon[targets]).astype(np.int64)
    dan_kc_ids = np.flatnonzero(is_dan[sources] & is_kc[targets]).astype(np.int64)
    dan_mbon_ids = np.flatnonzero(is_dan[sources] & is_mbon[targets]).astype(np.int64)
    mbon_out_ids = np.flatnonzero(is_mbon[sources]).astype(np.int64)

    kc_mbon = edge_frame(kc_mbon_ids, sources, targets, weights, metadata)
    dan_kc = edge_frame(dan_kc_ids, sources, targets, weights, metadata)
    dan_mbon = edge_frame(dan_mbon_ids, sources, targets, weights, metadata)
    mbon_downstream = edge_frame(mbon_out_ids, sources, targets, weights, metadata)
    for frame in (kc_mbon, dan_kc, dan_mbon, mbon_downstream):
        frame["source_annotation_type"] = metadata.loc[frame["source_neuron_idx"], "annotation_type"].to_numpy()
        frame["source_annotation_instance"] = metadata.loc[frame["source_neuron_idx"], "annotation_instance"].to_numpy()
        frame["source_compartment_verbatim"] = metadata.loc[frame["source_neuron_idx"], "mb_compartment_verbatim"].to_numpy()
        frame["target_annotation_class"] = metadata.loc[frame["target_neuron_idx"], "annotation_class"].to_numpy()
        frame["target_annotation_type"] = metadata.loc[frame["target_neuron_idx"], "annotation_type"].to_numpy()
        frame["target_annotation_instance"] = metadata.loc[frame["target_neuron_idx"], "annotation_instance"].to_numpy()
        frame["target_compartment_verbatim"] = metadata.loc[frame["target_neuron_idx"], "mb_compartment_verbatim"].to_numpy()

    dan_mbon["exact_compartment_match"] = (
        dan_mbon["source_compartment_verbatim"].notna()
        & dan_mbon["target_compartment_verbatim"].notna()
        & (dan_mbon["source_compartment_verbatim"] == dan_mbon["target_compartment_verbatim"])
    )
    confirmed_dan_mbon = dan_mbon[dan_mbon["exact_compartment_match"]].copy()
    confirmed_dan_mbon.to_parquet(exp / "mappings/dan-to-mbon-confirmed.parquet", index=False)
    dan_mbon.to_parquet(exp / "mappings/dan-to-mbon-all-actual-edges.parquet", index=False)
    dan_kc.to_parquet(exp / "mappings/dan-to-kc-actual-edges.parquet", index=False)
    kc_mbon.to_parquet(exp / "mappings/kc-to-mbon-actual-edges.parquet", index=False)
    mbon_downstream.to_parquet(exp / "mappings/mbon-to-downstream-actual-edges.parquet", index=False)

    reliable_mbon_nodes = set(confirmed_dan_mbon["target_neuron_idx"].astype(int).tolist())
    candidate = kc_mbon[kc_mbon["target_neuron_idx"].isin(reliable_mbon_nodes)].copy()
    candidate["candidate_basis"] = "actual_KC_to_MBON_edge; target_MBON_has_actual_DAN_edge_with_exact_verbatim_compartment_match"
    candidate["plasticity_polarity"] = int(cfg["plasticity"]["default_confirmed_compartment_polarity"])
    candidate.to_parquet(exp / "plastic-mask/plastic-edge-candidates.parquet", index=False)
    candidate_ids = candidate["edge_id"].to_numpy(dtype=np.int64)
    np.save(exp / "plastic-mask/plastic-edge-ids.int64.npy", candidate_ids, allow_pickle=False)
    packed = np.zeros(graph_actual["directed_edges"], dtype=np.bool_)
    packed[candidate_ids] = True
    np.save(exp / "plastic-mask/plastic-edge-mask.bool.npy", packed, allow_pickle=False)

    reliable_compartments = sorted(set(candidate["target_compartment_verbatim"].dropna().astype(str)))
    polarity_rows = []
    for compartment in reliable_compartments:
        polarity_rows.append({
            "compartment_verbatim": compartment,
            "polarity_compartment": int(cfg["plasticity"]["default_confirmed_compartment_polarity"]),
            "rule_family": cfg["plasticity"]["first_version_rule_family"],
            "interpretation": "model hypothesis: coincident KC and compartmental DAN activity depresses eligible KC-to-MBON synapses",
            "not_interpreted_as": "PAM-positive/PPL1-negative cluster rule",
            "provenance": ";".join(cfg["plasticity"]["polarity_provenance"]),
        })
    polarity = pd.DataFrame(polarity_rows, columns=["compartment_verbatim", "polarity_compartment", "rule_family",
                                                           "interpretation", "not_interpreted_as", "provenance"])
    polarity.to_parquet(exp / "mappings/compartment-plasticity-polarity.parquet", index=False)

    dan_nodes = metadata.loc[is_dan, ["neuron_idx", "body_id", "phase4h_population", "annotation_type",
                                      "annotation_instance", "mb_compartment_verbatim", "mb_compartment_mapping_status"]].copy()
    dan_nodes.to_parquet(exp / "mappings/dan-to-compartment.parquet", index=False)
    mbon_nodes = metadata.loc[is_mbon, ["neuron_idx", "body_id", "annotation_type", "annotation_instance",
                                        "neurotransmitter", "nt_consensus_nt", "mb_compartment_verbatim",
                                        "mb_compartment_mapping_status"]].copy()
    mbon_nodes.to_parquet(exp / "mappings/mbon-to-compartment.parquet", index=False)

    unknown = {
        "DAN_without_reliable_explicit_compartment": int(dan_nodes["mb_compartment_verbatim"].isna().sum()),
        "MBON_without_reliable_explicit_compartment": int(mbon_nodes["mb_compartment_verbatim"].isna().sum()),
        "actual_DAN_to_MBON_edges_not_exactly_compartment_matched": int((~dan_mbon["exact_compartment_match"]).sum()),
        "KC_to_MBON_edges_excluded_from_candidate_mask": int(len(kc_mbon) - len(candidate)),
        "KC_to_MBON_target_MBONs_without_confirmed_DAN_match": int(kc_mbon.loc[~kc_mbon["target_neuron_idx"].isin(reliable_mbon_nodes), "target_neuron_idx"].nunique()),
    }
    unknown["reported_relationship_unknown_total_non_disjoint"] = int(sum(unknown.values()))

    population_counts = {
        "KC": int(is_kc.sum()), "MBON": int(is_mbon.sum()), "PAM": int(is_pam.sum()),
        "PPL1": int(is_ppl1.sum()), "other_DAN": int(is_other_dan.sum()), "DAN_total": int(is_dan.sum()),
    }
    audit_pass = bool(
        frozen_before["all_passed"]
        and len(reliable_compartments) > 0
        and len(candidate) > 0
        and len(confirmed_dan_mbon) > 0
        and np.all(weights[candidate_ids] > 0)
    )
    report = {
        "phase": "4H-stage1",
        "status": "PASS-AUDIT" if audit_pass else "FAIL-INSUFFICIENT-BIOLOGICAL-MAPPING",
        "stage2_authorized_by_audit": audit_pass,
        "no_weight_update_or_conditioning_performed": True,
        "graph": graph_actual,
        "population_counts": population_counts,
        "reliable_mb_compartments": {"count": len(reliable_compartments), "verbatim_values": reliable_compartments},
        "KC_to_MBON": {"edges": int(len(kc_mbon)), "synapse_weight": int(kc_mbon["raw_synapse_count"].sum())},
        "DAN_to_KC": {"edges": int(len(dan_kc)), "synapse_weight": int(dan_kc["raw_synapse_count"].sum())},
        "DAN_to_MBON": {
            "actual_edges": int(len(dan_mbon)), "actual_synapse_weight": int(dan_mbon["raw_synapse_count"].sum()),
            "confirmed_exact_compartment_edge_mappings": int(len(confirmed_dan_mbon)),
            "confirmed_unique_DAN_MBON_pairs": int(confirmed_dan_mbon[["source_neuron_idx", "target_neuron_idx"]].drop_duplicates().shape[0]),
        },
        "DAN_to_compartment": {
            "confirmed_neurons": int(dan_nodes["mb_compartment_verbatim"].notna().sum()),
            "unknown_neurons": int(dan_nodes["mb_compartment_verbatim"].isna().sum()),
            "unique_explicit_compartments": int(dan_nodes["mb_compartment_verbatim"].dropna().nunique()),
        },
        "MBON_to_downstream": {"edges": int(len(mbon_downstream)), "synapse_weight": int(mbon_downstream["raw_synapse_count"].sum()),
                                "unique_targets": int(mbon_downstream["target_neuron_idx"].nunique())},
        "unknown": unknown,
        "plastic_edge_candidate_mask": {
            "edges": int(len(candidate)),
            "fraction_of_all_edges": float(len(candidate) / graph_actual["directed_edges"]),
            "percent_of_all_edges": float(100.0 * len(candidate) / graph_actual["directed_edges"]),
            "synapse_weight": int(candidate["raw_synapse_count"].sum()),
            "full_edge_boolean_mask_sha256": sha256(exp / "plastic-mask/plastic-edge-mask.bool.npy"),
            "edge_ids_sha256": sha256(exp / "plastic-mask/plastic-edge-ids.int64.npy"),
        },
        "compartment_polarity": polarity.to_dict(orient="records"),
        "mapping_policy": cfg["audit"],
        "biological_basis_conclusion": (
            "sufficient for a conservative, explicitly compartment-matched conditioning prototype"
            if audit_pass else "insufficient; conditioning must not start"
        ),
        "limitations": [
            "The neuron-level CSR has no synapse-level ROI/location, so candidate membership is conservative and does not prove that every KC-MBON contact lies inside the parsed compartment.",
            "Parenthesized annotation_instance tokens are preserved verbatim; missing or ambiguous tokens are not inferred from names.",
            "Direct DAN-to-MBON connectivity plus exact compartment-string equality is required; extrasynaptic dopamine is not inferred.",
            "polarity=-1 is an explicit anti-Hebbian/LTD model hypothesis with cited provenance, not a PAM/PPL1 valence rule.",
        ],
        "input_sha256": {name: sha256(path) for name, path in paths.items()},
        "frozen_checksum_audit": frozen_before,
    }
    write_json(exp / "audit/stage1-audit-report.json", report)
    shutil.copy2(args.config, exp / "audit/phase4h-config.yaml")
    shutil.copy2(Path(__file__).resolve(), exp / "audit/phase4h_audit.py")
    write_json(model / "STAGE1_AUDIT_GATE.json", {
        "status": report["status"], "stage2_authorized": audit_pass,
        "audit_report": str(exp / "audit/stage1-audit-report.json"),
        "audit_report_sha256": sha256(exp / "audit/stage1-audit-report.json"),
    })

    # Seal the stage-1 products without touching any pre-4H artifact.
    manifest_path = exp / "manifests/stage1-manifest.json"
    checksum_path = exp / "manifests/stage1-checksums.sha256"
    print(json.dumps(report, indent=2, sort_keys=True, default=json_default))
    sys.stdout.flush()
    manifest = tree_manifest(exp, {manifest_path, checksum_path})
    write_json(manifest_path, manifest)
    targets = sorted(path for path in exp.rglob("*") if path.is_file() and path != checksum_path)
    checksum_path.write_text("".join(f"{sha256(path)}  {path.relative_to(exp)}\n" for path in targets), encoding="utf-8")


if __name__ == "__main__":
    main()