Api

class BaseModel(_BaseModel, metaclass=ModelMetaclass):
    """Base model class for all models."""

    type: str = Field(default="", validate_default=True)
    _cache: dict = PrivateAttr(default_factory=dict)

    model_config = ConfigDict(
        extra="ignore",
        frozen=True,
    )

    def __init_subclass__(cls) -> None:
        MODELS[cls.__name__] = cls

    @field_validator("type")
    @classmethod
    def _validate_field_type(cls, field: str) -> str:
        field = cls.__name__
        return field

    @model_validator(mode="before")
    @classmethod
    def _validate_model(cls, obj: Any) -> Any:
        if isinstance(obj, dict):
            tpe = obj.get("type", cls.__name__)
            if tpe != cls.__name__:
                obj["type"] = tpe
                cls = MODELS.get(tpe, cls)  # noqa: PLW0642
                obj = cls.model_validate(obj)
        return obj

    @classmethod
    def model_validate(  # type: ignore[reportIncompatibleMethodOverride]  # ty: ignore[invalid-method-override]
        cls,
        obj: Any,
        *,
        strict: bool | None = None,
        from_attributes: bool | None = None,
        context: dict[str, Any] | None = None,
    ) -> Self:
        """Validate a pydantic model instance.

        Args:
            obj: The object to validate.
            strict: Whether to enforce types strictly.
            from_attributes: Whether to extract data from object attributes.
            context: Additional context to pass to the validator.

        Raises:
            ValidationError: If the object could not be validated.

        Returns:
            The validated model instance.
        """
        __tracebackhide__ = True

        if isinstance(obj, cls):
            return obj

        if isinstance(obj, dict):
            cls = MODELS.get(obj.get("type", cls.__name__), cls)  # noqa: PLW0642

        return cls.__pydantic_validator__.validate_python(
            obj, strict=strict, from_attributes=from_attributes, context=context
        )

    @classmethod
    def model_validate_json(  # type: ignore[reportIncompatibleMethodOverride]  # ty: ignore[invalid-method-override]
        cls,
        json_data: str | bytes | bytearray,
        *,
        strict: bool | None = None,
        context: dict[str, Any] | None = None,
    ) -> Self:
        """Validate a pydantic model instance.

        Args:
            json_data: The object to validate.
            strict: Whether to enforce types strictly.
            from_attributes: Whether to extract data from object attributes.
            context: Additional context to pass to the validator.

        Raises:
            ValidationError: If the object could not be validated.

        Returns:
            The validated model instance.
        """
        if isinstance(json_data, str):
            json_data = json_data.encode()
        dct = orjson.loads(json_data)
        return cls.model_validate(dct, strict=strict, context=context)

    def __hash__(self) -> int:
        to_hash = {}
        for k, v in self:
            if isinstance(v, np.ndarray):
                with suppress(ValueError):
                    v = v.item()
            if isinstance(v, np.ndarray):
                to_hash[k] = np.round(v, 9).tobytes()
            elif isinstance(v, str):
                to_hash[k] = v.encode()
            elif isinstance(v, float):
                to_hash[k] = f"{v:.9f}".encode()
            else:
                to_hash[k] = str(v).encode()
        bts = b""
        for k, v in to_hash.items():
            bts += k.encode() + b":" + v + b"|"
        return abs(int.from_bytes(md5(bts[:-1]).digest(), byteorder="big") % 10**18)  # noqa: S324

    def __eq__(self, other: object) -> bool:
        if isinstance(other, dict):
            try:
                other = self.__class__.model_validate(other)
            except Exception:  # noqa: BLE001
                return False
        elif isinstance(other, str):
            try:
                other = self.__class__.model_validate_json(other)
            except Exception:  # noqa: BLE001
                return False
        return _eq(self, other)

    def _repr(self, *, indent: int = 0, shift: int = 2) -> str:
        start = f"{self.__class__.__name__}("
        dct: dict[str, Any] = dict(self)  # don't model_dump() here.
        return _dict_repr(
            dct,
            indent=indent,
            shift=shift,
            start=start,
            end=")",
            eq="=",
        )

    def __repr__(self) -> str:
        return self._repr()

    def __str__(self) -> str:
        return self._repr()

    def _visualize(self, **_: Any) -> None:
        msg = f"visualization for {self.__class__.__name__!r} not (yet) implemented."
        raise NotImplementedError(msg)

