#!/usr/bin/env python3
from __future__ import annotations

import argparse
import json
import math
from datetime import datetime, timezone
from pathlib import Path
from typing import Any

import numpy as np
from shapely import contains_xy, line_merge, union_all
from shapely.geometry import LineString

import cm95_bathymetry_build_v1 as v1

ROOT = Path(__file__).resolve().parents[1]
DEFAULT_CONFIG = (
    ROOT
    / "config"
    / "bathymetry"
    / "cm95_bathymetry_north_kyushu_v2_smooth.yaml"
)

def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(
        description="Build smooth independent bathymetry PBF using local IDW grid interpolation."
    )
    parser.add_argument("--config", type=Path, default=DEFAULT_CONFIG)
    parser.add_argument("--output-dir", type=Path, default=None)
    return parser.parse_args()


def _grid_axis(start: float, stop: float, spacing: float) -> np.ndarray:
    count = int(math.ceil((stop - start) / spacing)) + 1
    return start + np.arange(count, dtype=np.float64) * spacing


def _bin_key(
    x: float,
    y: float,
    minx: float,
    miny: float,
    bin_size: float,
) -> tuple[int, int]:
    return (
        int(math.floor((x - minx) / bin_size)),
        int(math.floor((y - miny) / bin_size)),
    )

def build_idw_grid(
    points: dict[tuple[float, float], float],
    water_mask,
    cfg: dict[str, Any],
) -> tuple[np.ndarray, np.ndarray, np.ndarray, dict[str, Any]]:
    west, south, east, north = [float(v) for v in cfg["bbox_wgs84"]]
    minx, miny = v1.lonlat_to_mercator(west, south)
    maxx, maxy = v1.lonlat_to_mercator(east, north)

    spacing = float(cfg["grid_spacing_m"])
    radius = float(cfg["idw_search_radius_m"])
    max_nearest = float(cfg["idw_max_nearest_m"])
    min_neighbors = int(cfg["idw_min_neighbors"])
    max_neighbors = int(cfg["idw_max_neighbors"])
    power = float(cfg["idw_power"])

    xs = _grid_axis(minx, maxx, spacing)
    ys = _grid_axis(miny, maxy, spacing)
    grid = np.full((len(ys), len(xs)), np.nan, dtype=np.float64)

    point_rows = np.array(
        [(x, y, depth) for (x, y), depth in points.items()],
        dtype=np.float64,
    )
    bins: dict[tuple[int, int], list[int]] = {}

    bin_size = radius
    for index, (x, y, _depth) in enumerate(point_rows):
        bins.setdefault(
            _bin_key(float(x), float(y), minx, miny, bin_size),
            [],
        ).append(index)

    radius2 = radius * radius
    max_nearest2 = max_nearest * max_nearest
    interpolated = 0
    rejected_no_neighbors = 0
    rejected_nearest = 0
    water_nodes = 0

    for iy, y in enumerate(ys):
        row_mask = contains_xy(
            water_mask,
            xs,
            np.full(xs.shape, y, dtype=np.float64),
        )
        water_indexes = np.flatnonzero(row_mask)
        water_nodes += int(len(water_indexes))

        for ix in water_indexes:
            x = float(xs[ix])
            bx, by = _bin_key(x, float(y), minx, miny, bin_size)
            candidate_indexes: list[int] = []

            for ox in (-1, 0, 1):
                for oy in (-1, 0, 1):
                    candidate_indexes.extend(
                        bins.get((bx + ox, by + oy), ())
                    )

            if not candidate_indexes:
                rejected_no_neighbors += 1
                continue

            candidates = point_rows[candidate_indexes]
            dx = candidates[:, 0] - x
            dy = candidates[:, 1] - float(y)
            dist2 = dx * dx + dy * dy
            inside = np.flatnonzero(dist2 <= radius2)

            if len(inside) < min_neighbors:
                rejected_no_neighbors += 1
                continue

            inside = inside[np.argsort(dist2[inside])]
            if float(dist2[inside[0]]) > max_nearest2:
                rejected_nearest += 1
                continue

            inside = inside[:max_neighbors]
            chosen_d2 = dist2[inside]
            chosen_depths = candidates[inside, 2]

            if float(chosen_d2[0]) < 1.0:
                value = float(chosen_depths[0])

            else:
                weights = 1.0 / np.power(
                    np.maximum(chosen_d2, 1.0),
                    power / 2.0,
                )
                value = float(
                    np.sum(weights * chosen_depths)
                    / np.sum(weights)
                )

            grid[iy, ix] = value
            interpolated += 1

    stats = {
        "grid_width": int(len(xs)),
        "grid_height": int(len(ys)),
        "grid_nodes": int(grid.size),
        "water_nodes": water_nodes,
        "interpolated_nodes": interpolated,
        "rejected_no_neighbors": rejected_no_neighbors,
        "rejected_nearest": rejected_nearest,
        "finite_ratio": round(
            interpolated / float(grid.size),
            6,
        ),
    }
    return xs, ys, grid, stats


