#!/usr/bin/env python3
"""Generate a compact road-class grid for the RED4ext GPS edge-cost shim."""

from __future__ import annotations

import argparse
import json
import math
from dataclasses import dataclass
from pathlib import Path
from typing import Any, Iterable


FLAGS = {
    "Pavement": 0x0008,
    "Road": 0x0010,
    "GPSOnly": 0x0200,
    "Highway": 0x4000,
}

CATEGORY_CODES = {
    "unknown": ".",
    "pavement": "P",
    "gpsonly": "G",
    "road": "R",
    "highway": "H",
}

CATEGORY_PRIORITY = {
    ".": 0,
    "P": 1,
    "G": 2,
    "R": 3,
    "H": 4,
}


@dataclass(frozen=True)
class GridSpec:
    min_x: float
    min_y: float
    width: int
    height: int
    cell_size: float


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


def lane_category(flags: int) -> str:
    if flags & FLAGS["Highway"]:
        return "highway"
    if flags & FLAGS["GPSOnly"]:
        return "gpsonly"
    if flags & FLAGS["Road"]:
        return "road"
    if flags & FLAGS["Pavement"]:
        return "pavement"
    return "unknown"


def load_lanes(path: Path) -> list[dict[str, Any]]:
    data = json.loads(path.read_text(encoding="utf-8"))
    return data["Data"]["RootChunk"]["data"]["lanes"]


def load_lane_polygons(path: Path) -> list[list[tuple[float, float]]]:
    data = json.loads(path.read_text(encoding="utf-8"))
    rows = data["Data"]["RootChunk"]["data"]
    polygons: list[list[tuple[float, float]]] = []
    for row in rows:
        polygon = row.get("value", {}).get("polygon") or []
        polygons.append([(float(point["X"]), float(point["Y"])) for point in polygon])
    return polygons


def point_in_polygon(x: float, y: float, polygon: list[tuple[float, float]]) -> bool:
    inside = False
    count = len(polygon)
    if count < 3:
        return False

    previous_x, previous_y = polygon[-1]
    for current_x, current_y in polygon:
        if (current_y > y) != (previous_y > y):
            edge_x = (previous_x - current_x) * (y - current_y) / (previous_y - current_y) + current_x
            if x < edge_x:
                inside = not inside
        previous_x, previous_y = current_x, current_y

    return inside


def orientation(ax: float, ay: float, bx: float, by: float, cx: float, cy: float) -> float:
    return (by - ay) * (cx - bx) - (bx - ax) * (cy - by)


def on_segment(ax: float, ay: float, bx: float, by: float, cx: float, cy: float) -> bool:
    return min(ax, cx) <= bx <= max(ax, cx) and min(ay, cy) <= by <= max(ay, cy)


def segments_intersect(
    ax: float,
    ay: float,
    bx: float,
    by: float,
    cx: float,
    cy: float,
    dx: float,
    dy: float,
) -> bool:
    o1 = orientation(ax, ay, bx, by, cx, cy)
    o2 = orientation(ax, ay, bx, by, dx, dy)
    o3 = orientation(cx, cy, dx, dy, ax, ay)
    o4 = orientation(cx, cy, dx, dy, bx, by)

    if (o1 > 0) != (o2 > 0) and (o3 > 0) != (o4 > 0):
        return True

    eps = 1e-5
    if abs(o1) <= eps and on_segment(ax, ay, cx, cy, bx, by):
        return True
    if abs(o2) <= eps and on_segment(ax, ay, dx, dy, bx, by):
        return True
    if abs(o3) <= eps and on_segment(cx, cy, ax, ay, dx, dy):
        return True
    if abs(o4) <= eps and on_segment(cx, cy, bx, by, dx, dy):
        return True
    return False


def polygon_intersects_cell(
    polygon: list[tuple[float, float]],
    min_x: float,
    min_y: float,
    max_x: float,
    max_y: float,
) -> bool:
    center_x = (min_x + max_x) * 0.5
    center_y = (min_y + max_y) * 0.5
    if point_in_polygon(center_x, center_y, polygon):
        return True

    corners = ((min_x, min_y), (max_x, min_y), (max_x, max_y), (min_x, max_y))
    if any(point_in_polygon(x, y, polygon) for x, y in corners):
        return True

    if any(min_x <= x <= max_x and min_y <= y <= max_y for x, y in polygon):
        return True

    cell_edges = (
        (min_x, min_y, max_x, min_y),
        (max_x, min_y, max_x, max_y),
        (max_x, max_y, min_x, max_y),
        (min_x, max_y, min_x, min_y),
    )
    previous_x, previous_y = polygon[-1]
    for current_x, current_y in polygon:
        for edge in cell_edges:
            if segments_intersect(previous_x, previous_y, current_x, current_y, *edge):
                return True
        previous_x, previous_y = current_x, current_y

    return False


def grid_index(spec: GridSpec, x: float, y: float) -> tuple[int, int]:
    column = math.floor((x - spec.min_x) / spec.cell_size)
    row = math.floor((y - spec.min_y) / spec.cell_size)
    return int(column), int(row)


def mark_cell(grid: list[list[str]], column: int, row: int, code: str) -> bool:
    if row < 0 or row >= len(grid) or column < 0 or column >= len(grid[row]):
        return False
    if CATEGORY_PRIORITY[code] <= CATEGORY_PRIORITY[grid[row][column]]:
        return False
    grid[row][column] = code
    return True


