import numpy as np
from tqdm.notebook import tqdm

import meow as mw

def create_structures(length=10.0):
    oxide = mw.Structure(
        material=mw.silicon_oxide,
        geometry=mw.Prism(
            poly=np.array([(0, -2.1), (length, -2.1), (length, 2.1), (0, 2.1)]),
            h_min=-3,
            h_max=0,
            axis="y",
        ),
    )

    poly = np.array([(0, -0.45 / 2), (length, -1), (length, 1), (0, 0.45 / 2)])

    core = mw.Structure(
        material=mw.silicon,
        geometry=mw.Prism(
            poly=poly,
            h_min=0,
            h_max=0.22,
            axis="y",
        ),
    )

    structures = [oxide, core]
    return structures


structures = create_structures()

mw.visualize(structures)

num_cells = 20

cells = mw.create_cells(
    structures=structures,
    mesh=mw.Mesh2D(
        x=np.linspace(-2, 2, 101),
        y=np.linspace(-2, 2, 101),
    ),
    Ls=np.array([10.0 / num_cells] * num_cells),
)
mw.visualize(cells[-1])

env = mw.Environment(wl=1.55, T=25.0)
css = [mw.CrossSection.from_cell(cell=cell, env=env) for cell in cells]
mw.visualize(css[0])

num_modes = 20

modes = [mw.compute_modes(cs, num_modes=num_modes) for cs in tqdm(css)]

taper_lengths = [5.0, 10.0, 15.0, 20.0]

for length in taper_lengths:
    cell_lengths = [length / num_cells] * num_cells
    S, pm = mw.compute_s_matrix(modes, cell_lengths=cell_lengths)
    T_00 = abs(S[0, num_modes]) ** 2
    print(f"taper length {length:5.1f} | transmission: {T_00:.6f}")