import numpy as np

from gdshelpers.geometry.chip import Cell
from gdshelpers.parts.mode_converter import StripToSlotModeConverter
from gdshelpers.parts.waveguide import Waveguide
from gdshelpers.parts.port import Port

# waveguide 1: strip waveguide
wg_1 = Waveguide.make_at_port(Port(origin=(0, 0), angle=np.pi / 2, width=1))
wg_1.add_straight_segment(length=10)

# waveguide 2: slot waveguide
wg_2 = Waveguide.make_at_port(Port(origin=(5, 0), angle=np.pi / 2, width=[0.4, 0.2, 0.4]))
wg_2.add_straight_segment(length=10)

# waveguide 3: slot waveguide with tapering
wg_3 = Waveguide.make_at_port(Port(origin=(10, 0), angle=np.pi / 2, width=[0.5, 0.3, 0.5]))
wg_3.add_straight_segment(length=10, final_width=[0.2, 0.4, 0.2])

# waveguide 4: slot waveguide with three rails and two slots
wg_4 = Waveguide.make_at_port(Port(origin=(15, 0), angle=np.pi / 2, width=[0.2, 0.2, 0.3, 0.2, 0.4]))
wg_4.add_straight_segment(length=10)

# waveguide 5: slot waveguide with bends and strip to slot mode converter
wg_5_1 = Waveguide.make_at_port(Port(origin=(-6.5, 10), angle=-np.pi / 2, width=[0.4, 0.2, 0.4]))
wg_5_1.add_straight_segment(length=10)
wg_5_1.add_bend(angle=np.pi / 2, radius=5)
mc_1 = StripToSlotModeConverter.make_at_port(port=wg_5_1.current_port, taper_length=5, final_width=1,
                                             pre_taper_length=2, pre_taper_width=0.2)
wg_5_2 = Waveguide.make_at_port(port=mc_1.out_port)
wg_5_2.add_straight_segment(5)
mc_2 = StripToSlotModeConverter.make_at_port(port=wg_5_2.current_port, taper_length=5, final_width=[0.4, 0.2, 0.4],
                                             pre_taper_length=2, pre_taper_width=0.2)
wg_5_3 = Waveguide.make_at_port(port=mc_2.out_port)
wg_5_3.add_bend(angle=np.pi / 2, radius=5)
wg_5_3.add_straight_segment(length=10)

cell = Cell('Cell')
cell.add_to_layer(1, wg_1)  # red
cell.add_to_layer(2, wg_2)  # green
cell.add_to_layer(3, wg_3)  # blue
cell.add_to_layer(4, wg_4)  # jungle green
cell.add_to_layer(5, wg_5_1, mc_1, wg_5_2, mc_2, wg_5_3) # pink
cell.show()