class PRUC:
    """Class representing a Polarization-Rotating Unit Cell (PRUC) for reflectarray antennas.

    This class defines the geometry of a unit cell containing four arrow-shaped patches
    connected by vias through a ground plane. The cell can be rotated 0° or 180° to
    encode binary phase information.

    Attributes
    ----------
    S : float
        Size of the square unit cell aperture in microns.
    s : float
        Spacing between arrow patches and unit cell boundary in microns.
    l_head : float
        Length of the arrow head in microns.
    l_stem : float
        Length from via center to arrow tip in microns.
    w_stem : float
        Width of the arrow stem in microns.
    w_head : float
        Width of the arrow head in microns.
    w_via_pad : float
        Width of the square via pad on top layer in microns.
    w_cnx : float
        Width of the bottom layer connections between vias in microns.
    d_via : float
        Diameter of the metal vias in microns.
    d_hole : float
        Diameter of the holes in the ground plane in microns.
    metal_t : float
        Thickness of all metal layers in microns.
    z_top : float
        Z-coordinate of top substrate surface (excluding metal) in microns.
    z_ground : float
        Z-coordinate of ground plane substrate surface in microns.
    z_bottom : float
        Z-coordinate of bottom substrate surface (excluding metal) in microns.
    bit_version : bool
        If True, rotates bottom connections 180° to encode binary '1' state.
        If False, keeps 0° rotation for binary '0' state.
    """

    def __init__(
        self,
        S: float,
        s: float,
        l_head: float,
        l_stem: float,
        w_stem: float,
        w_head: float,
        w_via_pad: float,
        w_cnx: float,
        d_via: float,
        d_hole: float,
        metal_t: float,
        z_top: float,
        z_ground: float,
        z_bottom: float,
        bit_version: bool,
    ):
        self.S = S
        self.s = s
        self.l_head = l_head
        self.l_stem = l_stem
        self.w_stem = w_stem
        self.w_head = w_head
        self.w_via_pad = w_via_pad
        self.w_cnx = w_cnx
        self.d_via = d_via
        self.d_hole = d_hole
        self.metal_t = metal_t
        self.z_top = z_top
        self.z_ground = z_ground
        self.z_bottom = z_bottom
        self.bit_version = bit_version

    def calc_via_offset(self):
        """Calculates the offset of the via in the unit cell"""
        return self.S / 2 - self.s - self.l_stem / np.sqrt(2)

    def make_single_arrow(self) -> list[tuple[float, float]]:
        """Returns the 2D coordinates of a single arrow in the unit cell corresponding with the first quadrant."""

        def swap_xy(xy: tuple[float, float]):
            """Applies a reflection along the y = x line."""
            return (xy[1], xy[0])

        via_offset = self.l_stem / np.sqrt(2)

        # Create points along half of an arrow
        tr = (0, 0)
        br = (0, -self.l_head)
        br2 = (-self.w_head, -self.l_head)
        in_corner = (-self.w_head, -self.w_head - self.w_stem / np.sqrt(2))
        bl = (
            -via_offset + self.w_stem / np.sqrt(8),
            -via_offset - self.w_stem / np.sqrt(8),
        )
        # Complete arrow is created by reflecting points across y = x line
        arrow_points = [
            tr,
            br,
            br2,
            in_corner,
            bl,
            swap_xy(bl),
            swap_xy(in_corner),
            swap_xy(br2),
            swap_xy(br),
            swap_xy(tr),
        ]
        # Use shapely polygon operations to create the union of the arrow and via pad.
        arrow = shapely.Polygon(arrow_points)
        via_pad = shapely.box(
            -via_offset - 0.5 * self.w_via_pad,
            -via_offset - 0.5 * self.w_via_pad,
            -via_offset + 0.5 * self.w_via_pad,
            -via_offset + 0.5 * self.w_via_pad,
        )
        arrow_shapely = shapely.union(arrow, via_pad)
        return arrow_shapely.exterior.coords

    def make_all_arrows(self) -> list[td.PolySlab]:
        """Returns a list of all arrows in the unit cell by using rotations and translations."""
        slab_bounds = (self.z_top, self.z_top + self.metal_t)
        arrow_tr_verts = np.array(self.make_single_arrow())
        arrow_tl_verts = np.copy(arrow_tr_verts)
        arrow_tl_verts[:, 0] *= -1
        arrow_bl_verts = np.copy(arrow_tl_verts)
        arrow_bl_verts[:, 1] *= -1
        arrow_br_verts = np.copy(arrow_tr_verts)
        arrow_br_verts[:, 1] *= -1

        corner_offset = self.S / 2 - self.s

        arrow_tr_verts += np.array([corner_offset, corner_offset])
        arrow_tr = td.PolySlab(vertices=arrow_tr_verts, axis=2, slab_bounds=slab_bounds)
        arrow_tl_verts += np.array([-corner_offset, corner_offset])
        arrow_tl = td.PolySlab(vertices=arrow_tl_verts, axis=2, slab_bounds=slab_bounds)
        arrow_bl_verts += np.array([-corner_offset, -corner_offset])
        arrow_bl = td.PolySlab(vertices=arrow_bl_verts, axis=2, slab_bounds=slab_bounds)
        arrow_br_verts += np.array([corner_offset, -corner_offset])
        arrow_br = td.PolySlab(vertices=arrow_br_verts, axis=2, slab_bounds=slab_bounds)
        return [arrow_tr, arrow_tl, arrow_bl, arrow_br]

    def make_bottom_layer(self) -> td.Geometry:
        """Returns the bottom layer of the unit cell"""
        via_center = self.calc_via_offset()

        # Horizontal connections
        zcenter = self.z_bottom - self.metal_t / 2
        thickness = self.metal_t
        strip = td.Box(
            center=(0, 0, zcenter),
            size=(
                2 * np.sqrt(2) * (via_center) + self.d_via + (self.w_cnx - self.d_via),
                self.w_cnx,
                thickness,
            ),
        )
        angle = -np.pi / 4
        if self.bit_version:
            angle = np.pi / 4

        rotate_45 = td.Transformed.rotation(axis=2, angle=angle)
        strip = td.Transformed(geometry=strip, transform=rotate_45)

        geo_tuple = tuple([strip])

        metal_regions = td.GeometryGroup(geometries=geo_tuple)
        return metal_regions

    def make_ground_layer(self, ground_width: float = None) -> td.Geometry:
        """Returns the ground layer of the unit cell"""
        zcenter = self.z_ground + self.metal_t / 2
        thickness = self.metal_t
        via_center = self.calc_via_offset()

        ground_size = (self.S, self.S, thickness)
        if ground_width is not None:
            ground_size = (ground_width, ground_width, thickness)

        ground = td.Box(center=(0, 0, zcenter), size=ground_size)
        hole_size = (self.d_hole, self.d_hole, thickness)
        hole_tr = td.Box(center=(via_center, via_center, zcenter), size=hole_size)
        hole_tl = td.Box(center=(-via_center, via_center, zcenter), size=hole_size)
        hole_bl = td.Box(center=(-via_center, -via_center, zcenter), size=hole_size)
        hole_br = td.Box(center=(via_center, -via_center, zcenter), size=hole_size)
        holes = td.GeometryGroup(geometries=(hole_tr, hole_tl, hole_bl, hole_br))
        ground_with_holes = td.ClipOperation(
            operation="difference", geometry_a=ground, geometry_b=holes
        )
        return ground_with_holes

    def make_vias(self) -> tuple[td.Cylinder, td.Cylinder, td.Cylinder, td.Cylinder]:
        """Returns the vias in the unit cell"""
        via_center = self.calc_via_offset()
        center_z = (self.z_top + self.z_bottom) / 2
        length = self.z_top - self.z_bottom
        radius = self.d_via / 2
        tr = td.Cylinder(
            center=(via_center, via_center, center_z),
            length=length,
            radius=radius,
            axis=2,
        )
        tl = td.Cylinder(
            center=(-via_center, via_center, center_z),
            length=length,
            radius=radius,
            axis=2,
        )
        bl = td.Cylinder(
            center=(-via_center, -via_center, center_z),
            length=length,
            radius=radius,
            axis=2,
        )
        br = td.Cylinder(
            center=(via_center, -via_center, center_z),
            length=length,
            radius=radius,
            axis=2,
        )
        return tr, tl, bl, br

    def make_unit_cell_geometry(
        self, dl_hole: float, ground_width: float = None
    ) -> tuple[
        td.GeometryGroup,
        list[td.MeshOverrideStructure],
        list[tuple[float, float, float]],
    ]:
        """Returns the geometry of the metallic portion of the unit cell,
        mesh overrides around the holes, and the vertices of the arrows."""
        arrows = self.make_all_arrows()
        bottom_geometry = self.make_bottom_layer()
        ground_with_holes = self.make_ground_layer(ground_width)
        via_geometry = self.make_vias()
        all_geometries = tuple(
            arrows + [bottom_geometry] + [ground_with_holes] + list(via_geometry)
        )
        geometry_group = td.GeometryGroup(geometries=all_geometries)

        # Make regions of mesh refinement
        mesh_overrides = [
            td.MeshOverrideStructure(geometry=hole, dl=[dl_hole, dl_hole, self.metal_t])
            for hole in ground_with_holes.geometry_b.geometries
        ]

        # Get arrow vertices for snapping points
        arrow_verts = [
            (vert[0], vert[1], self.z_top)
            for arrow in arrows
            for vert in arrow.vertices
        ]
        return geometry_group, mesh_overrides, arrow_verts