class Box(Geometry3DBase):
    """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[Geometry2DBase]:
        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: Any,
        material_name: str,
        mesh_order: int,
        unit: float = 1e-6,
        xyz: str = "yzx",
    ) -> None:
        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: str | None = None, scale: tuple[float, float, float] | None = None
    ) -> Any:
        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[invalid-assignment]  # ty: ignore[invalid-assignment]
        return prism

class Cell(BaseModel):
    """An EME Cell.

    This defines an interval in z (the direction of propagation) within
    the simulation domain. The intersecting Structure3Ds are discretized by
    a given mesh at the center of the Cell

    """

    structures: list[Structure3D] = Field(
        description="the 3D structures which will be sliced by the cell"
    )
    mesh: Mesh2D = Field(description="the mesh to discretize the structures with")
    z_min: float = Field(description="the starting z-coordinate of the cell")
    z_max: float = Field(description="the ending z-coordinate of the cell")

    @property
    def z(self) -> float:
        """The z-position of the center of the cell."""
        return 0.5 * (self.z_min + self.z_max)

    @property
    def length(self) -> float:
        """The length of the cell."""
        return np.abs(self.z_max - self.z_min)

    @cached_property
    def materials(self) -> dict[Material, int]:
        """A mapping of the materials in the cell to their indices."""
        materials = {}
        for i, structure in enumerate(sort_structures(self.structures), start=1):
            if structure.material not in materials:
                materials[structure.material] = i
        return materials

    @cached_property
    def structures_2d(self) -> list[Structure2D]:
        """The 2D structures in the cell."""
        z = 0.5 * (self.z_min + self.z_max)
        list_of_list = [s._project(z) for s in self.structures]
        structures = [s for ss in list_of_list for s in ss]
        return structures

    @cached_property
    def m_full(self) -> IntArray2D:
        """The full material mask for the cell."""
        return _create_full_material_array(
            mesh=self.mesh,
            structures=self.structures_2d,
            materials=self.materials,
        )

    def _visualize(
        self,
        *,
        ax: Any = None,
        cbar: bool = True,
        show: bool = True,
        **_: Any,
    ) -> None:
        import matplotlib.pyplot as plt
        from matplotlib.colors import ListedColormap, to_rgba
        from mpl_toolkits.axes_grid1 import make_axes_locatable

        colors = [(0, 0, 0, 0)] + [
            to_rgba(m.meta.get("color", (0, 0, 0, 0))) for m in self.materials
        ]
        cmap = ListedColormap(colors=colors)
        if ax is not None:
            plt.sca(ax)
        else:
            ax = plt.gca()
        plt.pcolormesh(
            self.mesh.X_full,
            self.mesh.Y_full,
            self.m_full,
            cmap=cmap,
            vmin=0,
            vmax=len(self.materials) + 1,
        )
        plt.axis("scaled")
        plt.grid(visible=True)
        if cbar:
            divider = make_axes_locatable(ax)
            cax = divider.append_axes("right", size="5%", pad=0.05)
            _cbar = plt.colorbar(
                ticks=np.concatenate(
                    [np.unique(self.m_full) + 0.5, [len(self.materials) + 1.5]]
                ),
                cax=cax,
            )
            labels = [""] + [m.name for m in self.materials] + [""]
            _cbar.ax.set_yticklabels(labels, rotation=90, va="center")
            plt.sca(ax)
        if show:
            plt.show()

class CrossSection(BaseModel):
    """A `CrossSection` is built from a `Cell` with an `Environment`.

    This uniquely defines the refractive index everywhere.
    """

    structures: list[Structure2D] = Field(
        description="the 2D structures in the CrossSection"
    )
    mesh: Mesh2D = Field(description="the mesh to discretize the structures with")
    env: Environment = Field(
        description="the environment for which the cross section was calculated"
    )
    subpixel_smoothing: bool = Field(
        default=True,
        description="use subpixel smoothing at interfaces; if False: winner-takes-all",
    )
    _cell: Cell | None = PrivateAttr(default=None)

    @classmethod
    def from_cell(
        cls,
        *,
        cell: Cell,
        env: Environment,
        subpixel_smoothing: bool = True,
    ) -> Self:
        """Create a CrossSection from a Cell and Environment."""
        return cls(
            structures=cell.structures_2d,
            mesh=cell.mesh,
            env=env,
            subpixel_smoothing=subpixel_smoothing,
            _cell=cell,
        )

    @cached_property
    def materials(self) -> dict[Material, int]:
        """Return a dictionary mapping materials to their indices."""
        materials: dict[Material, int] = {}
        for i, structure in enumerate(sort_structures(self.structures), start=1):
            if structure.material not in materials:
                materials[structure.material] = i
        return materials

    @cached_property
    def _m_full(self) -> IntArray2D:
        """Return the material index array for the full mesh."""
        return _create_full_material_array(self.mesh, self.structures, self.materials)

    @cached_property
    def n_full(self) -> ComplexArray2D:
        """Return the refractive index array for the full mesh."""
        n_full = np.ones_like(self.mesh.X_full, dtype=np.complex128)
        for material, idx in self.materials.items():
            n_full = np.where(self._m_full == idx, material(self.env), n_full)
        return n_full

    @cached_property
    def nx(self) -> ComplexArray2D:
        """Return the smoothed refractive index on the Ex positions."""
        if self.subpixel_smoothing:
            return _compute_smoothed_n(
                self.mesh, self._m_full, self.materials, self.env, self.structures, "x"
            )
        return _compute_winner_takes_all_n(
            self.mesh, self._m_full, self.materials, self.env, self.structures, "x"
        )

    @cached_property
    def ny(self) -> ComplexArray2D:
        """Return the smoothed refractive index on the Ey positions."""
        if self.subpixel_smoothing:
            return _compute_smoothed_n(
                self.mesh, self._m_full, self.materials, self.env, self.structures, "y"
            )
        return _compute_winner_takes_all_n(
            self.mesh, self._m_full, self.materials, self.env, self.structures, "y"
        )

    @cached_property
    def nz(self) -> ComplexArray2D:
        """Return the smoothed refractive index on the Ez positions."""
        if self.subpixel_smoothing:
            return _compute_smoothed_n(
                self.mesh, self._m_full, self.materials, self.env, self.structures, "z"
            )
        return _compute_winner_takes_all_n(
            self.mesh, self._m_full, self.materials, self.env, self.structures, "z"
        )

    def _visualize(
        self,
        *,
        ax: Any = None,
        n_cmap: Any = None,
        cbar: bool = True,
        show: bool = True,
        **ignored: Any,
    ) -> None:
        import matplotlib.pyplot as plt  # fmt: skip
        from matplotlib import colors  # fmt: skip
        from mpl_toolkits.axes_grid1 import make_axes_locatable  # fmt: skip

        debug_grid = ignored.pop("debug_grid", False)
        if n_cmap is None:
            n_cmap = colors.LinearSegmentedColormap.from_list(
                name="c_cmap", colors=["#ffffff", "#86b5dc"]
            )
        if ax is not None:
            plt.sca(ax)
        else:
            ax = plt.gca()
        n_full = np.real(self.n_full).copy()
        n_full[0, 0] = 1.0
        plt.pcolormesh(self.mesh.X_full, self.mesh.Y_full, n_full, cmap=n_cmap)
        plt.axis("scaled")
        if not debug_grid:
            plt.grid(visible=True)
        else:
            dx = self.mesh.dx
            dy = self.mesh.dy
            x_ticks = np.sort(np.unique(self.mesh.X_full.ravel()))[::2]
            y_ticks = np.sort(np.unique(self.mesh.Y_full.ravel()))[::2]
            plt.xticks(x_ticks - 0.25 * dx, ["" for _ in x_ticks - 0.25 * dx])
            plt.yticks(y_ticks - 0.25 * dy, ["" for _ in y_ticks - 0.25 * dy])
            plt.xticks(
                x_ticks + 0.25 * dx, ["" for _ in x_ticks + 0.25 * dx], minor=True
            )
            plt.yticks(
                y_ticks + 0.25 * dy, ["" for _ in y_ticks + 0.25 * dy], minor=True
            )
            plt.grid(visible=True, which="major", ls="-")
            plt.grid(visible=True, which="minor", ls=":")
        if cbar:
            divider = make_axes_locatable(ax)
            cax = divider.append_axes("right", size="5%", pad=0.05)
            values = np.unique(np.real(self.n_full))
            _cbar = plt.colorbar(ticks=values, cax=cax)
            # material_names = ['air'] + [mat.name for mat in self.cell.materials]
            # labels = [f"\n{n}\n{v:.3f}" for n, v in zip(material_names, values)]
            labels = [f"{v:.3f}" for v in values]
            _cbar.ax.set_yticklabels(labels, rotation=90, va="center", ha="center")
            plt.sca(ax)
        if show:
            plt.show()

class Environment(BaseModel):
    """An environment contains all variables that don't depend on the structure."""

    wl: float = Field(default=1.5, description="the wavelength of the environment")
    T: float = Field(default=25.0, description="the temperature of the environment")

    model_config = ConfigDict(
        extra="allow",
        frozen=True,
    )

