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

from ..structures import *
from ..routing import *

from ..primitives.pic import *

## Class for splitting tree 
class SplittingTree():
    '''
    Class for generating splitting tree. 

    Paras:  
        1. ybranch          [class] (Default None)
            - length    (Default: 28)   Length of the ybranch 
            - width     (Default: 2)    Pitch between two output waveguides 
            - w_wg      (Default: 0.45) Width of output waveguide 
            - cell      (Default: box)  
                - a1    [Pin] Input waveguide 
                - b1    [Pin] Output waveguide1 
                - b2    [Pin] Output waveguide2 
        2. output_number    [-]     (Default: 16)
            Number of output channels(Need to be 2^N)
        3. bend_radius      [um]    (Default: 10)
            Bend radius used to connect the different layer of Y branch 
        4. output_pitch     [um]    (Default: None)
            Can define the pitch of output channel(If ==None, then pitch=self.width, which is the minimum pitch)
    '''

    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