import warnings
from secrets import token_hex
from typing import Dict, List, Literal, Tuple, Union, cast

import matplotlib.pyplot as plt
import numpy as np
import shapely.geometry as sg
from matplotlib.patches import Rectangle as MplRect
from pydantic.v1 import Field, validator

from meow.base_model import BaseModel

GEOMETRIES_3D: Dict[str, type] = {}
GEOMETRIES_2D: Dict[str, type] = {}
AxisDirection = Union[Literal["x"], Literal["y"], Literal["z"]]


class Geometry2D(BaseModel):
    type: str = ""

    def __init_subclass__(cls):
        GEOMETRIES_2D[cls.__name__] = cls

    def __new__(cls, **kwargs):
        cls = GEOMETRIES_2D.get(kwargs.get("type", cls.__name__), cls)
        return BaseModel.__new__(cls)  # type: ignore

    @validator("type", pre=True, always=True)
    def validate_type(cls, value):
        if not value:
            value = getattr(cls, "__name__", "Geometry")
        if value not in GEOMETRIES_2D:
            raise ValueError(
                f"Invalid Geometry type. Got: {value!r}. Valid types: {GEOMETRIES_2D}."
            )
        return value

    def _mask(self, X, Y):
        raise NotImplementedError(f"{self.__class__.__name__!r} cannot be masked.")

    def _visualize(self, color=None):
        raise NotImplementedError(f"{self.__class__.__name__!r} cannot be visualized.")


class Rectangle(Geometry2D):
    """a Rectangle"""

    x_min: float = Field(description="the minimum x-value of the box")
    x_max: float = Field(description="the maximum x-value of the box")
    y_min: float = Field(description="the minimum y-value of the box")
    y_max: float = Field(description="the maximum y-value of the box")

    def _mask(self, X, Y):
        mask = (
            (self.x_min <= X)
            & (X <= self.x_max)
            & (self.y_min <= Y)
            & (Y <= self.y_max)
        )
        return mask

    def _visualize(self, ax=None, show=True, color=None):
        if ax is None:
            ax = plt.gca()
        if color is None:
            color = "grey"
        width = self.x_max - self.x_min
        height = self.y_max - self.y_min
        mpl_rect = MplRect(
            xy=(self.x_min, self.y_min),
            width=width,
            height=height,
            color=color,
        )
        ax.add_patch(mpl_rect)
        ax.set_xlim(self.x_min - 0.1 * width, self.x_max + 0.1 * width)
        ax.set_ylim(self.y_min - 0.1 * width, self.y_max + 0.1 * width)
        ax.set_xlabel("x")
        ax.set_ylabel("y")
        plt.grid(True)
        if show:
            plt.show()


class Geometry3D(BaseModel):
    type: str = ""

    def __init_subclass__(cls):
        GEOMETRIES_3D[cls.__name__] = cls

    def __new__(cls, **kwargs):
        cls = GEOMETRIES_3D.get(kwargs.get("type", cls.__name__), cls)
        return BaseModel.__new__(cls)  # type: ignore

    @validator("type", pre=True, always=True)
    def validate_type(cls, value):
        if not value:
            value = getattr(cls, "__name__", "Geometry")
        if value not in GEOMETRIES_3D:
            raise ValueError(
                f"Invalid Geometry type. Got: {value!r}. Valid types: {GEOMETRIES_3D}."
            )
        return value

    def _project(self, z: float) -> List[Geometry2D]:
        raise NotImplementedError(f"{self.__class__.__name__!r} cannot be projected.")

    def _visualize(self, scale=None):
        from trimesh.scene import Scene  # fmt: skip
        from trimesh.transformations import rotation_matrix  # fmt: skip

        scene = Scene()
        scene.add_geometry(self._trimesh(scale=scale))
        scene.apply_transform(rotation_matrix(-np.pi / 6, (0, 1, 0)))
        return scene.show()

    def _lumadd(self, sim, material_name, mesh_order, unit=1e-6, xyz="yzx"):
        raise NotImplementedError(
            f"{self.__class__.__name__!r} cannot be added to Lumerical."
        )

    def _trimesh(self, **kwargs):
        raise NotImplementedError(f"{self.__class__.__name__!r} cannot be visualized.")