class GdsExtrusionRule(BaseModel):
    """A `GdsExtrusionRule` describes a single extrusion rule.

    Multiple of such rules can later be associated with a gds layer tuple.
    """

    material: Material = Field(description="the material of the extrusion")
    h_min: float = Field(description="the extrusion starting height")
    h_max: float = Field(description="the extrusion ending height")
    buffer: float = Field(
        default=0.0,
        description="an extra buffer (grow / shrink) operation applied to the polygon",
    )
    mesh_order: int = Field(
        default=5, description="the mesh order of the resulting `Structure3D`"
    )

    def __call__(self, poly: FloatArray2D) -> Structure3D:
        """Apply the extrusion rule to a polygon."""
        if self.buffer > 0:
            try:
                poly = np.asarray(sg.Polygon(poly).buffer(self.buffer).boundary.coords)
            except NotImplementedError as e:
                import gdspy  # type: ignore[unresolved-import]

                polygonset = gdspy.offset(gdspy.Polygon(poly), 0.25)
                if polygonset is None:
                    msg = (
                        "The polygon could not be offset."
                        "Please check the input polygon."
                    )
                    raise ValueError(msg) from e
                poly = polygonset.polygons[0]
        return Structure3D(
            material=self.material,
            geometry=Prism(
                poly=poly,
                h_min=self.h_min,
                h_max=self.h_max,
                axis="y",
            ),
            mesh_order=self.mesh_order,
        )

