"""Optical splitting tree composite layouts."""

from typing import Any
import nazca as nd
import numpy as np
import nazca.interconnects as IC

from ..geometry import *
from ..routing import *

from ..primitives.pic import *

## Class for splitting tree 
class SplittingTree():
    """Binary splitter tree built from repeated Y-branch cells.

    Parameters
    ----------
    ybranch : Any, optional
        Y-branch object with a ``cell`` attribute and ``a1``, ``b1``, and
        ``b2`` pins. If omitted, a simple box-based Y-branch is generated.
    output_number : int, optional
        Number of output channels. Values are coerced to the nearest lower
        power of two when needed.
    bend_radius : int, optional
        Bend radius used to route between splitter levels.
    output_pitch : Any, optional
        Output channel pitch in microns. If omitted, the Y-branch output
        pitch is used.
    show_pins : bool, optional
        Show Nazca pin stubs in the generated layout.

    Attributes
    ----------
    cell : nazca.Cell
        Generated splitting tree layout cell.
    """

    def __init__(self, ybranch: Any=None, output_number: int=16, bend_radius: int=10, output_pitch: Any=None, show_pins: bool=False) -> None:
        if ybranch == None:
            self._generate_default_ybranch_()
        else:
            self.yb_cell = ybranch.cell

            
            # self.yb_length = ybranch.length
            self.yb_length = np.abs(self.yb_cell.pin['a1'].x - self.yb_cell.pin['b1'].x)
            # self.yb_width = ybranch.width
            self.yb_width = np.abs(self.yb_cell.pin['b1'].y - self.yb_cell.pin['b2'].y)
            
            self.w_wg = ybranch.w_wg 
        self.output_number = output_number 
        self.bend_radius = bend_radius 
        if output_pitch == None:
            self.output_pitch = self.yb_width 
        else:
            self.output_pitch = output_pitch
        self.show_pins = show_pins
        
        self.cell = self.generate_gds()

    def generate_gds(self):
        '''
        Generate the gds of splitting tree. 
        '''
        ## initialize the parameters of splitting tree 
        self.level_number = np.log2(self.output_number)
        if self.level_number != int(np.log2(self.output_number)):
            print("WARNNING:: Please check the output number of your generated splitting tree, which is not 2^N. ")
            self.level_number = int(self.level_number)
            self.output_number = np.power(2, self.level_number)
        ## Generate the splitting tree 
        with nd.Cell(name="SplittingTree_N"+str(self.output_number), instantiate=False) as C:
            stripe_class = Route(xs="strip", width=self.w_wg, radius=self.bend_radius)
            cell_dic = {}
            cell_dic['00'] = self.yb_cell.put('a0', 0, 0)
            x_cur = 0 
            for level_index in range(1, int(self.level_number)):
                # Calculate the x location for current level 
                y_pitch = self.output_pitch * np.power(2, self.level_number-level_index)
                if self.bend_radius > (y_pitch/2-self.yb_width/2)/2:
                    x_cur = x_cur + self.yb_length + 2*np.sqrt(np.power(self.bend_radius,2)-np.power(self.bend_radius-(y_pitch/2)/2, 2)) + 2
                else:
                    x_cur = x_cur + self.yb_length + self.bend_radius*2 + 2
                for yb_index in range(np.power(2, level_index)):
                    y_cur = y_pitch * ((np.power(2, level_index)-1)/2 - yb_index) 
                    # Put the Y-branch 
                    cell_dic[str(level_index)+str(yb_index)] = self.yb_cell.put('a1', x_cur, y_cur)    
                    # Do the routing 
                    stripe_class.sbend_p2p(
                        pin1=cell_dic[str(level_index-1)+str(yb_index//2)].pin['b'+str(yb_index%2+1)], 
                        pin2=cell_dic[str(level_index)+str(yb_index)].pin['a1'], 
                        Lstart=1, 
                        arrow=False
                    ).put()       
            ## Put pins 
            nd.Pin(name="a0").put(0, 0, 180)
            nd.Pin(name="a1",width=self.w_wg).put(0, 0, 180)        
            level_index = int(self.level_number-1) 
            for yb_index in range(np.power(2, level_index)):
                nd.Pin(
                    name="b"+str(2*yb_index+1),
                    width=self.w_wg
                ).put(cell_dic[str(level_index)+str(yb_index)].pin['b1'])
                nd.Pin(
                    name="b"+str(2*yb_index+2),
                    width=self.w_wg
                ).put(cell_dic[str(level_index)+str(yb_index)].pin['b2'])
            nd.Pin(name="b0").put(
                cell_dic[str(level_index)+str(yb_index)].pin['b2'].x,
                0,
                0
            )
            if self.show_pins:
                nd.put_stub()
        return C

    def _generate_default_ybranch_(self, length=28, width=2, w_wg=0.45):
        '''
        Generate a ybranch built by box for quick showing. 
        '''
        self.yb_length = length 
        self.yb_width = width 
        self.w_wg = w_wg 
        with nd.Cell(name="Ybranch", instantiate=False) as C: 
            ## Put block to show the size of the device and the location of input&output waveguide 
            nd.strt(length=self.yb_length, width=self.yb_width+2, layer=(1001, 1)).put(0, 0)
            nd.strt(length=0.2, width=self.w_wg, layer=(1001, 2)).put(0, 0)
            nd.strt(length=0.2, width=self.w_wg, layer=(1001, 2)).put(self.yb_length, self.yb_width/2, 180)
            nd.strt(length=0.2, width=self.w_wg, layer=(1001, 2)).put(self.yb_length, -self.yb_width/2, 180)
            ## Put pins 
            nd.Pin(name="a0", width=self.w_wg).put(0, 0, 180)
            nd.Pin(name="a1", width=self.w_wg).put(0, 0, 180)
            nd.Pin(name="b1", width=self.w_wg).put(self.yb_length, self.yb_width/2, 0)
            nd.Pin(name="b2", width=self.w_wg).put(self.yb_length, -self.yb_width/2, 0)
        self.yb_cell = C