def iter_cells_for_polygon(spec: GridSpec, polygon: list[tuple[float, float]]) -> Iterable[tuple[int, int]]:
    xs = [point[0] for point in polygon]
    ys = [point[1] for point in polygon]
    min_col, min_row = grid_index(spec, min(xs), min(ys))
    max_col, max_row = grid_index(spec, max(xs), max(ys))

    for row in range(max(0, min_row), min(spec.height - 1, max_row) + 1):
        for column in range(max(0, min_col), min(spec.width - 1, max_col) + 1):
            cell_min_x = spec.min_x + column * spec.cell_size
            cell_min_y = spec.min_y + row * spec.cell_size
            if polygon_intersects_cell(
                polygon,
                cell_min_x,
                cell_min_y,
                cell_min_x + spec.cell_size,
                cell_min_y + spec.cell_size,
            ):
                yield column, row


def build_grid(
    lanes: list[dict[str, Any]],
    polygons: list[list[tuple[float, float]]],
    cell_size: float,
    inflate_cells: int,
) -> tuple[GridSpec, list[list[str]], dict[str, int]]:
    points = [point for polygon in polygons for point in polygon]
    if not points:
        raise ValueError("no polygon points found")

    min_x = math.floor(min(point[0] for point in points) / cell_size) * cell_size
    min_y = math.floor(min(point[1] for point in points) / cell_size) * cell_size
    max_x = math.ceil(max(point[0] for point in points) / cell_size) * cell_size
    max_y = math.ceil(max(point[1] for point in points) / cell_size) * cell_size
    spec = GridSpec(
        min_x=min_x,
        min_y=min_y,
        width=int(round((max_x - min_x) / cell_size)),
        height=int(round((max_y - min_y) / cell_size)),
        cell_size=cell_size,
    )

    grid = [[CATEGORY_CODES["unknown"] for _ in range(spec.width)] for _ in range(spec.height)]
    for index, (lane, polygon) in enumerate(zip(lanes, polygons, strict=True)):
        if len(polygon) < 3:
            continue
        code = CATEGORY_CODES[lane_category(flag_value(lane.get("flags", 0)))]
        if code == CATEGORY_CODES["unknown"]:
            continue

        changed_cells: set[tuple[int, int]] = set()
        for column, row in iter_cells_for_polygon(spec, polygon):
            for dy in range(-inflate_cells, inflate_cells + 1):
                for dx in range(-inflate_cells, inflate_cells + 1):
                    changed_cells.add((column + dx, row + dy))

        for column, row in changed_cells:
            mark_cell(grid, column, row, code)

    counts = {code: sum(row.count(code) for row in grid) for code in CATEGORY_PRIORITY}
    return spec, grid, counts


def write_header(path: Path, spec: GridSpec, grid: list[list[str]], counts: dict[str, int]) -> None:
    lines = [
        "#pragma once",
        "",
        "#include <array>",
        "#include <cstdint>",
        "#include <string_view>",
        "",
        "namespace EdgeWeightGPS::Generated",
        "{",
        f"constexpr float kSpatialRoadGridMinX = {spec.min_x:.1f}f;",
        f"constexpr float kSpatialRoadGridMinY = {spec.min_y:.1f}f;",
        f"constexpr float kSpatialRoadGridCellSize = {spec.cell_size:.1f}f;",
        f"constexpr uint32_t kSpatialRoadGridWidth = {spec.width};",
        f"constexpr uint32_t kSpatialRoadGridHeight = {spec.height};",
        "",
        "// Cell codes: H=highway, R=road, G=GPS-only, P=pavement, .=unknown.",
        "// Generated by tools/generate_spatial_edge_grid.py from all.traffic_persistent and all.lane_polygons.",
        f"// Counts: H={counts['H']} R={counts['R']} G={counts['G']} P={counts['P']} unknown={counts['.']}.",
        "constexpr std::array<std::string_view, kSpatialRoadGridHeight> kSpatialRoadGridRows = {{",
    ]

    for row in grid:
        lines.append(f'    "{("").join(row)}",')

    lines.extend(
        [
            "}};",
            "",
            "} // namespace EdgeWeightGPS::Generated",
            "",
        ]
    )
    path.write_text("\n".join(lines), encoding="utf-8")


def main() -> int:
    parser = argparse.ArgumentParser()
    parser.add_argument("--traffic-json", type=Path, default=Path("work/raw-segment-json/all.traffic_persistent.json"))
    parser.add_argument(
        "--lane-polygons-json",
        type=Path,
        default=Path("work/traffic-companions/json/all.lane_polygons.json"),
    )
    parser.add_argument(
        "--output",
        type=Path,
        default=Path("red4ext/EdgeWeightGPS/src/GeneratedSpatialRoadGrid.hpp"),
    )
    parser.add_argument("--cell-size", type=float, default=16.0)
    parser.add_argument("--inflate-cells", type=int, default=0)
    args = parser.parse_args()

    lanes = load_lanes(args.traffic_json)
    polygons = load_lane_polygons(args.lane_polygons_json)
    if len(lanes) != len(polygons):
        raise ValueError(f"lane/polygon count mismatch: {len(lanes)} != {len(polygons)}")

    spec, grid, counts = build_grid(lanes, polygons, args.cell_size, args.inflate_cells)
    args.output.parent.mkdir(parents=True, exist_ok=True)
    write_header(args.output, spec, grid, counts)

    total = spec.width * spec.height
    covered = total - sum(row.count(".") for row in grid)
    print(
        f"wrote {args.output} cells={total} covered={covered} "
        f"size={spec.width}x{spec.height} origin=({spec.min_x},{spec.min_y}) cell={spec.cell_size}"
    )
    print("counts " + " ".join(f"{code}={counts[code]}" for code in ("H", "R", "G", "P", ".")))
    return 0


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