"""Marching squares algorithm for finding iso-contours on 2D grids.

This module implements the marching squares algorithm for extracting contour
lines (iso-lines) from a 2D scalar field. The contours are returned in
fractional grid coordinates for sub-pixel precision.

Example::

    from gri_utils import terrain
    import numpy as np

    # Create a scalar field (e.g., distance from center)
    x = np.linspace(-1, 1, 100)
    y = np.linspace(-1, 1, 100)
    X, Y = np.meshgrid(x, y)
    values = X**2 + Y**2 - 0.5  # Circle of radius sqrt(0.5)

    # Find the zero contour
    contours = terrain.find_contours(values, level=0.0)

    # Each contour is a (K, 2) array of (row, col) fractional coordinates
    for contour in contours:
        print(f"Contour with {len(contour)} points")
"""

from __future__ import annotations

from typing import TYPE_CHECKING

import numpy as np

if TYPE_CHECKING:
    from collections.abc import Sequence

    from numpy.typing import NDArray


# Marching squares lookup table
# Each cell has 4 corners: top-left, top-right, bottom-right, bottom-left
# Binary encoding: bit 0 = TL, bit 1 = TR, bit 2 = BR, bit 3 = BL
# Value 1 means corner is above the level, 0 means below
# Table gives list of edge pairs to connect: (edge1, edge2)
# Edges: 0=top, 1=right, 2=bottom, 3=left
_EDGE_TABLE: dict[int, list[tuple[int, int]]] = {
    0: [],  # All below
    1: [(0, 3)],  # TL above
    2: [(0, 1)],  # TR above
    3: [(1, 3)],  # TL, TR above
    4: [(1, 2)],  # BR above
    5: [(0, 1), (2, 3)],  # TL, BR above (saddle)
    6: [(0, 2)],  # TR, BR above
    7: [(2, 3)],  # TL, TR, BR above
    8: [(2, 3)],  # BL above
    9: [(0, 2)],  # TL, BL above
    10: [(0, 3), (1, 2)],  # TR, BL above (saddle)
    11: [(1, 2)],  # TL, TR, BL above
    12: [(1, 3)],  # BR, BL above
    13: [(0, 1)],  # TL, BR, BL above
    14: [(0, 3)],  # TR, BR, BL above
    15: [],  # All above
}


def find_contours(
    values: NDArray[np.floating] | Sequence,
    level: float = 0.0,
) -> list[NDArray[np.floating]]:
    """Find iso-contours at given level using marching squares.

    Args:
        values: Scalar field values. Shape (M, N).
        level: Iso-contour level to find. Default 0.0.

    Returns:
        List of contour polylines. Each contour is an array of shape (K, 2)
        containing (row, col) coordinates in fractional grid indices.
        Contours may be:
        - Open: starting/ending at grid boundary
        - Closed: forming a complete loop
    """
    values = np.asarray(values, dtype=np.float64)
    m, n = values.shape

    if m < 2 or n < 2:  # noqa: PLR2004
        return []

    # Subtract level to find zero crossings
    field = values - level

    # Build all edge segments
    segments = _build_segments(field)

    if not segments:
        return []

    # Chain segments into contours
    return _chain_segments(segments)


def _build_segments(
    field: NDArray[np.floating],
) -> list[tuple[tuple[float, float], tuple[float, float]]]:
    """Build all edge segments from the scalar field."""
    m, n = field.shape
    segments = []

    for i in range(m - 1):
        for j in range(n - 1):
            # Get corner values
            tl = field[i, j]
            tr = field[i, j + 1]
            br = field[i + 1, j + 1]
            bl = field[i + 1, j]

            # Compute cell index
            cell_type = 0
            if tl > 0:
                cell_type |= 1
            if tr > 0:
                cell_type |= 2
            if br > 0:
                cell_type |= 4
            if bl > 0:
                cell_type |= 8

            # Handle saddle points by checking center value
            # Cases 5 and 10 are saddle points in marching squares
            # The lookup table already handles saddle point edge connections,
            # so we don't need to modify cell_type based on center value

            # Get edges to draw
            edges = _EDGE_TABLE[cell_type]

            for edge1, edge2 in edges:
                # Interpolate positions on edges
                p1 = _edge_position(i, j, edge1, tl, tr, br, bl)
                p2 = _edge_position(i, j, edge2, tl, tr, br, bl)
                segments.append((p1, p2))

    return segments


