Port

class BasePort(BaseModel, arbitrary_types_allowed=True):
    """Class representing the base port.

    This does not have any knowledge of units.
    """

    name: str
    kcl: KCLayout
    cross_section: SymmetricalCrossSection
    trans: kdb.Trans | None = None
    dcplx_trans: kdb.DCplxTrans | None = None
    info: Info = Info()
    port_type: str

    @model_validator(mode="after")
    def check_exclusivity(self) -> Self:
        """Check if the port has a valid transformation."""
        if self.trans is None and self.dcplx_trans is None:
            raise ValueError("Both trans and dcplx_trans cannot be None.")
        if self.trans is not None and self.dcplx_trans is not None:
            raise ValueError("Only one of trans or dcplx_trans can be set.")
        return self

    def __copy__(self) -> BasePort:
        """Copy the BasePort."""
        return BasePort(
            name=self.name,
            kcl=self.kcl,
            cross_section=self.cross_section,
            trans=self.trans.dup() if self.trans else None,
            dcplx_trans=self.dcplx_trans.dup() if self.dcplx_trans else None,
            info=self.info.model_copy(),
            port_type=self.port_type,
        )

    def transformed(
        self,
        trans: kdb.Trans | kdb.DCplxTrans = kdb.Trans.R0,
        post_trans: kdb.Trans | kdb.DCplxTrans = kdb.Trans.R0,
    ) -> BasePort:
        """Get a transformed copy of the BasePort."""
        base = self.__copy__()
        if (
            base.trans is not None
            and isinstance(trans, kdb.Trans)
            and isinstance(post_trans, kdb.Trans)
        ):
            base.trans = trans * base.trans * post_trans
            base.dcplx_trans = None
            return base
        if isinstance(trans, kdb.Trans):
            trans = kdb.DCplxTrans(trans.to_dtype(self.kcl.dbu))
        if isinstance(post_trans, kdb.Trans):
            post_trans = kdb.DCplxTrans(post_trans.to_dtype(self.kcl.dbu))
        dcplx_trans = self.dcplx_trans or kdb.DCplxTrans(
            t=self.trans.to_dtype(self.kcl.dbu)  # ty:ignore[unresolved-attribute]
        )

        base.trans = None
        base.dcplx_trans = trans * dcplx_trans * post_trans
        return base

    def transform(
        self,
        trans: kdb.Trans | kdb.DCplxTrans = kdb.Trans.R0,
        post_trans: kdb.Trans | kdb.DCplxTrans = kdb.Trans.R0,
    ) -> Self:
        """Transform self."""
        base = self
        if (
            base.trans is not None
            and isinstance(trans, kdb.Trans)
            and isinstance(post_trans, kdb.Trans)
        ):
            base.trans = trans * base.trans * post_trans
            base.dcplx_trans = None
            return self
        if isinstance(trans, kdb.Trans):
            trans = kdb.DCplxTrans(trans.to_dtype(self.kcl.dbu))
        if isinstance(post_trans, kdb.Trans):
            post_trans = kdb.DCplxTrans(post_trans.to_dtype(self.kcl.dbu))
        dcplx_trans = self.dcplx_trans or kdb.DCplxTrans(
            t=self.trans.to_dtype(self.kcl.dbu)  # ty:ignore[unresolved-attribute]
        )

        base.trans = None
        base.dcplx_trans = trans * dcplx_trans * post_trans
        return self

    @model_serializer()
    def ser_model(self) -> BasePortDict:
        """Serialize the BasePort."""
        trans = self.trans.dup() if self.trans is not None else None
        dcplx_trans = self.dcplx_trans.dup() if self.dcplx_trans is not None else None
        return BasePortDict(
            name=self.name,
            kcl=self.kcl,
            cross_section=self.cross_section,
            trans=trans,
            dcplx_trans=dcplx_trans,
            info=self.info.model_copy(),
            port_type=self.port_type,
        )

    def get_trans(self) -> kdb.Trans:
        """Get the transformation."""
        if self.trans is not None:
            return self.trans
        assert self.dcplx_trans is not None, "Both trans and dcplx_trans are None"
        return kdb.ICplxTrans(trans=self.dcplx_trans, dbu=self.kcl.dbu).s_trans()

    def get_dcplx_trans(self) -> kdb.DCplxTrans:
        """Get the complex transformation."""
        if self.dcplx_trans is not None:
            return self.dcplx_trans
        assert self.trans is not None, "Both trans and dcplx_trans are None"
        return kdb.DCplxTrans(self.trans.to_dtype(self.kcl.dbu))

    def __eq__(self, other: object) -> bool:
        """Check if two ports are equal."""
        if not isinstance(other, BasePort):
            return False
        return (
            (self.trans is None and other.trans is None)
            or (
                (
                    self.trans is not None
                    and other.trans is not None
                    and self.trans == other.trans
                )
                and (self.dcplx_trans is None and other.dcplx_trans is None)
            )
            or (
                (
                    self.dcplx_trans is not None
                    and other.dcplx_trans is not None
                    and self.dcplx_trans == other.dcplx_trans
                )
                and self.name == other.name
                and self.kcl == other.kcl
                and self.cross_section == other.cross_section
                and self.port_type == other.port_type
                and self.info == other.info
            )
        )

    def check_connection(
        self,
        other: BasePort,
        tolerance: float = 0.1,
        angle_tolerance: float = 0.01,
        snapped: bool = False,
    ) -> int:
        tol_um = self.kcl.dbu * tolerance
        check: int = 0
        if snapped or (self.trans is not None and other.trans is not None):
            t1 = self.get_trans()
            t2 = other.get_trans()
            if t1.disp == t2.disp:
                check += PortCheck.position
            orientation = (t1.angle - t2.angle) % 4
            if orientation == 2:
                check += PortCheck.opposite
            elif orientation == 0:
                check += PortCheck.same
        else:
            dt1 = self.get_dcplx_trans()
            dt2 = other.get_dcplx_trans()
            if (dt1.disp - dt2.disp).length() < tol_um:
                check += PortCheck.position
            angle_diff = (dt1.angle - dt2.angle) % 360
            if abs(angle_diff - 180) < angle_tolerance:
                check += PortCheck.opposite
            elif abs(angle_diff) < angle_tolerance:
                check += PortCheck.same
        if self.cross_section == other.cross_section:
            check += PortCheck.cross_section
            check += PortCheck.layer
            check += PortCheck.width
        else:
            if self.cross_section.main_layer.is_equivalent(
                other.cross_section.main_layer
            ):
                check += PortCheck.layer
            if self.cross_section.width == other.cross_section.width:
                check += PortCheck.width
        if self.port_type == other.port_type:
            check += PortCheck.port_type

        return check

