from __future__ import annotations

import math
from typing import Literal, cast, Any
import numpy
import shapely
from shapely.geometry import Polygon, box, MultiPolygon
from shapely.ops import unary_union

from .primitives import Port



# Search Grid Snap (5.0 µm default)
SEARCH_GRID_SNAP_UM = 5.0


def snap_search_grid(value: float, snap_size: float = SEARCH_GRID_SNAP_UM) -> float:
    """
    Snap a coordinate to the nearest search grid unit.
    """
    return round(value / snap_size) * snap_size


class ComponentResult:
    """
    Standard container for generated move geometry and state.
    """
    __slots__ = (
        'geometry', 'dilated_geometry', 'proxy_geometry', 'actual_geometry',
        'end_port', 'length', 'move_type', 'bounds', 'dilated_bounds', 
        'total_bounds', 'total_dilated_bounds', 'total_bounds_box', 'total_dilated_bounds_box', '_t_cache'
    )

    def __init__(
            self,
            geometry: list[Polygon],
            end_port: Port,
            length: float,
            dilated_geometry: list[Polygon] | None = None,
            proxy_geometry: list[Polygon] | None = None,
            actual_geometry: list[Polygon] | None = None,
            skip_bounds: bool = False,
            move_type: str = 'Unknown'
            ) -> None:
        self.geometry = geometry
        self.dilated_geometry = dilated_geometry
        self.proxy_geometry = proxy_geometry
        self.actual_geometry = actual_geometry
        self.end_port = end_port
        self.length = length
        self.move_type = move_type
        self._t_cache = {}
        if not skip_bounds:
            # Vectorized bounds calculation
            self.bounds = shapely.bounds(geometry)
            # Total bounds across all polygons in the move
            self.total_bounds = numpy.array([
                numpy.min(self.bounds[:, 0]),
                numpy.min(self.bounds[:, 1]),
                numpy.max(self.bounds[:, 2]),
                numpy.max(self.bounds[:, 3])
            ])
            self.total_bounds_box = box(*self.total_bounds)
            
            if dilated_geometry is not None:
                self.dilated_bounds = shapely.bounds(dilated_geometry)
                self.total_dilated_bounds = numpy.array([
                    numpy.min(self.dilated_bounds[:, 0]),
                    numpy.min(self.dilated_bounds[:, 1]),
                    numpy.max(self.dilated_bounds[:, 2]),
                    numpy.max(self.dilated_bounds[:, 3])
                ])
                self.total_dilated_bounds_box = box(*self.total_dilated_bounds)
            else:
                self.dilated_bounds = None
                self.total_dilated_bounds = None
                self.total_dilated_bounds_box = None

    def translate(self, dx: float, dy: float) -> ComponentResult:
        """
        Create a new ComponentResult translated by (dx, dy).
        """
        dxr, dyr = round(dx, 3), round(dy, 3)
        if (dxr, dyr) == (0.0, 0.0):
            return self
        if (dxr, dyr) in self._t_cache:
            return self._t_cache[(dxr, dyr)]

        # Vectorized translation
        geoms = list(self.geometry)
        num_geom = len(self.geometry)
        
        offsets = [num_geom]
        if self.dilated_geometry is not None:
            geoms.extend(self.dilated_geometry)
        offsets.append(len(geoms))
        
        if self.proxy_geometry is not None:
            geoms.extend(self.proxy_geometry)
        offsets.append(len(geoms))
        
        if self.actual_geometry is not None:
            geoms.extend(self.actual_geometry)
        offsets.append(len(geoms))
        
        import shapely
        coords = shapely.get_coordinates(geoms)
        translated = shapely.set_coordinates(geoms, coords + [dx, dy])
        
        new_geom = list(translated[:offsets[0]])
        new_dil = list(translated[offsets[0]:offsets[1]]) if self.dilated_geometry is not None else None
        new_proxy = list(translated[offsets[1]:offsets[2]]) if self.proxy_geometry is not None else None
        new_actual = list(translated[offsets[2]:offsets[3]]) if self.actual_geometry is not None else None
        
        new_port = Port(self.end_port.x + dx, self.end_port.y + dy, self.end_port.orientation)
        res = ComponentResult(new_geom, new_port, self.length, new_dil, new_proxy, new_actual, skip_bounds=True, move_type=self.move_type)
        
        # Optimize: reuse and translate bounds
        res.bounds = self.bounds + [dx, dy, dx, dy]
        res.total_bounds = self.total_bounds + [dx, dy, dx, dy]
        res.total_bounds_box = box(*res.total_bounds)
        
        if self.dilated_bounds is not None:
            res.dilated_bounds = self.dilated_bounds + [dx, dy, dx, dy]
            res.total_dilated_bounds = self.total_dilated_bounds + [dx, dy, dx, dy]
            res.total_dilated_bounds_box = box(*res.total_dilated_bounds)
        else:
            res.total_dilated_bounds = None
            res.total_dilated_bounds_box = None
            
        self._t_cache[(dxr, dyr)] = res
        return res