class PRUCArray:
    """Class representing an array of Polarization-Rotating Unit Cells (PRUCs).

    This class handles the construction of a reflectarray consisting of multiple PRUCs,
    including substrate layers and mesh refinement specifications.

    Attributes
    ----------
    unit_cell_0 : PRUC
        Unit cell representing binary '0' state with 0° rotation.
    unit_cell_1 : PRUC
        Unit cell representing binary '1' state with 180° rotation.
    Nx : int
        Number of unit cells along x-axis.
    Ny : int
        Number of unit cells along y-axis.
    array_sx : float
        Size of array in x-direction in microns. Allows for substrates that extend farther than the unit cells.
    array_sy : float
        Size of array in y-direction in microns. Allows for substrates that extend farther than the unit cells.
    prepreg_t : float
        Thickness of prepreg layers in microns.
    l1_t : float
        Thickness of top substrate layer in microns.
    l2_t : float
        Thickness of middle substrate layer in microns.
    l3_t : float
        Thickness of bottom substrate layer in microns.
    substrate_medium : Medium
        Material properties of the substrate layers.
    prepreg_medium : Medium
        Material properties of the prepreg layers.
    metal_medium : Medium
        Material properties of the metal layers.
    dl_hole : float
        Mesh cell size for refinement around holes in microns.
    dl_xy : float
        Mesh cell size in xy-plane in microns.
    dl_z : float
        Mesh cell size in z-direction in microns.
    """

    def __init__(
        self,
        unit_cell_0: PRUC,
        unit_cell_1: PRUC,
        Nx,
        Ny,
        array_sx,
        array_sy,
        prepreg_t,
        l1_t,
        l2_t,
        l3_t,
        substrate_medium,
        prepreg_medium,
        metal_medium,
        dl_hole: float,
        dl_xy: float,
        dl_z: float,
    ):
        self.unit_cell_0 = unit_cell_0
        self.unit_cell_1 = unit_cell_1
        self.Nx = Nx
        self.Ny = Ny
        self.array_sx = array_sx
        self.array_sy = array_sy
        self.prepreg_t = prepreg_t
        self.l1_t = l1_t
        self.l2_t = l2_t
        self.l3_t = l3_t
        self.substrate_medium = substrate_medium
        self.prepreg_medium = prepreg_medium
        self.metal_medium = metal_medium
        self.dl_hole = dl_hole
        self.dl_xy = dl_xy
        self.dl_z = dl_z

    def make_substrates(
        self,
    ) -> tuple[list[td.Structure], list[td.MeshOverrideStructure]]:
        """Create the layered substrate structure for the reflectarray along
        with mesh overrides that refine the mesh along the normal axis of the layers."""
        top_l1 = 0
        top_l2 = top_l1 - self.l1_t - self.prepreg_t
        top_l3 = top_l2 - self.l2_t - self.prepreg_t
        bottom = top_l3 - self.l3_t

        metal_t = self.unit_cell_0.metal_t

        array_sx = self.array_sx
        array_sy = self.array_sy

        sub1 = td.Structure(
            geometry=td.Box(
                center=(0, 0, top_l1 - self.l1_t / 2),
                size=(array_sx, array_sy, self.l1_t),
            ),
            medium=self.substrate_medium,
        )
        prepreg1 = td.Structure(
            geometry=td.Box(
                center=(0, 0, top_l1 - self.l1_t - self.prepreg_t / 2),
                size=(array_sx, array_sy, self.prepreg_t),
            ),
            medium=self.prepreg_medium,
        )
        sub2 = td.Structure(
            geometry=td.Box(
                center=(0, 0, top_l2 - self.l2_t / 2),
                size=(array_sx, array_sy, self.l2_t),
            ),
            medium=self.substrate_medium,
        )
        prepreg2 = td.Structure(
            geometry=td.Box(
                center=(0, 0, top_l2 - self.l2_t - self.prepreg_t / 2),
                size=(array_sx, array_sy, self.prepreg_t),
            ),
            medium=self.prepreg_medium,
        )
        sub3 = td.Structure(
            geometry=td.Box(
                center=(0, 0, top_l3 - self.l3_t / 2),
                size=(array_sx, array_sy, self.l3_t),
            ),
            medium=self.substrate_medium,
        )

        substrates = [sub1, prepreg1, sub2, prepreg2, sub3]

        patch_box = td.Box(
            center=(0, 0, top_l1 + metal_t / 2), size=(array_sx, array_sy, metal_t)
        )
        ground_box = td.Box(
            center=(0, 0, top_l3 + metal_t / 2),
            size=(array_sx, array_sy, metal_t),
        )
        bottom_box = td.Box(
            center=(0, 0, bottom - metal_t / 2),
            size=(array_sx, array_sy, metal_t),
        )

        dl_metals = (self.dl_xy, self.dl_xy, metal_t)
        dl_subs = (self.dl_xy, self.dl_xy, self.dl_z)

        substrate_overrides = [
            td.MeshOverrideStructure(geometry=patch_box, dl=dl_metals),
            td.MeshOverrideStructure(geometry=ground_box, dl=dl_metals),
            td.MeshOverrideStructure(geometry=bottom_box, dl=dl_metals),
            td.MeshOverrideStructure(geometry=prepreg1.geometry, dl=dl_subs),
            td.MeshOverrideStructure(geometry=prepreg2.geometry, dl=dl_subs),
            td.MeshOverrideStructure(geometry=sub1.geometry, dl=dl_subs),
            td.MeshOverrideStructure(geometry=sub2.geometry, dl=dl_subs),
            td.MeshOverrideStructure(geometry=sub3.geometry, dl=dl_subs),
        ]
        return substrates, substrate_overrides

    def make_array(
        self, bit_mask: Union[xr.DataArray, bool], ground_width: float = None
    ) -> tuple[
        td.GeometryGroup,
        list[td.MeshOverrideStructure],
        list[tuple[float, float, float]],
    ]:
        """Returns all Tidy3D structures, mesh overrides and snapping points for the array."""
        S = self.unit_cell_0.S
        xstart = -self.Nx / 2.0 * S + 0.5 * S
        ystart = -self.Ny / 2.0 * S + 0.5 * S

        substrate_structures, substrate_overrides = self.make_substrates()

        geo_unit_cell_0, cell_overrides, cell_snapping_points = (
            self.unit_cell_0.make_unit_cell_geometry(
                dl_hole=self.dl_hole, ground_width=ground_width
            )
        )
        geo_unit_cell_1, _, _ = self.unit_cell_1.make_unit_cell_geometry(
            dl_hole=self.dl_hole, ground_width=ground_width
        )

        def get_cell_variant(i, j):
            if type(bit_mask) is bool:
                return bit_mask
            return bit_mask[i, j].item()

        unit_cell_geos = []
        mesh_overrides = []
        arrow_snapping_points = []
        for i in range(0, self.Nx):
            dx = xstart + i * S
            for j in range(0, self.Ny):
                dy = ystart + j * S
                # Add translated geometries
                translate = td.Transformed.translation(dx, dy, 0)
                if get_cell_variant(i, j):
                    transformed_cell = td.Transformed(
                        geometry=geo_unit_cell_1, transform=translate
                    )
                else:
                    transformed_cell = td.Transformed(
                        geometry=geo_unit_cell_0, transform=translate
                    )
                unit_cell_geos.append(transformed_cell)

                # Add mesh overrides
                mesh_overrides += [
                    override_hole.updated_copy(
                        path="geometry",
                        center=(
                            dx + override_hole.geometry.center[0],
                            dy + override_hole.geometry.center[1],
                            override_hole.geometry.center[2],
                        ),
                    )
                    for override_hole in cell_overrides
                ]

                # Add snapping points
                arrow_snapping_points += [
                    (v[0] + dx, v[1] + dy, v[2]) for v in cell_snapping_points
                ]

        all_cells = td.GeometryGroup(geometries=unit_cell_geos)
        all_cells_structure = td.Structure(geometry=all_cells, medium=self.metal_medium)

        all_structures = substrate_structures + [all_cells_structure]
        all_mesh_overrides = substrate_overrides + mesh_overrides

        return all_structures, all_mesh_overrides, arrow_snapping_points

