"""Export measured MaleCNS somata for the viewer; no invented coordinates."""

import argparse
import hashlib
import json
import tempfile
from pathlib import Path
from urllib.request import urlretrieve

import numpy as np
from pyarrow import feather, parquet

ROOT = Path(__file__).resolve().parents[1]
SOURCE = "https://storage.googleapis.com/flyem-male-cns/v1.0/connectome-data/flat-connectome/body-annotations-male-cns-v1.0-minconf-0.5.feather"


def build(path):
    rows = feather.read_table(
        path, columns=["bodyId", "somaLocation", "status"]
    ).to_pylist()
    rows = sorted(
        [
            row
            for row in rows
            if row["status"] == "Traced"
            and row["somaLocation"] is not None
            and len(row["somaLocation"]) == 3
        ],
        key=lambda row: row["bodyId"],
    )
    points = np.array([row["somaLocation"] for row in rows], dtype="<f4")
    points.tofile(ROOT / "web/anatomy.bin")
    by_id = {row["bodyId"]: row["somaLocation"] for row in rows}
    nodes = parquet.read_table(ROOT / "data/nodes.parquet").to_pylist()
    model = [by_id.get(node["node_id"]) for node in nodes]
    (ROOT / "web/neuron-positions.json").write_text(json.dumps(model))
    manifest = {
        "source": SOURCE,
        "source_sha256": hashlib.sha256(Path(path).read_bytes()).hexdigest(),
        "license": "CC-BY-4.0",
        "count": len(rows),
        "format": "little-endian float32 native XYZ triples; 8 nm voxels",
        "filter": "Traced status and measured somaLocation; sorted by bodyId",
        "missing_model_positions": sum(position is None for position in model),
        "note": "Static anatomy context. Only the controller subgraph has simulated activity.",
    }
    (ROOT / "data/anatomy-provenance.json").write_text(json.dumps(manifest, indent=2))
    print(json.dumps(manifest, indent=2))


if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    parser.add_argument("--source-file", type=Path)
    args = parser.parse_args()
    if args.source_file:
        build(args.source_file)
    else:
        with tempfile.TemporaryDirectory() as directory:
            path = Path(directory) / "annotations.feather"
            urlretrieve(SOURCE, path)
            build(path)