class Box(Geometry3D):
    """A Box is a simple rectangular cuboid"""

    x_min: float = Field(description="the minimum x-value of the box")
    x_max: float = Field(description="the maximum x-value of the box")
    y_min: float = Field(description="the minimum y-value of the box")
    y_max: float = Field(description="the maximum y-value of the box")
    z_min: float = Field(description="the minimum z-value of the box")
    z_max: float = Field(description="the maximum z-value of the box")

    def _project(self, z: float) -> List[Geometry2D]:
        if z < self.z_min or z > self.z_max:
            return []
        rect = Rectangle(
            x_min=self.x_min,
            x_max=self.x_max,
            y_min=self.y_min,
            y_max=self.y_max,
        )
        return [rect]

    def _lumadd(self, sim, material_name, mesh_order, unit, xyz):
        x, y, z = xyz
        name = token_hex(4)
        kwargs = {
            f"{x}_min": float(self.x_min * unit),
            f"{x}_max": float(self.x_max * unit),
            f"{y}_min": float(self.y_min * unit),
            f"{y}_max": float(self.y_max * unit),
            f"{z}_min": float(self.z_min * unit),
            f"{z}_max": float(self.z_max * unit),
        }
        sim.addrect(
            name=name,
            material=material_name,
            override_mesh_order_from_material_database=True,
            mesh_order=mesh_order,
            **kwargs,
        )

    def _trimesh(self, color=None, scale=None):
        from trimesh import Trimesh  # fmt: skip
        from trimesh.creation import extrude_polygon  # fmt: skip

        sx, sy, sz = scale or (1, 1, 1)
        poly = sg.Polygon(
            [
                (self.x_min * sx, self.y_min * sy),
                (self.x_min * sx, self.y_max * sy),
                (self.x_max * sx, self.y_max * sy),
                (self.x_max * sx, self.y_min * sy),
            ],
        )
        prism = extrude_polygon(poly, self.z_max * sz - self.z_min * sz)
        prism = cast(Trimesh, prism.apply_translation((0, 0, self.z_min * sz)))
        if color is not None:
            prism.visual.face_colors = _to_rgba(color)  # type: ignore
        return prism


