#!/usr/bin/env python3
"""Reproduce NeuroAI's descriptive MaleCNS v1.0 analysis (not an AI training run).

Requirements: Python >=3.10, pandas >=2.2, pyarrow >=21.
Usage: python analyze_connectome.py --input-dir DATA --output-dir OUTPUT

Download the four URLs in SOURCES into DATA using their `local` filenames.
Inputs are public, CC BY 4.0. The source data are from FlyEM (HHMI Janelia),
the University of Cambridge, MRC Laboratory of Molecular Biology and Google
Research. This script creates aggregate tables and a derived 3D centreline view.
No individual-level human data are used. No model is trained or benchmarked.
"""

import argparse
import csv
import hashlib
import json
import math
from collections import Counter
from pathlib import Path

import pandas as pd

ROOT = "https://storage.googleapis.com/flyem-male-cns/v1.0/"
SOURCES = [
    {"local": "body-annotations.feather", "url": ROOT + "connectome-data/flat-connectome/body-annotations-male-cns-v1.0-minconf-0.5.feather"},
    {"local": "body-neurotransmitters.feather", "url": ROOT + "connectome-data/flat-connectome/body-neurotransmitters-male-cns-v1.0.feather"},
    {"local": "12781.swc", "url": ROOT + "segmentation/skeletons-malecns/skeletons-swc/12781.swc"},
    {"local": "556329.swc", "url": ROOT + "segmentation/skeletons-malecns/skeletons-swc/556329.swc"},
]
DATE = "2026-09-05"
THRESHOLDS = [0.5, 0.7, 0.9]
CLASS_LABELS = {
    "ol_intrinsic": "Optic-lobe intrinsic",
    "cb_intrinsic": "Central-brain intrinsic",
    "vnc_intrinsic": "Ventral-cord intrinsic",
    "visual_projection": "Visual projection",
    "vnc_sensory": "Ventral-cord sensory",
    "cb_sensory": "Central-brain sensory",
}
NT_LABELS = {"acetylcholine": "Acetylcholine", "glutamate": "Glutamate", "gaba": "GABA", "dopamine": "Dopamine", "serotonin": "Serotonin", "octopamine": "Octopamine"}


def write_json(path, content):
    path.write_text(json.dumps(content, separators=(",", ":"), allow_nan=False) + "\n")


def rows_for_counts(counts, denominator, labels=None):
    return [{"key": str(key), "label": (labels or {}).get(key, str(key)), "count": int(count),
             "percent": round(int(count) / denominator * 100, 4)} for key, count in counts.items()]


def read_swc(path):
    nodes = {}
    for line in path.read_text().splitlines():
        if not line.strip() or line.startswith("#"):
            continue
        index, kind, x, y, z, radius, parent = line.split()
        nodes[int(index)] = {"position": [float(x), float(y), float(z)], "parent": int(parent)}
    for node in nodes.values():
        assert node["parent"] < 0 or node["parent"] in nodes, "Missing SWC parent"
        assert all(math.isfinite(v) for v in node["position"])
    return nodes


def skeleton_metrics(nodes):
    children = Counter(node["parent"] for node in nodes.values() if node["parent"] >= 0)
    # Sum Euclidean edge lengths in original 8-nanometre coordinate units.
    length = sum(math.dist(node["position"], nodes[node["parent"]]["position"])
                 for node in nodes.values() if node["parent"] >= 0)
    return {"nodes": len(nodes), "branchPoints": sum(v > 1 for v in children.values()),
            "tips": sum(index not in children for index in nodes),
            "cableLengthMm": round(length * 8 / 1_000_000, 5)}


def simplify(nodes, centre, span):
    # Keep every branch, tip and root; decimate only degree-two paths (stride 4).
    # This preserves tree topology. The view is a simplified centreline, not a mesh.
    children = Counter(node["parent"] for node in nodes.values() if node["parent"] >= 0)
    keep = {index for index, node in nodes.items()
            if node["parent"] < 0 or children[index] != 1 or index % 4 == 0}
    ordered = sorted(keep)
    index_map = {index: i for i, index in enumerate(ordered)}
    points = []
    for index in ordered:
        node = nodes[index]
        parent = node["parent"]
        visited = {index}
        while parent >= 0 and parent not in keep:
            assert parent not in visited, "SWC cycle"
            visited.add(parent)
            parent = nodes[parent]["parent"]
        points.append([*[round((v - centre[i]) / span, 5) for i, v in enumerate(node["position"])],
                       index_map[parent] if parent >= 0 else -1])
    return points