class Straight:
    """
    Move generator for straight waveguide segments.
    """
    @staticmethod
    def generate(
            start_port: Port,
            length: float,
            width: float,
            snap_to_grid: bool = True,
            dilation: float = 0.0,
            snap_size: float = SEARCH_GRID_SNAP_UM,
            ) -> ComponentResult:
        """
        Generate a straight waveguide segment.
        """
        rad = numpy.radians(start_port.orientation)
        cos_val = numpy.cos(rad)
        sin_val = numpy.sin(rad)

        ex = start_port.x + length * cos_val
        ey = start_port.y + length * sin_val

        if snap_to_grid:
            ex = snap_search_grid(ex, snap_size)
            ey = snap_search_grid(ey, snap_size)

        end_port = Port(ex, ey, start_port.orientation)
        actual_length = numpy.sqrt((end_port.x - start_port.x)**2 + (end_port.y - start_port.y)**2)

        # Create polygons using vectorized points
        half_w = width / 2.0
        pts_raw = numpy.array([
            [0, half_w], 
            [actual_length, half_w], 
            [actual_length, -half_w], 
            [0, -half_w]
        ])
        
        # Rotation matrix (standard 2D rotation)
        rot_matrix = numpy.array([[cos_val, -sin_val], [sin_val, cos_val]])
        
        # Transform points: P' = R * P + T
        poly_points = (pts_raw @ rot_matrix.T) + [start_port.x, start_port.y]
        geom = [Polygon(poly_points)]

        dilated_geom = None
        if dilation > 0:
            # Direct calculation of dilated rectangle instead of expensive buffer()
            half_w_dil = half_w + dilation
            pts_dil = numpy.array([
                [-dilation, half_w_dil], 
                [actual_length + dilation, half_w_dil], 
                [actual_length + dilation, -half_w_dil], 
                [-dilation, -half_w_dil]
            ])
            poly_points_dil = (pts_dil @ rot_matrix.T) + [start_port.x, start_port.y]
            dilated_geom = [Polygon(poly_points_dil)]

        # For straight segments, geom IS the actual geometry
        return ComponentResult(geometry=geom, end_port=end_port, length=actual_length, dilated_geometry=dilated_geom, actual_geometry=geom, move_type='Straight')


def _get_num_segments(radius: float, angle_deg: float, sagitta: float = 0.01) -> int:
    """
    Calculate number of segments for an arc to maintain a maximum sagitta.
    """
    if radius <= 0:
        return 1
    ratio = max(0.0, min(1.0, 1.0 - sagitta / radius))
    theta_max = 2.0 * numpy.arccos(ratio)
    if theta_max < 1e-9:
        return 16
    num = int(numpy.ceil(numpy.radians(abs(angle_deg)) / theta_max))
    return max(8, num)


def _get_arc_polygons(
        cx: float,
        cy: float,
        radius: float,
        width: float,
        t_start: float,
        t_end: float,
        sagitta: float = 0.01,
        dilation: float = 0.0,
        ) -> list[Polygon]:
    """
    Helper to generate arc-shaped polygons using vectorized NumPy operations.
    """
    num_segments = _get_num_segments(radius, float(numpy.degrees(abs(t_end - t_start))), sagitta)
    angles = numpy.linspace(t_start, t_end, num_segments + 1)
    
    cos_a = numpy.cos(angles)
    sin_a = numpy.sin(angles)
    
    inner_radius = radius - width / 2.0 - dilation
    outer_radius = radius + width / 2.0 + dilation
    
    inner_points = numpy.stack([cx + inner_radius * cos_a, cy + inner_radius * sin_a], axis=1)
    outer_points = numpy.stack([cx + outer_radius * cos_a[::-1], cy + outer_radius * sin_a[::-1]], axis=1)
    
    # Concatenate inner and outer points to form the polygon ring
    poly_points = numpy.concatenate([inner_points, outer_points])
    
    return [Polygon(poly_points)]