class Geometry2DBase(BaseModel):
    """Base class for 2D geometries."""

    def _mask(self, X: FloatArray2D, Y: FloatArray2D) -> BoolArray2D:
        msg = f"{self.__class__.__name__!r} cannot be masked."
        raise NotImplementedError(msg)

    def _shapely_polygon(self) -> sg.Polygon:
        msg = f"{self.__class__.__name__!r} has no shapely polygon."
        raise NotImplementedError(msg)

    def _visualize(
        self,
        *,
        ax: Any = None,
        show: bool = True,
        color: str | None = None,
        **ignored: Any,
    ) -> None:
        msg = f"{self.__class__.__name__!r} cannot be visualized."
        raise NotImplementedError(msg)

class Geometry3DBase(BaseModel):
    """Base class for 3D geometries."""

    def _project(self, z: float) -> list[Geometry2DBase]:
        msg = f"{self.__class__.__name__!r} cannot be projected to 2D."

        raise NotImplementedError(msg)

    def _visualize(
        self,
        scale: tuple[float, float, float] | None = None,
        **ignored: Any,  # noqa: ARG002
    ) -> None:
        from trimesh.scene import Scene
        from trimesh.transformations import rotation_matrix

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

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

    def _trimesh(
        self, color: str | None = None, scale: tuple[float, float, float] | None = None
    ) -> Any:
        msg = f"{self.__class__.__name__!r} cannot be visualized."
        raise NotImplementedError(msg)

class IndexMaterial(MaterialBase):
    """A material with a constant refractive index."""

    n: complex = Field(description="the refractive index of the material")

    def __call__(self, env: Environment) -> np.ndarray:  # noqa: ARG002
        """Get the refractive index of the material for the given environment."""
        return np.squeeze(np.asarray(self.n, dtype=np.complex128))

class MaterialBase(BaseModel):
    """a `Material` defines the index of a `Structure3D` in an `Environment`."""

    name: str = Field(description="the name of the material")
    meta: dict[str, Any] = Field(
        default_factory=dict, description="metadata for the material"
    )

    @model_validator(mode="after")
    def _validate_model(self) -> Self:
        MATERIALS[self.name] = self
        return self

    def __call__(self, env: Environment) -> np.ndarray:
        """Get the refractive index of the material for the given environment."""
        msg = "Please use one of the Material child classes"
        raise NotImplementedError(msg)

    def _lumadd(self, sim: Any, env: Environment, unit: float) -> str:
        from matplotlib.cm import get_cmap

        n = self(env)
        wl = np.asarray(env.wl * unit, dtype=complex).ravel()
        eps = np.asarray(n, dtype=complex).ravel() ** 2
        data = np.stack([c / wl, eps], 1)  # permittivity, not index

        if not sim.materialexists(self.name):
            sim.setmaterial(sim.addmaterial("Sampled data"), "name", self.name)
            color = np.asarray(
                self.meta.get("color") or get_cmap("jet")(np.abs(eps) / 15.0)
            )
            sim.setmaterial(self.name, "color", color)

        sim.setmaterial(self.name, "sampled data", data)

        return self.name