def main():
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--input-dir", type=Path, required=True)
    parser.add_argument("--output-dir", type=Path, required=True)
    args = parser.parse_args()
    args.output_dir.mkdir(parents=True, exist_ok=True)
    annotations = pd.read_feather(args.input_dir / SOURCES[0]["local"])
    assert annotations.bodyId.is_unique, "Duplicate body IDs in annotations"
    cohort = annotations.loc[annotations.status.eq("Traced")].copy()
    total = len(cohort)
    assert total > 0
    counts = cohort.superclass.fillna("Unassigned").value_counts()
    full_classes = rows_for_counts(counts, total)
    displayed = {key: int(counts.get(key, 0)) for key in CLASS_LABELS}
    displayed["Other annotated classes"] = int(counts.drop([*CLASS_LABELS, "Unassigned"]).sum())
    displayed["Unassigned"] = int(counts.get("Unassigned", 0))
    assert sum(displayed.values()) == total

    nt = pd.read_feather(args.input_dir / SOURCES[1]["local"])
    assert nt.body.is_unique, "Duplicate body IDs in transmitter predictions"
    merged = cohort[["bodyId"]].merge(nt, left_on="bodyId", right_on="body", how="left", validate="one_to_one")
    thresholds = []
    for threshold in THRESHOLDS:
        valid = (merged.predicted_nt_confidence.ge(threshold) & merged.total_nt_predictions.gt(0)
                 & merged.predicted_nt.isin(NT_LABELS))
        raw_counts = merged.loc[valid, "predicted_nt"].value_counts()
        values = {name: int(raw_counts.get(name, 0)) for name in NT_LABELS}
        values["Below threshold / unavailable"] = int((~valid).sum())
        assert sum(values.values()) == total
        thresholds.append({"threshold": threshold, "retained": int(valid.sum()),
                           "coveragePercent": round(valid.sum() / total * 100, 4),
                           "rows": rows_for_counts(values, total, NT_LABELS)})

    trees = [read_swc(args.input_dir / source["local"]) for source in SOURCES[2:]]
    coordinates = [node["position"] for tree in trees for node in tree.values()]
    minimum = [min(v[i] for v in coordinates) for i in range(3)]
    maximum = [max(v[i] for v in coordinates) for i in range(3)]
    centre = [(a + b) / 2 for a, b in zip(minimum, maximum)]
    span = max(b - a for a, b in zip(minimum, maximum))
    skeletons = []
    for source, tree in zip(SOURCES[2:], trees):
        body_id = int(source["local"].split(".")[0])
        row = cohort.loc[cohort.bodyId.eq(body_id)].iloc[0]
        skeletons.append({"bodyId": body_id, "label": row["instance"], "side": row.somaSide,
                          "type": row["type"], **skeleton_metrics(tree),
                          "points": simplify(tree, centre, span)})

    provenance = {"dataset": "MaleCNS", "release": "v1.0", "analysedOn": DATE,
                  "filter": "status == 'Traced'", "inputRows": len(annotations),
                  "excludedRows": len(annotations) - total, "total": total,
                  "uniqueTypes": int(cohort.type.nunique()),
                  "sourcePage": "https://male-cns.janelia.org/download/",
                  "license": "https://creativecommons.org/licenses/by/4.0/",
                  "attribution": "FlyEM (HHMI Janelia), University of Cambridge, MRC Laboratory of Molecular Biology, Google Research",
                  "sources": [{**source, "sha256": hashlib.sha256((args.input_dir / source["local"]).read_bytes()).hexdigest()}
                              for source in SOURCES]}
    summary = {**provenance, "classes": rows_for_counts(displayed, total, CLASS_LABELS),
               "allClasses": full_classes, "neurotransmitters": thresholds,
               "skeletons": [{key: value for key, value in item.items() if key != "points"} for item in skeletons]}
    write_json(args.output_dir / "malecns-summary.json", summary)
    write_json(args.output_dir / "malecns-skeletons.json",
               {"release": "v1.0", "coordinateUnitNm": 8, "normalizationCentre": centre,
                "normalizationSpan": span, "simplificationStride": 4, "skeletons": skeletons})
    with (args.output_dir / "malecns-analysis.csv").open("w", newline="") as target:
        writer = csv.writer(target)
        writer.writerow(["dataset", "release", "filter", "analysis", "category", "threshold", "value", "unit", "denominator", "percent", "source_url"])
        for row in full_classes:
            writer.writerow(["MaleCNS", "v1.0", provenance["filter"], "superclass", row["key"], "", row["count"], "neurons", total, row["percent"], SOURCES[0]["url"]])
        for threshold in thresholds:
            for row in threshold["rows"]:
                writer.writerow(["MaleCNS", "v1.0", provenance["filter"], "predicted_nt", row["key"], threshold["threshold"], row["count"], "neurons", total, row["percent"], SOURCES[1]["url"]])
        for item, source in zip(skeletons, SOURCES[2:]):
            writer.writerow(["MaleCNS", "v1.0", provenance["filter"], "skeleton_length", item["label"], "", item["cableLengthMm"], "mm", "", "", source["url"]])
    print(json.dumps({"total": total, "classes": summary["classes"],
                      "thresholds": [{k: v for k, v in t.items() if k != "rows"} for t in thresholds],
                      "skeletons": summary["skeletons"]}, indent=2))


if __name__ == "__main__":
    main()