def compute_flux_polarization_components(mon_data) -> tuple[float,float]:
    """ Compute flux in the +z axis direction, but decomposed
    into x and y electric field polarizations"""
    Ex = mon_data.Ex
    Ey = mon_data.Ey
    Hx = mon_data.Hx
    Hy = mon_data.Hy

    Px = 0.5 * np.real(Ex * Hy.conj())
    Py = 0.5 * np.real(-Ey * Hx.conj())

    flux_x = Px.integrate(["x", "y"])
    flux_y = Py.integrate(["x", "y"])
    return flux_x, flux_y

def find_frequency_bounds(value_dB: xr.DataArray, threshold = -1) -> tuple[float,float]:
    """Find the frequency bounds where value_dB is above the threshold."""
    value_dB = value_dB.squeeze()
    peak_amp = value_dB.max().item()
    peak_frequency = value_dB.idxmax(dim="f").item()
    if peak_amp < threshold:
        return np.nan, np.nan

    below_threshold = (value_dB < threshold).values
    low_freq = value_dB.f[below_threshold].where(value_dB.f < peak_frequency).max().item()
    high_freq = value_dB.f[below_threshold].where(value_dB.f > peak_frequency).min().item()

    return low_freq, high_freq

def calculate_fractional_bandwith_metric(mon_data: td.FieldData, threshold=-1) -> float:
    """Calculate the fractional bandwidth metric given a simulation result"""
    flux_x, _ = compute_flux_polarization_components(mon_data)
    # The co-polarization reflection coefficient in dB [1]. The source is normalized to 1 Watt.
    Rxy_dB = 10 * np.log10(flux_x).squeeze()

    peak = Rxy_dB.max().item()
    if peak < threshold:
        return 0.0

    # Find the max value, and from there find where Rxy crosses the threshold of -1 dB.
    low_freq, high_freq = find_frequency_bounds(Rxy_dB, threshold=threshold)
    freq0 = (low_freq + high_freq) / 2
    bandwidth = high_freq - low_freq
    # Fractional bandwidth in percent
    fractional_bandwidth = 100 * bandwidth / freq0
    return fractional_bandwidth

