#!/usr/bin/env python3
"""Summarize GPS-relevant structure in all.traffic_persistent JSON."""

from __future__ import annotations

import argparse
import collections
import json
from pathlib import Path
from typing import Any


FLAGS = {
    "FromRoadSpline": 1,
    "Bidirectional": 2,
    "PatrolRoute": 4,
    "Pavement": 8,
    "Road": 16,
    "Intersection": 32,
    "NeverDeadEnd": 64,
    "TrafficDisabled": 128,
    "CrossWalk": 256,
    "GPSOnly": 512,
    "ShowDebug": 1024,
    "Blockade": 2048,
    "Yield": 4096,
    "NoAIDriving": 8192,
    "Highway": 16384,
    "NoAutodrive": 32768,
}


def flag_value(flags: Any) -> int:
    if isinstance(flags, int):
        return flags
    if isinstance(flags, str):
        return int(flags)
    if isinstance(flags, dict):
        if "Value" in flags:
            return int(flags["Value"])
        if "$value" in flags:
            return int(flags["$value"])
    return int(flags or 0)


def has(flags: int, name: str) -> bool:
    return bool(flags & FLAGS[name])


def category(flags: int) -> str:
    if has(flags, "Highway"):
        return "highway"
    if has(flags, "GPSOnly"):
        return "gpsonly"
    if has(flags, "Road"):
        return "road"
    if has(flags, "Pavement"):
        return "pavement"
    return "other"


def ref_index(ref: Any) -> int | None:
    if isinstance(ref, int):
        return ref
    if isinstance(ref, str):
        try:
            return int(ref)
        except ValueError:
            return None
    if isinstance(ref, dict):
        for key in ("Index", "index", "laneIndex", "LaneIndex", "value", "$value"):
            if key in ref:
                return ref_index(ref[key])
    return None


def lane_refs(value: Any) -> list[int]:
    if value is None:
        return []
    if isinstance(value, list):
        out: list[int] = []
        for item in value:
            idx = ref_index(item)
            if idx is not None:
                out.append(idx)
        return out
    if isinstance(value, dict) and "Elements" in value:
        return lane_refs(value["Elements"])
    return []


def gps_info(lane: dict[str, Any]) -> tuple[int | None, int | None]:
    info = lane.get("playerGPSInfo") or {}
    return info.get("subGraphId"), info.get("stronglyConnectedComponentId")


def main() -> int:
    parser = argparse.ArgumentParser()
    parser.add_argument("json_file", type=Path)
    parser.add_argument("--lane-connections", type=Path)
    parser.add_argument("--samples", type=int, default=12)
    args = parser.parse_args()

    data = json.loads(args.json_file.read_text(encoding="utf-8"))
    root = data["Data"]["RootChunk"]["data"]
    lanes = root["lanes"]
    neighbor_groups = root.get("neighborGroups", [])

    categories = []
    flags_by_lane = []
    for lane in lanes:
        flags = flag_value(lane.get("flags", 0))
        flags_by_lane.append(flags)
        categories.append(category(flags))

    print(f"lanes: {len(lanes)}")
    print(f"neighborGroups: {len(neighbor_groups)}")
    print("categories:", dict(collections.Counter(categories).most_common()))

    gps_counter = collections.Counter(gps_info(lane) for lane in lanes)
    print(f"gps components: {len(gps_counter)}")
    print("top gps components:")
    for (subgraph, scc), count in gps_counter.most_common(20):
        cats = collections.Counter(
            categories[i] for i, lane in enumerate(lanes) if gps_info(lane) == (subgraph, scc)
        )
        print(f"  subGraph={subgraph} scc={scc}: {count} {dict(cats.most_common())}")

    connection_rows: list[dict[str, Any]] | None = None
    if args.lane_connections:
        connection_data = json.loads(args.lane_connections.read_text(encoding="utf-8"))
        connection_rows = connection_data["Data"]["RootChunk"]["data"]

    transitions = collections.Counter()
    probabilities: dict[tuple[str, str], collections.Counter[int]] = collections.defaultdict(
        collections.Counter
    )
    sharp_angles = collections.Counter()
    out_degree = collections.Counter()
    missing_refs = 0
    rows = enumerate(lanes) if connection_rows is None else (
        (row["index"], row["value"]) for row in connection_rows
    )
    for i, lane_or_connections in rows:
        refs = (
            lane_refs(lane_or_connections.get("outLanes"))
            if connection_rows is None
            else lane_refs(lane_or_connections.get("outlanes"))
        )
        out_degree[len(refs)] += 1
        raw_outs = (
            lane_or_connections.get("outLanes", [])
            if connection_rows is None
            else lane_or_connections.get("outlanes", [])
        )
        for out, ref in zip(raw_outs, refs):
            if ref < 0 or ref >= len(lanes):
                missing_refs += 1
                continue
            pair = (categories[i], categories[ref])
            transitions[pair] += 1
            if isinstance(out, dict) and "exitProbabilityCompressed" in out:
                probabilities[pair][out["exitProbabilityCompressed"]] += 1
            if isinstance(out, dict) and "isSharpAngle" in out:
                sharp_angles[(pair, out["isSharpAngle"])] += 1
    print("out degree:", dict(sorted(out_degree.items())))
    print(f"missing out refs: {missing_refs}")
    print("category transitions:")
    for (src, dst), count in transitions.most_common(30):
        extra = ""
        if probabilities[(src, dst)]:
            extra = f" prob={probabilities[(src, dst)].most_common(6)}"
        sharp = [(key[1], value) for key, value in sharp_angles.items() if key[0] == (src, dst)]
        if sharp:
            extra += f" sharp={sharp}"
        print(f"  {src:8s} -> {dst:8s}: {count}{extra}")

    print("highway connector samples:")
    shown = 0
    for i, lane in enumerate(lanes):
        if categories[i] != "highway":
            continue
        if connection_rows is None:
            refs = [ref for ref in lane_refs(lane.get("outLanes")) if 0 <= ref < len(lanes)]
        else:
            row = connection_rows[i]["value"]
            refs = [ref for ref in lane_refs(row.get("outlanes")) if 0 <= ref < len(lanes)]
        ref_cats = collections.Counter(categories[ref] for ref in refs)
        if ref_cats and any(cat != "highway" for cat in ref_cats):
            subgraph, scc = gps_info(lane)
            print(
                f"  lane={i} flags={flags_by_lane[i]} len={lane.get('length')} "
                f"speed={lane.get('maxSpeed')} gps=({subgraph},{scc}) out={dict(ref_cats)}"
            )
            shown += 1
            if shown >= args.samples:
                break

    return 0


if __name__ == "__main__":
    raise SystemExit(main())