class BasePortDict(TypedDict):
    """TypedDict for the BasePort."""

    name: str
    kcl: KCLayout
    cross_section: SymmetricalCrossSection
    trans: kdb.Trans | None
    dcplx_trans: kdb.DCplxTrans | None
    info: Info
    port_type: str

class DIRECTION(IntEnum):
    """Alias for KLayout direction to compass directions."""

    E = 0
    N = 1
    W = 2
    S = 3

class DPort(ProtoPort[float]):
    """A port is the photonics equivalent to a pin in electronics.

    In addition to the location and layer
    that defines a pin, a port also contains an orientation and a width.
    This can be fully represented with a transformation, integer and layer_index.


    Attributes:
        name: String to name the port.
        width: The width of the port in dbu.
        trans: Transformation in dbu. If the port can be represented in 90° intervals
            this is the safe way to do so.
        dcplx_trans: Transformation in micrometer. The port will autoconvert between
            trans and dcplx_trans on demand.
        port_type: A string defining the type of the port
        layer: Index of the layer or a LayerEnum that acts like an integer, but can
            contain layer number and datatype
        info: A dictionary with additional info. Not reflected in GDS. Copy will make a
            (shallow) copy of it.
        d: Access port info in micrometer basis such as width and center / angle.
        kcl: Link to the layout this port resides in.
    """

    @overload
    def __init__(
        self,
        name: str,
        *,
        width: float,
        layer: LayerEnum | int,
        trans: kdb.Trans | str,
        kcl: KCLayout | None = None,
        port_type: str = "optical",
        info: dict[str, int | float | str] = ...,
    ) -> None: ...

    @overload
    def __init__(
        self,
        name: str,
        *,
        width: float,
        layer: LayerEnum | int,
        dcplx_trans: kdb.DCplxTrans | str,
        kcl: KCLayout | None = None,
        port_type: str = "optical",
        info: dict[str, int | float | str] = ...,
    ) -> None: ...

    @overload
    def __init__(
        self,
        name: str,
        *,
        width: float,
        layer: LayerEnum | int,
        port_type: str = "optical",
        orientation: float,
        center: tuple[float, float] = (0, 0),
        mirror_x: bool = False,
        kcl: KCLayout | None = None,
        info: dict[str, int | float | str] = ...,
    ) -> None: ...

    @overload
    def __init__(
        self,
        name: str,
        *,
        width: float,
        layer_info: kdb.LayerInfo,
        trans: kdb.Trans | str,
        kcl: KCLayout | None = None,
        port_type: str = "optical",
        info: dict[str, int | float | str] = ...,
    ) -> None: ...

    @overload
    def __init__(
        self,
        name: str,
        *,
        width: float,
        layer_info: kdb.LayerInfo,
        dcplx_trans: kdb.DCplxTrans | str,
        kcl: KCLayout | None = None,
        port_type: str = "optical",
        info: dict[str, int | float | str] = ...,
    ) -> None: ...

    @overload
    def __init__(
        self,
        name: str,
        *,
        width: float,
        layer_info: kdb.LayerInfo,
        port_type: str = "optical",
        orientation: float,
        center: tuple[float, float] = (0, 0),
        mirror_x: bool = False,
        kcl: KCLayout | None = None,
        info: dict[str, int | float | str] = ...,
    ) -> None: ...

    @overload
    def __init__(
        self,
        name: str,
        *,
        cross_section: DCrossSection | SymmetricalCrossSection,
        port_type: str = "optical",
        orientation: float,
        center: tuple[float, float],
        mirror_x: bool = False,
        kcl: KCLayout | None = None,
        info: dict[str, int | float | str] = ...,
    ) -> None: ...

    @overload
    def __init__(
        self,
        name: str,
        *,
        cross_section: DCrossSection | SymmetricalCrossSection,
        trans: kdb.Trans | str,
        kcl: KCLayout | None = None,
        info: dict[str, int | float | str] = ...,
        port_type: str = "optical",
    ) -> None: ...

    @overload
    def __init__(
        self,
        name: str,
        *,
        cross_section: DCrossSection | SymmetricalCrossSection,
        dcplx_trans: kdb.DCplxTrans | str,
        kcl: KCLayout | None = None,
        info: dict[str, int | float | str] = ...,
        port_type: str = "optical",
    ) -> None: ...

    @overload
    def __init__(self, *, base: BasePort) -> None: ...

    @overload
    def __init__(self, *, port: ProtoPort[Any]) -> None: ...

    def __init__(
        self,
        name: str | None = None,
        *,
        width: float | None = None,
        layer: int | None = None,
        layer_info: kdb.LayerInfo | None = None,
        port_type: str = "optical",
        trans: kdb.Trans | str | None = None,
        dcplx_trans: kdb.DCplxTrans | str | None = None,
        orientation: float = 0,
        center: tuple[float, float] = (0, 0),
        mirror_x: bool = False,
        port: ProtoPort[Any] | None = None,
        kcl: KCLayout | None = None,
        info: dict[str, int | float | str] | None = None,
        cross_section: DCrossSection | SymmetricalCrossSection | None = None,
        base: BasePort | None = None,
    ) -> None:
        """Create a port from dbu or um based units."""
        if info is None:
            info = {}
        if base is not None:
            self._base = base
            return
        if port is not None:
            self._base = port.base.__copy__()
            return

        if name is None:
            raise ValueError(
                "DPort must have a name. Only when passing another port or port base"
                " name can be None."
            )
        info_ = Info(**info)

        from .layout import get_default_kcl

        kcl_ = kcl or get_default_kcl()
        if cross_section is None:
            if layer_info is None:
                if layer is None:
                    raise ValueError("layer or layer_info for a port must be defined")
                layer_info = kcl_.layout.get_info(layer)
            if width is None:
                raise ValueError(
                    "If a cross_section is not given a width must be defined."
                )
            width_ = kcl_.to_dbu(width)
            if width_ % 2:
                raise ValueError(
                    f"width needs to be even to snap to grid. Got {width}."
                    "Ports must have a grid width of multiples of 2."
                )
            cross_section_ = kcl_.get_symmetrical_cross_section(
                CrossSectionSpec(layer=layer_info, width=kcl_.to_dbu(width))
            )
        elif isinstance(cross_section, SymmetricalCrossSection):
            cross_section_ = cross_section
        else:
            cross_section_ = cross_section.base
        if trans is not None:
            trans_ = kdb.Trans.from_s(trans) if isinstance(trans, str) else trans.dup()
            self._base = BasePort(
                name=name,
                kcl=kcl_,
                cross_section=cross_section_,
                trans=trans_,
                info=info_,
                port_type=port_type,
            )
        elif dcplx_trans is not None:
            if isinstance(dcplx_trans, str):
                dcplx_trans_ = kdb.DCplxTrans.from_s(dcplx_trans)
            else:
                dcplx_trans_ = dcplx_trans.dup()
            self._base = BasePort(
                name=name,
                kcl=kcl_,
                cross_section=cross_section_,
                trans=kdb.Trans.R0,
                info=info_,
                port_type=port_type,
            )
            self.dcplx_trans = dcplx_trans_
        else:
            assert center is not None
            dcplx_trans_ = kdb.DCplxTrans.R0
            self._base = BasePort(
                name=name,
                kcl=kcl_,
                cross_section=cross_section_,
                dcplx_trans=dcplx_trans_,
                info=info_,
                port_type=port_type,
            )
            self.center = center
            self.orientation = orientation
            self.mirror_x = mirror_x

    def copy(
        self,
        trans: kdb.Trans | kdb.DCplxTrans = kdb.Trans.R0,
        post_trans: kdb.Trans | kdb.DCplxTrans = kdb.Trans.R0,
    ) -> DPort:
        """Get a copy of a port.

        Transformation order which results in `copy.trans`:
            - Trans: `trans * port.trans * post_trans`
            - DCplxTrans: `trans * port.dcplx_trans * post_trans`

        Args:
            trans: an optional transformation applied to the port to be copied.
            post_trans: transformation to apply to the port after copying.

        Returns:
            port: a copy of the port
        """
        return DPort(base=self._base.transformed(trans=trans, post_trans=post_trans))

    def copy_polar(
        self,
        d: float = 0,
        d_orth: float = 0,
        orientation: float = 180,
        mirror: bool = False,
    ) -> DPort:
        """Get a polar copy of the port.

        This will return a port which is transformed relatively to the original port's
        transformation (orientation, angle and position).

        Args:
            d: The distance to the old port
            d_orth: Orthogonal distance (positive is positive y for a port which is
                facing angle=0°)
            orientation: Relative angle to the original port, in degrees.
            mirror: Whether to mirror the port relative to the original port.
        """
        return self.copy(
            post_trans=kdb.DCplxTrans(rot=orientation, mirrx=mirror, x=d, y=d_orth)
        )

    @property
    def x(self) -> float:
        """X coordinate of the port in um."""
        return self.dx

    @x.setter
    def x(self, value: float) -> None:
        self.dx = value

    @property
    def y(self) -> float:
        """Y coordinate of the port in um."""
        return self.dy

    @y.setter
    def y(self, value: float) -> None:
        self.dy = value

    @property
    def width(self) -> float:
        """Width of the port in um."""
        return self.dwidth

    @property
    def cross_section(self) -> DCrossSection:
        """Get the cross section of the port."""
        return DCrossSection(kcl=self._base.kcl, base=self._base.cross_section)

    @cross_section.setter
    def cross_section(
        self, value: SymmetricalCrossSection | TCrossSection[Any]
    ) -> None:
        if isinstance(value, SymmetricalCrossSection):
            self._base.cross_section = value
            return
        self._base.cross_section = value.base