class Prism(Geometry3D):
    """A prism is a 2D Polygon extruded along a certain axis direction ('x', 'y', or 'z')."""

    poly: np.ndarray[Tuple[int, Literal[2]], np.dtype[np.float_]] = Field(
        description="the 2D array (Nx2) with polygon vertices"
    )
    h_min: float = Field(description="the start height of the extrusion")
    h_max: float = Field(description="the end height of the extrusion")
    axis: AxisDirection = Field(
        default="y",
        description="the axis along which the polygon will be extruded ('x', 'y', or 'z').",
    )

    def _project_axis_x(self, z):
        # x, y, z -> y, z, x
        poly = sg.Polygon(self.poly)
        y_min, _ = self.poly.min(0)
        y_max, _ = self.poly.max(0)
        line = sg.LineString([(y_min, z), (y_max, z)])
        with warnings.catch_warnings():
            warnings.filterwarnings("ignore", category=RuntimeWarning, module="shapely")
            intersections = poly.intersection(line)

        if not isinstance(intersections, sg.MultiLineString):
            intersection_array = np.asarray(intersections.coords)
            if not intersection_array.shape[0]:
                return []
            intersections = sg.MultiLineString([intersections])

        geoms_2d = []
        for intersection in intersections.geoms:
            intersection = np.asarray(intersection.coords)
            if not intersection.shape[0]:
                continue
            (y_min, _), (y_max, _) = intersection
            y_min, y_max = min(y_min, y_max), max(y_min, y_max)
            x_min, x_max = min(self.h_min, self.h_max), max(self.h_min, self.h_max)
            rect = Rectangle(
                x_min=x_min,
                x_max=x_max,
                y_min=y_min,
                y_max=y_max,
            )
            geoms_2d.append(rect)
        return geoms_2d

    def _project_axis_y(self, z):
        # x, y, z -> z, x, y
        poly = sg.Polygon(self.poly)
        _, x_min = self.poly.min(0)
        _, x_max = self.poly.max(0)
        line = sg.LineString([(z, x_min), (z, x_max)])
        with warnings.catch_warnings():
            warnings.filterwarnings("ignore", category=RuntimeWarning, module="shapely")
            intersections = poly.intersection(line)

        if not isinstance(intersections, sg.MultiLineString):
            intersection_array = np.asarray(intersections.coords)
            if not intersection_array.shape[0]:
                return []
            intersections = sg.MultiLineString([intersections])

        geoms_2d = []
        for intersection in intersections.geoms:
            intersection = np.asarray(intersection.coords)
            if not intersection.shape[0]:
                continue
            (_, x_min), (_, x_max) = intersection
            x_min, x_max = min(x_min, x_max), max(x_min, x_max)
            y_min, y_max = min(self.h_min, self.h_max), max(self.h_min, self.h_max)
            rect = Rectangle(
                x_min=x_min,
                x_max=x_max,
                y_min=y_min,
                y_max=y_max,
            )
            geoms_2d.append(rect)
        return geoms_2d

    def _project_axis_z(self, z):
        # x, y, z -> x, y, z
        if z < self.h_min or z < self.h_max:
            return []
        else:
            return [self.poly]

    def _project(self, z):
        if self.axis == "x":
            return self._project_axis_x(z)
        elif self.axis == "y":
            return self._project_axis_y(z)
        else:
            return self._project_axis_z(z)

    def _lumadd(self, sim, material_name, mesh_order, unit=1e-6, xyz="yzx"):
        name = token_hex(4)

        if xyz not in ("xyz", "yzx", "zxy"):
            raise ValueError(
                f"Prism axes should be positively oriented when adding to Lumerical. Got: {xyz!r}"
            )

        sim.addpoly(
            name=name,
            material=material_name,
            override_mesh_order_from_material_database=True,
            mesh_order=mesh_order,
            use_relative_coordinates=True,
            vertices=np.asarray(self.poly, float) * float(unit),
            x=0.0,
            y=0.0,
            z_min=float(self.h_min * unit),
            z_max=float(self.h_max * unit),
        )

        x, y, z = xyz
        if self.axis == "x":
            z, x, y = x, y, z
            # raise NotImplementedError(
            #    "Only prisms extruded perpendicular to the 'chip surface' are currently supported in Lumerical."
            # )
        elif self.axis == "y":
            y, z, x = x, y, z
        xyz = f"{x}{y}{z}"

        if xyz in ["yzx", "zxy"]:
            raise NotImplementedError(
                "Only prisms extruded perpendicular to the 'chip surface' are currently supported in Lumerical."
            )

    def _trimesh(self, color=None, scale=None):
        from trimesh.creation import extrude_polygon  # fmt: skip

        poly = sg.Polygon(self.poly)
        prism = extrude_polygon(poly, self.h_max - self.h_min)
        prism = prism.apply_translation((0, 0, self.h_min))
        if self.axis == "x":
            prism.vertices = np.roll(prism.vertices, shift=1, axis=1)  # type: ignore
        if self.axis == "y":
            prism.vertices = np.roll(prism.vertices, shift=-1, axis=1)  # type: ignore
        if scale is not None:
            sx, sy, sz = scale
            prism.vertices *= np.array([[sx, sy, sz]])  # type: ignore
        if color is not None:
            prism.visual.face_colors = _to_rgba(color)  # type: ignore
        return prism

    def _center(self):
        import shapely.geometry as sg  # fmt: skip

        a, b = np.array(sg.Polygon(self.poly).centroid.xy).ravel()
        c = 0.5 * (self.h_min + self.h_max)
        x, y, z = {
            "x": (c, a, b),
            "y": (b, c, a),
            "z": (a, b, c),
        }[self.axis]
        return x, y, z

    def _axis_tuple(self):
        return {
            "x": (1, 0, 0),
            "y": (0, 1, 0),
            "z": (0, 0, 1),
        }[self.axis]


def _to_rgba(c):
    from matplotlib.colors import to_rgba as _to_rgba_mpl  # fmt: skip
    r, g, b, a = _to_rgba_mpl(c)
    a = min(max(a, 0.1), 0.9)
    return float(r), float(g), float(b), float(a)