def _clip_bbox(
        cx: float,
        cy: float,
        radius: float,
        width: float,
        t_start: float,
        t_end: float,
        ) -> Polygon:
    """
    Generates a rotationally invariant bounding polygon for an arc.
    """
    sweep = abs(t_end - t_start)
    if sweep > 2 * numpy.pi:
        sweep = sweep % (2 * numpy.pi)
    
    mid_angle = (t_start + t_end) / 2.0
    # Handle wrap-around for mid_angle
    if abs(t_end - t_start) > numpy.pi:
        mid_angle += numpy.pi
        
    r_out = radius + width / 2.0
    r_in = max(0.0, radius - width / 2.0)
    
    half_sweep = sweep / 2.0
    
    # Define vertices in local space (center at 0,0, symmetry axis along +X)
    # 1. Start Inner
    # 2. Start Outer
    # 3. Peak Outer Start (tangent intersection approximation)
    # 4. Peak Outer End
    # 5. End Outer
    # 6. End Inner
    # 7. Peak Inner (ensures convexity and inner clipping)
    
    # To clip the outer corner, we use two peak vertices that follow the arc tighter.
    cos_hs = numpy.cos(half_sweep)
    cos_hs2 = numpy.cos(half_sweep / 2.0)
    tan_hs2 = numpy.tan(half_sweep / 2.0)
    
    # Distance to peak from center: r_out / cos(hs/2)
    # At angles +/- hs/2
    peak_r = r_out / cos_hs2
    
    local_verts = [
        [r_in * numpy.cos(-half_sweep), r_in * numpy.sin(-half_sweep)],
        [r_out * numpy.cos(-half_sweep), r_out * numpy.sin(-half_sweep)],
        [peak_r * numpy.cos(-half_sweep/2), peak_r * numpy.sin(-half_sweep/2)],
        [peak_r * numpy.cos(half_sweep/2), peak_r * numpy.sin(half_sweep/2)],
        [r_out * numpy.cos(half_sweep), r_out * numpy.sin(half_sweep)],
        [r_in * numpy.cos(half_sweep), r_in * numpy.sin(half_sweep)],
        [r_in, 0.0]
    ]
    
    # Rotate and translate to world space
    cos_m = numpy.cos(mid_angle)
    sin_m = numpy.sin(mid_angle)
    rot = numpy.array([[cos_m, -sin_m], [sin_m, cos_m]])
    
    world_verts = (numpy.array(local_verts) @ rot.T) + [cx, cy]
    
    return Polygon(world_verts)


def _apply_collision_model(
        arc_poly: Polygon,
        collision_type: Literal["arc", "bbox", "clipped_bbox"] | Polygon,
        radius: float,
        width: float,
        cx: float = 0.0,
        cy: float = 0.0,
        clip_margin: float = 10.0,
        t_start: float | None = None,
        t_end: float | None = None,
        ) -> list[Polygon]:
    """
    Applies the specified collision model to an arc geometry.
    """
    if isinstance(collision_type, Polygon):
        return [collision_type]

    if collision_type == "arc":
        return [arc_poly]
    
    if collision_type == "clipped_bbox" and t_start is not None and t_end is not None:
        return [_clip_bbox(cx, cy, radius, width, t_start, t_end)]

    # Bounding box of the high-fidelity arc (fallback for bbox or missing angles)
    minx, miny, maxx, maxy = arc_poly.bounds
    bbox_poly = box(minx, miny, maxx, maxy)

    if collision_type == "bbox":
        return [bbox_poly]
    
    return [arc_poly]


