#!/usr/bin/env python3
"""Snapshot the official MaleCNS retinotopy T-bar annotations at body level."""
from __future__ import annotations

import argparse
import json
import time
from collections import defaultdict
from pathlib import Path

import numpy as np
import pyarrow as pa
import pyarrow.parquet as pq
from neuroglancer.read_precomputed_annotations import AnnotationReader


def main() -> None:
    ap = argparse.ArgumentParser()
    ap.add_argument("--url", required=True)
    ap.add_argument("--output", type=Path, required=True)
    ap.add_argument("--summary", type=Path, required=True)
    args = ap.parse_args()

    reader = AnnotationReader(args.url)
    prop_ids = [p.id for p in reader.properties]
    expected = ["hex1", "hex2", "area_to_count", "r", "theta", "v", "h", "ht", "vic", "area_fit", "body", "roi"]
    if prop_ids != expected:
        raise RuntimeError(f"official schema changed: {prop_ids}")

    # Values are accumulated with area_to_count as the official contribution
    # weight.  This is a lossless body-level sufficient statistic for the
    # coordinate means used by Phase 4F; no body IDs are inferred.
    names = ["hex1", "hex2", "r", "theta", "v", "h", "ht", "vic", "area_fit"]
    accum: dict[int, dict[str, float]] = defaultdict(lambda: defaultdict(float))
    roi_sets: dict[int, set[int]] = defaultdict(set)
    t0 = time.perf_counter()
    annotations = 0
    for ann in reader.get_within_spatial_bounds(max_parallelism=128):
        p = ann.props
        body = int(p[10])
        weight = float(p[2])
        row = accum[body]
        row["weight_sum"] += weight
        row["annotation_count"] += 1.0
        for i, name in enumerate(names):
            value = float(p[i if i < 2 else i + 1])
            if np.isfinite(value):
                row[f"{name}_weighted_sum"] += weight * value
                row[f"{name}_finite_weight"] += weight
        roi_sets[body].add(int(p[11]))
        annotations += 1
        if annotations % 250_000 == 0:
            print(json.dumps({"annotations": annotations, "bodies": len(accum), "wall_seconds": time.perf_counter() - t0}), flush=True)

    bodies = np.asarray(sorted(accum), dtype=np.uint64)
    cols: dict[str, pa.Array] = {
        "body_id": pa.array(bodies, type=pa.uint64()),
        "tbar_annotation_count": pa.array([int(accum[int(b)]["annotation_count"]) for b in bodies], type=pa.int32()),
        "visual_input_contribution_sum": pa.array([accum[int(b)]["weight_sum"] for b in bodies], type=pa.float64()),
    }
    for name in names:
        values = []
        for body in bodies:
            row = accum[int(body)]
            den = row[f"{name}_finite_weight"]
            values.append(row[f"{name}_weighted_sum"] / den if den else None)
        cols[f"tbar_{name}_weighted_mean"] = pa.array(values, type=pa.float64())
    cols["tbar_roi_codes"] = pa.array([sorted(roi_sets[int(b)]) for b in bodies], type=pa.list_(pa.uint8()))
    table = pa.table(cols)
    args.output.parent.mkdir(parents=True, exist_ok=True)
    pq.write_table(table, args.output, compression="zstd", compression_level=9)
    summary = {
        "source_url": args.url,
        "official_property_schema": prop_ids,
        "coordinate_space": reader.metadata["dimensions"],
        "lower_bound": reader.lower_bound,
        "upper_bound": reader.upper_bound,
        "annotation_count": annotations,
        "unique_body_count": len(bodies),
        "aggregation": "body-level area_to_count-weighted means; official body property used verbatim",
        "wall_seconds": time.perf_counter() - t0,
        "output": str(args.output),
    }
    args.summary.write_text(json.dumps(summary, indent=2) + "\n", encoding="utf-8")
    print(json.dumps(summary), flush=True)


if __name__ == "__main__":
    main()