class Mesh2D(BaseModel):
    """A ``Mesh2D`` describes how a ``Structure3D`` is discritized into a ``Cell``."""

    x: FloatArray1D = Field(
        description="x-coords of the mesh (Ez locations, i.e. corners of the 2D cell)"
    )

    y: FloatArray1D = Field(
        description="y-coords of the mesh (Ez locations, i.e. corners of the 2D cell)"
    )

    angle_phi: float = Field(
        default=0.0,
        description=(
            "Azimuth angle of the propagation axis in the plane orthogonal to the mesh."
        ),
    )
    angle_theta: float = Field(
        default=0.0,
        description="Polar angle of the propagation axis from the injection axis.",
    )
    bend_radius: Annotated[
        float, BeforeValidator(lambda x: np.nan if x is None else x)
    ] = Field(
        default=np.nan,
        description=(
            "A curvature radius for simulation of waveguide bends. Can be negative, "
            "in which case the mode plane center has a smaller value than the "
            "curvature center along the tangential axis perpendicular to the bend axis."
        ),
    )
    bend_axis: Literal[0, 1] = Field(
        default=0,
        description=(
            "Index into the two tangential axes defining the normal to the plane in "
            "which the bend lies. This must be provided if ``bend_radius`` is not "
            "``None``. For example, for a ring in the global xy-plane, and a mode "
            "plane in either the xz or the yz plane, the ``bend_axis`` is always 1 "
            "(the global z axis)."
        ),
    )
    plane_center: tuple[float, float] = Field(
        default=(0.0, 0.0),
        description=(
            "If ``bend_radius`` is not ``None``, the position of the plane "
            "corresponding along the circonference of the circle"
        ),
    )

    num_pml: tuple[NonNegativeInt, NonNegativeInt] = Field(
        default=(0, 0),
        description="Number of standard pml layers to add in the two tangential axes.",
    )

    @cached_property
    def dx(self) -> FloatArray1D:
        """dx at Hz locations, i.e. center of the 2D cell."""
        return self.x[1:] - self.x[:-1]

    @cached_property
    def dy(self) -> FloatArray1D:
        """dy at Hz locations, i.e. center of the 2D cell."""
        return self.y[1:] - self.y[:-1]

    @cached_property
    def x_(self) -> FloatArray1D:
        """x at Hz locations, i.e. center of the 2D cell."""
        return 0.5 * (self.x[1:] + self.x[:-1])

    @cached_property
    def y_(self) -> FloatArray1D:
        """y at Hz locations, i.e. center of the 2D cell."""
        return 0.5 * (self.y[1:] + self.y[:-1])

    @cached_property
    def x_full(self) -> FloatArray1D:
        """x at half-integer locations."""
        return np.stack([self.x[:-1], self.x[:-1] + self.dx / 2], 1).ravel()

    @cached_property
    def y_full(self) -> FloatArray1D:
        """y at half-integer locations."""
        return np.stack([self.y[:-1], self.y[:-1] + self.dy / 2], 1).ravel()

    @cached_property
    def XY_full(self) -> tuple[FloatArray2D, FloatArray2D]:
        """X and Y at half-integer locations."""
        Y_full, X_full = np.meshgrid(self.y_full, self.x_full)
        return X_full, Y_full

    @property
    def X_full(self) -> FloatArray2D:
        """X at half-integer locations."""
        return self.XY_full[0]

    @property
    def Y_full(self) -> FloatArray2D:
        """Y at half-integer locations."""
        return self.XY_full[1]

    @property
    def Xx(self) -> FloatArray2D:
        """X at Ex locations."""
        return self.X_full[1::2, ::2]

    @property
    def Yx(self) -> FloatArray2D:
        """Y at Ex locations."""
        return self.Y_full[1::2, ::2]

    @property
    def Xy(self) -> FloatArray2D:
        """X at Ey locations."""
        return self.X_full[::2, 1::2]

    @property
    def Yy(self) -> FloatArray2D:
        """Y at Ey locations."""
        return self.Y_full[::2, 1::2]

    @property
    def Xz(self) -> FloatArray2D:
        """X at Ez locations."""
        return self.X_full[::2, ::2]

    @property
    def Yz(self) -> FloatArray2D:
        """Y at Ez locations."""
        return self.Y_full[::2, ::2]

    @property
    def Xz_(self) -> FloatArray2D:
        """X at Hz locations."""
        return self.X_full[1::2, 1::2]

    @property
    def Yz_(self) -> FloatArray2D:
        """Y at Hz locations."""
        return self.Y_full[1::2, 1::2]