def _edge_intersection(
    p1: tuple[float, float],
    p2: tuple[float, float],

    v1_value: float,
    v2_value: float,
    level: float,
) -> tuple[float, float] | None:
    dv = v2_value - v1_value
    if abs(dv) < 1e-12:
        return None
    t = (level - v1_value) / dv
    if t < 0.0 or t > 1.0:
        return None
    return (
        p1[0] + t * (p2[0] - p1[0]),
        p1[1] + t * (p2[1] - p1[1]),
    )


def _march_cell(
    x0: float,
    x1: float,
    y0: float,
    y1: float,
    values: tuple[float, float, float, float],
    level: float,
) -> list[LineString]:
    v0, v1_value, v2, v3 = values
    points = (
        (x0, y0),
        (x1, y0),
        (x1, y1),
        (x0, y1),
    )

    edge_vertices = ((0, 1), (1, 2), (2, 3), (3, 0))
    hits: dict[int, tuple[float, float]] = {}

    for edge_index, (a, b) in enumerate(edge_vertices):
        va, vb = values[a], values[b]
        if (va < level and vb < level) or (va > level and vb > level):
            continue
        hit = _edge_intersection(points[a], points[b], va, vb, level)
        if hit is not None:
            hits[edge_index] = hit

    if len(hits) == 2:
        ordered = list(hits.values())
        if ordered[0] != ordered[1]:
            return [LineString(ordered)]
        return []

    if len(hits) != 4:
        return []

    case_index = 0
    for bit, value in enumerate(values):
        if value >= level:
            case_index |= 1 << bit

    center_high = sum(values) * 0.25 >= level
    if case_index == 5:
        pairs = ((0, 1), (2, 3)) if center_high else ((3, 0), (1, 2))
    elif case_index == 10:

        pairs = ((3, 0), (1, 2)) if center_high else ((0, 1), (2, 3))
    else:
        pairs = ((0, 1), (2, 3))

    return [
        LineString((hits[a], hits[b]))
        for a, b in pairs
        if hits[a] != hits[b]
    ]


def chaikin_smooth(line: LineString, iterations: int) -> LineString:
    coords = [(float(x), float(y)) for x, y, *rest in line.coords]
    if len(coords) < 3 or iterations <= 0:
        return line

    closed = coords[0] == coords[-1]
    if closed:
        coords = coords[:-1]

    for _ in range(iterations):
        if len(coords) < 2:
            break
        output: list[tuple[float, float]] = []
        if not closed:
            output.append(coords[0])
        pair_count = len(coords) if closed else len(coords) - 1
        for index in range(pair_count):
            p0 = coords[index]
            p1 = coords[(index + 1) % len(coords)]

            q = (
                0.75 * p0[0] + 0.25 * p1[0],
                0.75 * p0[1] + 0.25 * p1[1],
            )
            r = (
                0.25 * p0[0] + 0.75 * p1[0],
                0.25 * p0[1] + 0.75 * p1[1],
            )
            output.extend((q, r))
        if not closed:
            output.append(coords[-1])
        coords = output

    if closed and coords:
        coords.append(coords[0])
    return LineString(coords)


def extract_smooth_contours(
    xs: np.ndarray,
    ys: np.ndarray,
    grid: np.ndarray,
    water_mask,
    cfg: dict[str, Any],
) -> tuple[dict[float, Any], dict[str, Any]]:
    levels = sorted(float(v) for v in cfg["contour_levels_m"])
    min_segment = float(cfg["min_segment_length_m"])
    simplify_m = float(cfg.get("regional_simplify_m", 0))
    smooth_iterations = int(cfg.get("chaikin_iterations", 0))

    contours: dict[float, Any] = {}
    stats: dict[str, Any] = {}

    for level in levels:
        raw_segments: list[LineString] = []
        for iy in range(len(ys) - 1):
            y0 = float(ys[iy])
            y1 = float(ys[iy + 1])
            for ix in range(len(xs) - 1):
                values = (
                    float(grid[iy, ix]),
                    float(grid[iy, ix + 1]),
                    float(grid[iy + 1, ix + 1]),
                    float(grid[iy + 1, ix]),
                )
                if not all(math.isfinite(value) for value in values):
                    continue
                lo, hi = min(values), max(values)
                if level < lo or level > hi or lo == hi:
                    continue
                raw_segments.extend(
                    _march_cell(
                        float(xs[ix]),
                        float(xs[ix + 1]),
                        y0,
                        y1,
                        values,
                        level,
                    )
                )

        if not raw_segments:
            stats[str(level)] = {
                "raw_segments": 0,
                "line_parts": 0,
                "length_m": 0.0,
            }
            continue

        merged = line_merge(union_all(raw_segments))
        parts = []
        for part in v1.line_parts(merged):
            if part.length < min_segment:
                continue
            smooth = chaikin_smooth(part, smooth_iterations)
            if simplify_m > 0:
                smooth = smooth.simplify(
                    simplify_m,
                    preserve_topology=False,
                )
            clipped = smooth.intersection(water_mask)
            parts.extend(
                candidate
                for candidate in v1.line_parts(clipped)
                if candidate.length >= min_segment
            )

        if not parts:
            stats[str(level)] = {
                "raw_segments": len(raw_segments),
                "line_parts": 0,
                "length_m": 0.0,
            }
            continue

        geom = union_all(parts)
        contours[level] = geom
        final_parts = v1.line_parts(geom)
        stats[str(level)] = {
            "raw_segments": len(raw_segments),
            "line_parts": len(final_parts),
            "length_m": round(
                sum(part.length for part in final_parts),
                3,
            ),
        }

    return contours, stats