def make_structures_and_gridspec(l_head : float, l_stem : float) -> tuple [td.Structure, td.GridSpec]:
    """Recreates the structures that are affected by a new arrow head length and stem length."""
    # Update unit cell parameters
    array.unit_cell_0.l_head = l_head
    array.unit_cell_1.l_head = l_head
    array.unit_cell_0.l_stem = l_stem
    array.unit_cell_1.l_stem = l_stem

    all_structures, all_mesh_overrides, arrow_snapping_points = array.make_array(
        bit_mask=bool_array, ground_width=substrate_width
    )

    grid_spec = td.GridSpec.auto(
        min_steps_per_wvl=10,
        wavelength=td.C_0 / freq_stop,
        override_structures=all_mesh_overrides,
        snapping_points=arrow_snapping_points,
        dl_min=0.018 * mm,
    )
    return all_structures, grid_spec


def pre(l_head, l_stem) -> td.Simulation:
    """Pre-processing function, which creates a ``Simulation`` from the design parameters."""
    all_structures, grid_spec = make_structures_and_gridspec(l_head, l_stem)
    return sim.updated_copy(structures=all_structures, grid_spec=grid_spec)


def post(data: td.SimulationData) -> float:
    """Post-processing function, which processes the tidy3d simulation data to return the function output."""
    mon_data = data[R_mon.name]
    frac_bw = calculate_fractional_bandwith_metric(mon_data, threshold=-1)
    return frac_bw


# Setup of the parameters for the optimization step. The span indicates the bounds of the associated parameter.
param_l1 = tdd.ParameterFloat(name="l_head", span=(2.2 * mm, 2.83 * mm), num_points=1)
param_l2 = tdd.ParameterFloat(name="l_stem", span=(2.2 * mm, 3.4 * mm), num_points=2)

method = tdd.MethodBayOpt(
    initial_iter=50,
    n_iter=10,
    acq_func="ucb",
    kappa=0.6,
    seed=1,
)

design_space = tdd.DesignSpace(parameters=[param_l1, param_l2], method=method)