class Port(ProtoPort[int]):
    """A port is the photonics equivalent to a pin in electronics.

    In addition to the location and layer
    that defines a pin, a port also contains an orientation and a width.
    This can be fully represented with a transformation, integer and layer_index.


    Attributes:
        name: String to name the port.
        width: The width of the port in dbu.
        trans: Transformation in dbu. If the port can be represented in 90° intervals
            this is the safe way to do so.
        dcplx_trans: Transformation in micrometer. The port will autoconvert between
            trans and dcplx_trans on demand.
        port_type: A string defining the type of the port
        layer: Index of the layer or a LayerEnum that acts like an integer, but can
            contain layer number and datatype
        info: A dictionary with additional info. Not reflected in GDS. Copy will make a
            (shallow) copy of it.
        d: Access port info in micrometer basis such as width and center / angle.
        kcl: Link to the layout this port resides in.
    """

    @overload
    def __init__(
        self,
        *,
        name: str,
        width: int,
        layer: LayerEnum | int,
        trans: kdb.Trans | str,
        kcl: KCLayout | None = None,
        port_type: str = "optical",
        info: dict[str, int | float | str] = ...,
    ) -> None: ...

    @overload
    def __init__(
        self,
        *,
        name: str,
        width: int,
        layer: LayerEnum | int,
        dcplx_trans: kdb.DCplxTrans | str,
        kcl: KCLayout | None = None,
        port_type: str = "optical",
        info: dict[str, int | float | str] = ...,
    ) -> None: ...

    @overload
    def __init__(
        self,
        name: str,
        *,
        width: int,
        layer: LayerEnum | int,
        port_type: str = "optical",
        angle: int,
        center: tuple[int, int],
        mirror_x: bool = False,
        kcl: KCLayout | None = None,
        info: dict[str, int | float | str] = ...,
    ) -> None: ...

    @overload
    def __init__(
        self,
        name: str,
        *,
        width: int,
        layer_info: kdb.LayerInfo,
        trans: kdb.Trans | str,
        kcl: KCLayout | None = None,
        port_type: str = "optical",
        info: dict[str, int | float | str] = ...,
    ) -> None: ...

    @overload
    def __init__(
        self,
        name: str,
        *,
        width: int,
        layer_info: kdb.LayerInfo,
        dcplx_trans: kdb.DCplxTrans | str,
        kcl: KCLayout | None = None,
        port_type: str = "optical",
        info: dict[str, int | float | str] = ...,
    ) -> None: ...

    @overload
    def __init__(
        self,
        name: str,
        *,
        width: int,
        layer_info: kdb.LayerInfo,
        port_type: str = "optical",
        angle: int,
        center: tuple[int, int],
        mirror_x: bool = False,
        kcl: KCLayout | None = None,
        info: dict[str, int | float | str] = ...,
    ) -> None: ...

    @overload
    def __init__(
        self,
        name: str,
        *,
        cross_section: CrossSection | SymmetricalCrossSection,
        port_type: str = "optical",
        angle: int,
        center: tuple[int, int],
        mirror_x: bool = False,
        kcl: KCLayout | None = None,
        info: dict[str, int | float | str] = ...,
    ) -> None: ...

    @overload
    def __init__(
        self,
        name: str,
        *,
        cross_section: CrossSection | SymmetricalCrossSection,
        trans: kdb.Trans | str,
        kcl: KCLayout | None = None,
        info: dict[str, int | float | str] = ...,
        port_type: str = "optical",
    ) -> None: ...

    @overload
    def __init__(
        self,
        name: str,
        *,
        cross_section: CrossSection | SymmetricalCrossSection,
        dcplx_trans: kdb.DCplxTrans | str,
        kcl: KCLayout | None = None,
        info: dict[str, int | float | str] = ...,
        port_type: str = "optical",
    ) -> None: ...

    @overload
    def __init__(self, *, base: BasePort) -> None: ...

    @overload
    def __init__(self, *, port: ProtoPort[Any]) -> None: ...

    def __init__(
        self,
        name: str | None = None,
        *,
        width: int | None = None,
        layer: int | None = None,
        layer_info: kdb.LayerInfo | None = None,
        port_type: str = "optical",
        trans: kdb.Trans | str | None = None,
        dcplx_trans: kdb.DCplxTrans | str | None = None,
        angle: int | None = None,
        center: tuple[int, int] | None = None,
        mirror_x: bool = False,
        port: ProtoPort[Any] | None = None,
        kcl: KCLayout | None = None,
        info: dict[str, int | float | str] | None = None,
        cross_section: CrossSection | SymmetricalCrossSection | None = None,
        base: BasePort | None = None,
    ) -> None:
        """Create a port from dbu or um based units."""
        if info is None:
            info = {}
        if base is not None:
            self._base = base
            return
        if port is not None:
            self._base = port.base.__copy__()
            return

        if name is None:
            raise ValueError(
                "Port must have a name. Only when passing another port or port base"
                " name can be None."
            )
        info_ = Info(**info)
        from .layout import get_default_kcl

        kcl_ = kcl or get_default_kcl()
        if cross_section is None:
            if layer_info is None:
                if layer is None:
                    raise ValueError("layer or layer_info for a port must be defined")
                layer_info = kcl_.layout.get_info(layer)
            if width is None:
                raise ValueError(
                    "any width and layer, or a cross_section must be given if the"
                    " 'port is None'"
                )
            cross_section_ = kcl_.get_symmetrical_cross_section(
                CrossSectionSpec(layer=layer_info, width=width)
            )
        elif isinstance(cross_section, SymmetricalCrossSection):
            cross_section_ = cross_section
        else:
            cross_section_ = cross_section.base
        if trans is not None:
            trans_ = kdb.Trans.from_s(trans) if isinstance(trans, str) else trans.dup()
            self._base = BasePort(
                name=name,
                kcl=kcl_,
                cross_section=cross_section_,
                trans=trans_,
                info=info_,
                port_type=port_type,
            )
        elif dcplx_trans is not None:
            if isinstance(dcplx_trans, str):
                dcplx_trans_ = kdb.DCplxTrans.from_s(dcplx_trans)
            else:
                dcplx_trans_ = dcplx_trans.dup()
            self._base = BasePort(
                name=name,
                kcl=kcl_,
                cross_section=cross_section_,
                trans=kdb.Trans.R0,
                info=info_,
                port_type=port_type,
            )
            self.dcplx_trans = dcplx_trans_
        elif angle is not None:
            assert center is not None
            trans_ = kdb.Trans(angle, mirror_x, *center)
            self._base = BasePort(
                name=name,
                kcl=kcl_,
                cross_section=cross_section_,
                trans=trans_,
                info=info_,
                port_type=port_type,
            )
        else:
            raise ValueError("Missing port parameters given")

    def copy(
        self,
        trans: kdb.Trans | kdb.DCplxTrans = kdb.Trans.R0,
        post_trans: kdb.Trans | kdb.DCplxTrans = kdb.Trans.R0,
    ) -> Port:
        """Get a copy of a port.

        Transformation order which results in `copy.trans`:
            - Trans: `trans * port.trans * post_trans`
            - DCplxTrans: `trans * port.dcplx_trans * post_trans`

        Args:
            trans: an optional transformation applied to the port to be copied.
            post_trans: transformation to apply to the port after copying.

        Returns:
            port: a copy of the port
        """
        return Port(base=self._base.transformed(trans=trans, post_trans=post_trans))

    def copy_polar(
        self, d: int = 0, d_orth: int = 0, angle: int = 2, mirror: bool = False
    ) -> Port:
        """Get a polar copy of the port.

        This will return a port which is transformed relatively to the original port's
        transformation (orientation, angle and position).

        Args:
            d: The distance to the old port
            d_orth: Orthogonal distance (positive is positive y for a port which is
                facing angle=0°)
            angle: Relative angle to the original port (0=0°,1=90°,2=180°,3=270°).
            mirror: Whether to mirror the port relative to the original port.

        Returns:
            Port copied relative to it's current position and angle/orientation.
        """
        return self.copy(post_trans=kdb.Trans(angle, mirror, d, d_orth))

    @property
    def x(self) -> int:
        """X coordinate of the port in dbu."""
        return self.ix

    @x.setter
    def x(self, value: int) -> None:
        self.ix = value

    @property
    def y(self) -> int:
        """Y coordinate of the port in dbu."""
        return self.iy

    @y.setter
    def y(self, value: int) -> None:
        self.iy = value

    @property
    def width(self) -> int:
        """Width of the port in um."""
        return self.iwidth

    @property
    def cross_section(self) -> CrossSection:
        """Get the cross section of the port."""
        return CrossSection(kcl=self._base.kcl, base=self._base.cross_section)

    @cross_section.setter
    def cross_section(
        self, value: SymmetricalCrossSection | TCrossSection[Any]
    ) -> None:
        if isinstance(value, SymmetricalCrossSection):
            self._base.cross_section = value
            return
        self._base.cross_section = value.base