def main() -> int:
    args = parse_args()
    config_path = args.config.resolve()
    cfg = v1.load_config(config_path)
    output_dir = (
        args.output_dir.resolve()
        if args.output_dir is not None
        else (ROOT / str(cfg["output_dir"])).resolve()
    )
    output_dir.mkdir(parents=True, exist_ok=True)

    started = datetime.now(timezone.utc)
    print(f"BUILD {cfg['build_id']}", flush=True)
    print(f"BBOX {cfg['bbox_wgs84']}", flush=True)
    print(f"OUTPUT {output_dir}", flush=True)

    conn = v1.connect_db()
    try:
        points, sounding_stats = v1.fetch_soundings(conn, cfg)
        print(
            "SOUNDINGS",
            json.dumps(sounding_stats, ensure_ascii=False),
            flush=True,
        )
        water_mask, water_rows = v1.fetch_water_mask(conn, cfg)
        print(f"WATER_AREAS {water_rows}", flush=True)
    finally:
        conn.close()

    xs, ys, grid, surface_stats = build_idw_grid(
        points,
        water_mask,
        cfg,
    )
    print(
        "SURFACE",
        json.dumps(surface_stats, ensure_ascii=False),
        flush=True,
    )

    contours, contour_stats = extract_smooth_contours(
        xs,
        ys,
        grid,
        water_mask,
        cfg,
    )

    if not contours:
        raise RuntimeError("no smooth contours generated")
    print(
        "CONTOURS",
        json.dumps(contour_stats, ensure_ascii=False),
        flush=True,
    )

    debug_feature_count = v1.write_debug_geojson(
        contours,
        cfg,
        output_dir,
    )
    v1.write_tilejson(cfg, output_dir)
    tile_stats = v1.build_tiles(contours, cfg, output_dir)
    print(
        "TILES",
        json.dumps(tile_stats, ensure_ascii=False),
        flush=True,
    )

    finished = datetime.now(timezone.utc)
    manifest = {
        "schema": "navsea.cm95.bathymetry-build-result/v2",
        "build_id": str(cfg["build_id"]),
        "generated_at": finished.isoformat(),
        "duration_seconds": round(
            (finished - started).total_seconds(),
            3,
        ),

        "dataset_id": int(cfg["dataset_id"]),
        "bbox_wgs84": [float(v) for v in cfg["bbox_wgs84"]],
        "source_layer": str(cfg["source_layer"]),
        "source_kind": str(cfg["source_kind"]),
        "reliability": str(cfg["reliability"]),
        "contour_levels_m": [
            float(v) for v in cfg["contour_levels_m"]
        ],
        "native_scale_max": int(cfg["native_scale_max"]),
        "water_object_classes": [
            str(v) for v in cfg["water_object_classes"]
        ],
        "surface_method": "local_idw_regular_grid",
        "surface_parameters": {
            key: cfg[key]
            for key in (
                "grid_spacing_m",
                "idw_search_radius_m",
                "idw_max_nearest_m",
                "idw_min_neighbors",
                "idw_max_neighbors",
                "idw_power",
                "chaikin_iterations",
            )
        },

        "sounding_stats": sounding_stats,
        "water_area_rows": water_rows,
        "surface_stats": surface_stats,
        "contour_stats": contour_stats,
        "debug_geojson_features": debug_feature_count,
        "tile_stats": tile_stats,
    }
    (output_dir / "build_manifest.json").write_text(
        json.dumps(
            manifest,
            ensure_ascii=False,
            indent=2,
        ),
        encoding="utf-8",
    )
    print(
        "DONE",
        output_dir / "build_manifest.json",
        flush=True,
    )
    return 0


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