class Bend90:
    """
    Move generator for 90-degree waveguide bends.
    """
    @staticmethod
    def generate(
            start_port: Port,
            radius: float,
            width: float,
            direction: Literal["CW", "CCW"],
            sagitta: float = 0.01,
            collision_type: Literal["arc", "bbox", "clipped_bbox"] | Polygon = "arc",
            clip_margin: float = 10.0,
            dilation: float = 0.0,
            snap_to_grid: bool = True,
            snap_size: float = SEARCH_GRID_SNAP_UM,
            ) -> ComponentResult:
        """
        Generate a 90-degree bend.
        """
        rad_start = numpy.radians(start_port.orientation)
        
        # Center of the arc
        if direction == "CCW":
            cx = start_port.x + radius * numpy.cos(rad_start + numpy.pi / 2)
            cy = start_port.y + radius * numpy.sin(rad_start + numpy.pi / 2)
            t_start = rad_start - numpy.pi / 2
            t_end = t_start + numpy.pi / 2
            new_ori = (start_port.orientation + 90) % 360
        else:
            cx = start_port.x + radius * numpy.cos(rad_start - numpy.pi / 2)
            cy = start_port.y + radius * numpy.sin(rad_start - numpy.pi / 2)
            t_start = rad_start + numpy.pi / 2
            t_end = t_start - numpy.pi / 2
            new_ori = (start_port.orientation - 90) % 360

        # Snap the end point to the grid
        ex_raw = cx + radius * numpy.cos(t_end)
        ey_raw = cy + radius * numpy.sin(t_end)
        
        if snap_to_grid:
            ex = snap_search_grid(ex_raw, snap_size)
            ey = snap_search_grid(ey_raw, snap_size)
        else:
            ex, ey = ex_raw, ey_raw
        
        # Slightly adjust radius and t_end to hit snapped point exactly
        actual_radius = numpy.sqrt((ex - cx)**2 + (ey - cy)**2)
        
        t_end_snapped = numpy.arctan2(ey - cy, ex - cx)
        # Ensure directionality and approx 90 degree sweep
        if direction == "CCW":
            while t_end_snapped <= t_start:
                t_end_snapped += 2 * numpy.pi
            while t_end_snapped > t_start + numpy.pi:
                t_end_snapped -= 2 * numpy.pi
        else:
            while t_end_snapped >= t_start:
                t_end_snapped -= 2 * numpy.pi
            while t_end_snapped < t_start - numpy.pi:
                t_end_snapped += 2 * numpy.pi
        t_end = t_end_snapped

        end_port = Port(ex, ey, new_ori)

        arc_polys = _get_arc_polygons(cx, cy, actual_radius, width, t_start, t_end, sagitta)
        collision_polys = _apply_collision_model(
            arc_polys[0], collision_type, actual_radius, width, cx, cy, clip_margin, t_start, t_end
        )

        proxy_geom = None
        if collision_type == "arc":
            # Auto-generate a clipped_bbox proxy for tiered collision checks
            proxy_geom = _apply_collision_model(
                arc_polys[0], "clipped_bbox", actual_radius, width, cx, cy, clip_margin, t_start, t_end
            )

        dilated_geom = None
        if dilation > 0:
            if collision_type == "arc":
                dilated_geom = _get_arc_polygons(cx, cy, actual_radius, width, t_start, t_end, sagitta, dilation=dilation)
            else:
                dilated_geom = [p.buffer(dilation) for p in collision_polys]

        return ComponentResult(
            geometry=collision_polys, 
            end_port=end_port, 
            length=actual_radius * numpy.abs(t_end - t_start),
            dilated_geometry=dilated_geom,
            proxy_geometry=proxy_geom,
            actual_geometry=arc_polys,
            move_type='Bend90'
        )