def _edge_position(  # noqa: PLR0913
    i: int,
    j: int,
    edge: int,
    tl: float,
    tr: float,
    br: float,
    bl: float,
) -> tuple[float, float]:
    """Get interpolated position on cell edge.

    Args:
        i: Cell row index (top-left corner).
        j: Cell column index (top-left corner).
        edge: Edge index (0=top, 1=right, 2=bottom, 3=left).
        tl: Top-left corner value.
        tr: Top-right corner value.
        br: Bottom-right corner value.
        bl: Bottom-left corner value.

    Returns:
        (row, col) fractional coordinates.
    """
    if edge == 0:  # Top edge (TL to TR)
        t = _interpolate(tl, tr)
        return (float(i), float(j) + t)
    if edge == 1:  # Right edge (TR to BR)
        t = _interpolate(tr, br)
        return (float(i) + t, float(j) + 1)
    if edge == 2:  # Bottom edge (BL to BR)  # noqa: PLR2004
        t = _interpolate(bl, br)
        return (float(i) + 1, float(j) + t)
    # Left edge (TL to BL)
    t = _interpolate(tl, bl)
    return (float(i) + t, float(j))


_INTERP_TOLERANCE = 1e-12


def _interpolate(v1: float, v2: float) -> float:
    """Linear interpolation to find zero crossing between two values."""
    if abs(v2 - v1) < _INTERP_TOLERANCE:
        return 0.5
    return -v1 / (v2 - v1)


def _chain_segments(  # noqa: PLR0912, PLR0915, C901
    segments: list[tuple[tuple[float, float], tuple[float, float]]],
) -> list[NDArray[np.floating]]:
    """Chain disconnected segments into contour polylines."""
    if not segments:
        return []

    # Build adjacency structure using rounded coordinates as keys
    # Use small tolerance for matching endpoints
    def key(p: tuple[float, float]) -> tuple[int, int]:
        return (round(p[0] * 10000), round(p[1] * 10000))

    # Create mapping from endpoint to segment indices
    endpoint_to_segments: dict[tuple[int, int], list[int]] = {}
    for idx, (p1, p2) in enumerate(segments):
        k1, k2 = key(p1), key(p2)
        endpoint_to_segments.setdefault(k1, []).append(idx)
        endpoint_to_segments.setdefault(k2, []).append(idx)

    used = [False] * len(segments)
    contours = []

    for start_idx in range(len(segments)):
        if used[start_idx]:
            continue

        # Start a new contour
        contour_points = []
        p1, p2 = segments[start_idx]
        contour_points.append(p1)
        contour_points.append(p2)
        used[start_idx] = True

        # Extend forward from p2
        current_end = p2
        while True:
            k = key(current_end)
            candidates = endpoint_to_segments.get(k, [])
            next_idx = None
            for idx in candidates:
                if not used[idx]:
                    next_idx = idx
                    break

            if next_idx is None:
                break

            used[next_idx] = True
            seg_p1, seg_p2 = segments[next_idx]

            # Determine which end connects to current_end
            if key(seg_p1) == k:
                contour_points.append(seg_p2)
                current_end = seg_p2
            else:
                contour_points.append(seg_p1)
                current_end = seg_p1

        # Extend backward from p1
        current_end = p1
        backward_points = []
        while True:
            k = key(current_end)
            candidates = endpoint_to_segments.get(k, [])
            next_idx = None
            for idx in candidates:
                if not used[idx]:
                    next_idx = idx
                    break

            if next_idx is None:
                break

            used[next_idx] = True
            seg_p1, seg_p2 = segments[next_idx]

            if key(seg_p1) == k:
                backward_points.append(seg_p2)
                current_end = seg_p2
            else:
                backward_points.append(seg_p1)
                current_end = seg_p1

        # Combine backward + forward
        if backward_points:
            backward_points.reverse()
            contour_points = backward_points + contour_points

        contours.append(np.array(contour_points, dtype=np.float64))

    return contours


def contours_to_xyz(
    contours: list[NDArray[np.floating]],
    sheet_xyz: NDArray[np.floating] | Sequence,
    *,
    order: int = 3,
) -> list[NDArray[np.floating]]:
    """Convert grid-coordinate contours to XYZ coordinates via sheet interpolation.

    Args:
        contours: List of contours from find_contours, each (K, 2) in grid coords.
        sheet_xyz: XYZ sheet in ECEF meters. Shape (M, N, 3).
        order: Interpolation order. 0=nearest, 1=bilinear, 3=bicubic (default).

    Returns:
        List of contours in XYZ coordinates, each (K, 3) in ECEF meters.
    """
    from .sheet import interpolate_sheet_batch  # noqa: PLC0415

    sheet_xyz = np.asarray(sheet_xyz, dtype=np.float64)
    xyz_contours = []

    for contour in contours:
        if len(contour) == 0:
            xyz_contours.append(np.zeros((0, 3), dtype=np.float64))
            continue

        row_fracs = contour[:, 0]
        col_fracs = contour[:, 1]
        xyz_points = interpolate_sheet_batch(
            sheet_xyz,
            row_fracs,
            col_fracs,
            order=order,
        )
        xyz_contours.append(xyz_points)

    return xyz_contours
