from typing import List, Tuple, Optional

import numpy as np

from backend.contour_models import ContourLine, ContourSegment


# ===== Contour Generation =====
class MarchingSquares:

    # ===== Public API =====
    @staticmethod
    def generate_contours(grid: np.ndarray, levels: List[float]) -> List[ContourLine]:
        """
        Generate contour lines for a 2D grid at specified levels.

        Args:
            grid: 2D numpy array (shape: [width, height])
            levels: List of contour levels to compute

        Returns:
            List of ContourLine objects
        """
        if grid.size == 0:
            return []

        width, height = grid.shape[0], grid.shape[1]
        contours = []

        for level in levels:
            contour = ContourLine(level=level)

            # Iterate over each cell (x, y) in the grid
            for y in range(height - 1):
                for x in range(width - 1):
                    v0 = float(grid[x, y])  # top-left
                    v1 = float(grid[x + 1, y])  # top-right
                    v2 = float(grid[x + 1, y + 1])  # bottom-right
                    v3 = float(grid[x, y + 1])  # bottom-left

                    if any(np.isnan([v0, v1, v2, v3])):
                        continue  # Skip cells with NaN values

                    state = (
                        (1 if v0 > level else 0)
                        | (2 if v1 > level else 0)
                        | (4 if v2 > level else 0)
                        | (8 if v3 > level else 0)
                    )

                    seg = MarchingSquares._get_segment(
                        x, y, v0, v1, v2, v3, level, state
                    )
                    if seg is not None:
                        contour.segments.append(seg)

            contours.append(contour)

        return contours

    # ===== Geometry Helpers =====
    @staticmethod
    def _lerp(
        x1: float,
        y1: float,
        x2: float,
        y2: float,
        val1: float,
        val2: float,
        level: float,
    ) -> Tuple[float, float]:
        """Linear interpolation between (x1,y1) and (x2,y2)."""
        if val2 == val1:
            return (x1 + x2) / 2, (y1 + y2) / 2
        t = (level - val1) / (val2 - val1)
        px = x1 + t * (x2 - x1)
        py = y1 + t * (y2 - y1)
        return px, py

    # ===== Segment Lookup =====
    @staticmethod
    def _get_segment(
        x: int,
        y: int,
        v0: float,
        v1: float,
        v2: float,
        v3: float,
        level: float,
        state: int,
    ) -> Optional[ContourSegment]:
        """Return a ContourSegment for the given cell and state."""
        if state in (10,):  # Ambiguous case — skip
            return None

        # Map C# states to Python logic
        if state == 1 or state == 14:
            start = MarchingSquares._lerp(x, y, x + 1, y, v0, v1, level)
            end = MarchingSquares._lerp(x, y, x, y + 1, v0, v3, level)
        elif state == 2 or state == 13:
            start = MarchingSquares._lerp(x + 1, y, x + 1, y + 1, v1, v2, level)
            end = MarchingSquares._lerp(x, y, x + 1, y, v0, v1, level)
        elif state == 3 or state == 12:
            start = MarchingSquares._lerp(x, y, x, y + 1, v0, v3, level)
            end = MarchingSquares._lerp(x + 1, y, x + 1, y + 1, v1, v2, level)
        elif state == 4 or state == 11:
            start = MarchingSquares._lerp(x + 1, y, x + 1, y + 1, v1, v2, level)
            end = MarchingSquares._lerp(x + 1, y + 1, x, y + 1, v2, v3, level)
        elif state == 5:
            start = MarchingSquares._lerp(x, y, x, y + 1, v0, v3, level)
            end = MarchingSquares._lerp(x + 1, y, x + 1, y + 1, v1, v2, level)
        elif state == 6 or state == 9:
            start = MarchingSquares._lerp(x, y, x + 1, y, v0, v1, level)
            end = MarchingSquares._lerp(x + 1, y + 1, x, y + 1, v2, v3, level)
        elif state == 7 or state == 8:
            start = MarchingSquares._lerp(x, y, x, y + 1, v0, v3, level)
            end = MarchingSquares._lerp(x + 1, y + 1, x, y + 1, v2, v3, level)
        else:
            return None

        return ContourSegment(
            start_x=start[0], start_y=start[1], end_x=end[0], end_y=end[1]
        )

    # ===== Color Mapping =====
    @staticmethod
    def color_from_level(
        value: float, min_val: float, max_val: float
    ) -> Tuple[int, int, int]:
        """Return (R, G, B) tuple for a value between min and max."""
        range_val = max_val - min_val
        if range_val == 0:
            t = 0.5
        else:
            t = max(0.0, min(1.0, (value - min_val) / range_val))

        r = int(255 * t)
        b = int(255 * (1 - t))
        return (r, 0, b)