class SBend:
    """
    Move generator for parametric S-bends.
    """
    @staticmethod
    def generate(
            start_port: Port,
            offset: float,
            radius: float,
            width: float,
            sagitta: float = 0.01,
            collision_type: Literal["arc", "bbox", "clipped_bbox"] | Polygon = "arc",
            clip_margin: float = 10.0,
            dilation: float = 0.0,
            snap_to_grid: bool = True,
            snap_size: float = SEARCH_GRID_SNAP_UM,
            ) -> ComponentResult:
        """
        Generate a parametric S-bend (two tangent arcs).
        """
        if abs(offset) >= 2 * radius:
            raise ValueError(f"SBend offset {offset} must be less than 2*radius {2 * radius}")

        theta_init = numpy.arccos(1 - abs(offset) / (2 * radius))
        dx_init = 2 * radius * numpy.sin(theta_init)
        rad_start = numpy.radians(start_port.orientation)
        
        # Target point
        ex_raw = start_port.x + dx_init * numpy.cos(rad_start) - offset * numpy.sin(rad_start)
        ey_raw = start_port.y + dx_init * numpy.sin(rad_start) + offset * numpy.cos(rad_start)

        if snap_to_grid:
            ex = snap_search_grid(ex_raw, snap_size)
            ey = snap_search_grid(ey_raw, snap_size)
        else:
            ex, ey = ex_raw, ey_raw

        end_port = Port(ex, ey, start_port.orientation)

        # Solve for theta and radius that hit (ex, ey) exactly
        local_dx = (ex - start_port.x) * numpy.cos(rad_start) + (ey - start_port.y) * numpy.sin(rad_start)
        local_dy = -(ex - start_port.x) * numpy.sin(rad_start) + (ey - start_port.y) * numpy.cos(rad_start)
        
        # tan(theta / 2) = local_dy / local_dx
        theta = 2 * numpy.arctan2(abs(local_dy), local_dx)
        
        if abs(theta) < 1e-9:
             # Practically straight, but offset implies we need a bend.
             # If offset is also tiny, return a straight?
             if abs(offset) < 1e-6:
                 # Degenerate case: effectively straight
                 return Straight.generate(start_port, numpy.sqrt(local_dx**2 + local_dy**2), width, snap_to_grid=False, dilation=dilation)
             raise ValueError("SBend calculation failed: theta close to zero")

        # Avoid division by zero if theta is 0 (though unlikely due to offset check)
        denom = (2 * (1 - numpy.cos(theta)))
        if abs(denom) < 1e-9:
             raise ValueError("SBend calculation failed: radius denominator zero")
             
        actual_radius = abs(local_dy) / denom
        
        # Limit radius to prevent giant arcs
        if actual_radius > 100000.0:
             raise ValueError("SBend calculation failed: radius too large")

        direction = 1 if local_dy > 0 else -1
        c1_angle = rad_start + direction * numpy.pi / 2
        cx1 = start_port.x + actual_radius * numpy.cos(c1_angle)
        cy1 = start_port.y + actual_radius * numpy.sin(c1_angle)
        ts1, te1 = c1_angle + numpy.pi, c1_angle + numpy.pi + direction * theta

        c2_angle = rad_start - direction * numpy.pi / 2
        cx2 = ex + actual_radius * numpy.cos(c2_angle)
        cy2 = ey + actual_radius * numpy.sin(c2_angle)
        te2 = c2_angle + numpy.pi
        ts2 = te2 + direction * theta

        arc1 = _get_arc_polygons(cx1, cy1, actual_radius, width, ts1, te1, sagitta)[0]
        arc2 = _get_arc_polygons(cx2, cy2, actual_radius, width, ts2, te2, sagitta)[0]
        arc_polys = [arc1, arc2]

        # Use the provided collision model for primary geometry
        col1 = _apply_collision_model(arc1, collision_type, actual_radius, width, cx1, cy1, clip_margin, ts1, te1)[0]
        col2 = _apply_collision_model(arc2, collision_type, actual_radius, width, cx2, cy2, clip_margin, ts2, te2)[0]
        collision_polys = [col1, col2]

        proxy_geom = None
        if collision_type == "arc":
            # Auto-generate proxies
            p1 = _apply_collision_model(arc1, "clipped_bbox", actual_radius, width, cx1, cy1, clip_margin, ts1, te1)[0]
            p2 = _apply_collision_model(arc2, "clipped_bbox", actual_radius, width, cx2, cy2, clip_margin, ts2, te2)[0]
            proxy_geom = [p1, p2]

        dilated_geom = None
        if dilation > 0:
            if collision_type == "arc":
                d1 = _get_arc_polygons(cx1, cy1, actual_radius, width, ts1, te1, sagitta, dilation=dilation)[0]
                d2 = _get_arc_polygons(cx2, cy2, actual_radius, width, ts2, te2, sagitta, dilation=dilation)[0]
                dilated_geom = [d1, d2]
            else:
                dilated_geom = [p.buffer(dilation) for p in collision_polys]

        return ComponentResult(
            geometry=collision_polys, 
            end_port=end_port, 
            length=2 * actual_radius * theta,
            dilated_geometry=dilated_geom,
            proxy_geometry=proxy_geom,
            actual_geometry=arc_polys,
            move_type='SBend'
        )
