""" Layout helper functions.

Authors: Thomas Ferreira de Lima @thomaslima
         Lukas Chrostowski @lukasc-ubc

The following functions are useful for scripted layout, or making
PDK Pcells.

Functions:

layout_waveguide2
layout_waveguide
layout_waveguide_sbend_bezier
make_pin
y_splitter_tree

TODO: enhance documentation
TODO: make some of the functions in util use these.
"""

from itertools import repeat
import pya
import numpy as np
from numpy import cos, sin, pi, sqrt
from functools import reduce
from .sampling import sample_function
from .geometry import rotate90, rotate, bezier_optimal, curve_length

'''
Create a waveguide, in a specific technology
inputs
- cell: into which Cell we add the waveguide
- dpath: DPath type
- waveguide_type: a name from Waveguides.XML
    can be a <compound_waveguide>
    or a primitive waveguide type containing <component> info
output
- compound waveguide, or regular waveguide
by Lukas Chrostowski
acknowledgements: Diedrik Vermeulen for the code to place the taper in the correct orientation
'''
def layout_waveguide4(cell, dpath, waveguide_type, debug=True):

    if debug:
        print ('SiEPIC.utils.layout.layout_waveguide4: ' )
        print (' - waveguide_type: %s' % (waveguide_type) )

    # get the path and clean it up
    layout = cell.layout()
    dbu = layout.dbu
    dpath = dpath.to_itype(dbu)
    dpath.unique_points()
    pts = dpath.get_points()
    dpts = dpath.get_dpoints()

    # Load the technology and all waveguide types
    from SiEPIC.utils import load_Waveguides_by_Tech
    technology_name = layout.technology_name
    waveguide_types = load_Waveguides_by_Tech(technology_name)
    if debug:
        print (' - technology_name: %s' % (technology_name) )
        print (' - waveguide_types: %s' % (waveguide_types) )

    # Load parameters for the chosen waveguide type
    params = [t for t in waveguide_types if t['name'] == waveguide_type]
    if type(params) == type([]) and len(params)>0:
        params = params[0]
    else:
        print('error: waveguide type not found in PDK waveguides')
        raise Exception('error: waveguide type (%s) not found in PDK waveguides' % waveguide_type)

    # compound waveguide types:
    if 'compound_waveguide' in params:
        # find the singlemode and multimode waveguides:
        if 'singlemode' in params['compound_waveguide']:
            singlemode = params['compound_waveguide']['singlemode']
        else:
            raise Exception('error: waveguide type (%s) does not have singlemode defined' % waveguide_type)
        if 'multimode' in params['compound_waveguide']:
            multimode = params['compound_waveguide']['multimode']
        else:
            raise Exception('error: waveguide type (%s) does not have multimode defined' % waveguide_type)
        params_singlemode = [t for t in waveguide_types if t['name'] == singlemode]
        params_multimode = [t for t in waveguide_types if t['name'] == multimode]
        if type(params_singlemode) == type([]) and len(params_singlemode)>0:
            params_singlemode = params_singlemode[0]
        else:
            raise Exception('error: waveguide type (%s) not found in PDK waveguides' % singlemode)
        if type(params_multimode) == type([]) and len(params_multimode)>0:
            params_multimode = params_multimode[0]
        else:
            raise Exception('error: waveguide type (%s) not found in PDK waveguides' % multimode)
        # find the taper
        if 'taper_library' in params['compound_waveguide'] and 'taper_cell' in params['compound_waveguide']:
            taper = layout.create_cell(params['compound_waveguide']['taper_cell'], params['compound_waveguide']['taper_library'])
            if not taper:
                raise Exception ('Cannot import cell %s : %s' % (params['compound_waveguide']['taper_cell'], params['compound_waveguide']['taper_library']))
        else:
            raise Exception('error: waveguide type (%s) does not have taper cell and library defined' % waveguide_type)
        from pya import Trans, CellInstArray

        '''
        find sections of waveguides that are larger than (2 x radius + 2 x taper_length)
         - insert two tapers
         - insert multimode straight section
         - insert singlemode waveguides (including bends) before
        '''
        import math
        from SiEPIC.extend import to_itype
        from pya import Point
        radius = to_itype(params_singlemode['radius'],dbu)
        taper_length = taper.find_pins()[0].center.distance(taper.find_pins()[1].center)
        min_length = 2*radius + 2*taper_length
        offset = radius
        wg_sm_segment_pts = []
        wg_last=0
        waveguide_length = 0
        for ii in range(1,len(dpts)):
            start_point = dpts[ii-1]
            end_point = dpts[ii]
            distance_points = end_point.distance(start_point)
            if distance_points < min_length:
                # single mode segment, keep track
                if ii==1:
                    wg_sm_segment_pts.append(pts[ii-1])
                wg_sm_segment_pts.append(pts[ii])
                if ii==len(pts)-1:
                    subcell = layout.create_cell("Waveguide_sm_%s" %ii)
                    cell.insert(CellInstArray(subcell.cell_index(), Trans()))
                    waveguide_length += layout_waveguide3(subcell, wg_sm_segment_pts, params_singlemode, debug=True)
            else:
                # insert two tapers and multimode waveguide
                angle = math.atan2((end_point.y-start_point.y),(end_point.x-start_point.x))/math.pi*180
                if ii==1:
                    wg_first=offset
                else:
                    wg_first=0
                if ii==len(pts)-1:
                    wg_last=offset
                if round(angle)%360 == 270.0:
                    t = Trans(Trans.R270, start_point.x, start_point.y-offset+wg_first)
                    t2 = Trans(Trans.R90, end_point.x, end_point.y+offset-wg_last)
                    wg_start_pt = Point(start_point.x, start_point.y-offset-taper_length+wg_first)
                    wg_end_pt = Point(end_point.x, end_point.y+offset+taper_length-wg_last)
                if round(angle)%360 == 90.0:
                    t = Trans(Trans.R90, start_point.x, start_point.y+offset-wg_first)
                    t2 = Trans(Trans.R270, end_point.x, end_point.y-offset+wg_last)
                    wg_start_pt = Point(start_point.x, start_point.y+offset+taper_length-wg_first)
                    wg_end_pt = Point(end_point.x, end_point.y-offset-taper_length+wg_last)
                if round(angle)%360 == 180.0:
                    t = Trans(Trans.R180, start_point.x-offset+wg_first, start_point.y)
                    t2 = Trans(Trans.R0, end_point.x+offset-wg_last, end_point.y)
                    wg_start_pt = Point(start_point.x-offset-taper_length+wg_first, start_point.y)
                    wg_end_pt = Point(end_point.x+offset+taper_length-wg_last, end_point.y)
                if round(angle)%360 == 0.0:
                    t = Trans(Trans.R0, start_point.x+offset-wg_first, start_point.y)
                    t2 = Trans(Trans.R180, end_point.x-offset+wg_last, end_point.y)
                    wg_start_pt = Point(start_point.x+offset+taper_length-wg_first, start_point.y)
                    wg_end_pt = Point(end_point.x-offset-taper_length+wg_last, end_point.y)
                inst_taper = cell.insert(CellInstArray(taper.cell_index(), t))
                inst_taper = cell.insert(CellInstArray(taper.cell_index(), t2))
                waveguide_length += taper_length*2
                subcell = layout.create_cell("Waveguide_mm_%s" %ii)
                cell.insert(CellInstArray(subcell.cell_index(), Trans()))
                waveguide_length += layout_waveguide3(subcell, [wg_start_pt, wg_end_pt], params_multimode, debug=True)
                # compound segment
                if ii>1:
                    wg_sm_segment_pts.append(t.disp.to_p())
                    subcell = layout.create_cell("Waveguide_sm_%s" %ii)
                    cell.insert(CellInstArray(subcell.cell_index(), Trans()))
                    waveguide_length += layout_waveguide3(subcell, wg_sm_segment_pts, params_singlemode, debug=True)
                    wg_sm_segment_pts = [t2.disp.to_p(), pts[ii]]
                else:
                    wg_sm_segment_pts = [t2.disp.to_p(), pts[ii]]

    else:
        # primitive waveguide type
        waveguide_length = layout_waveguide3(cell, pts, params, debug=True)

    return waveguide_length

'''
Create a waveguide, in a specific technology
inputs
- cell: into which Cell we add the waveguide
    from SiEPIC.utils import get_layout_variables
    TECHNOLOGY, lv, layout, cell = get_layout_variables()
- pts
- params, obtained from load_Waveguides_by_Tech and Waveguides.XML
    must be a primitive waveguide type containing <component> info
output:
- waveguide
- DevRec, PinRec
by Lukas Chrostowski
'''
def layout_waveguide3(cell, pts, params, debug=True):

    if debug:
        print ('SiEPIC.utils.layout.layout_waveguide3: ' )

    layout = cell.layout()
    dbu = layout.dbu
    technology_name = layout.technology_name
    from SiEPIC.utils import get_technology_by_name
    TECHNOLOGY = get_technology_by_name(technology_name)

    from SiEPIC.extend import to_itype
    wg_width = to_itype(params['width'],dbu)
    radius = float(params['radius'])
    model = params['model']
    cellName = 'Waveguide2'
    CML = params['CML']

    if debug:
        print (' - waveguide params: %s' % (params) )

    if 'compound_waveguide' in params:
        print('error: this function cannot handle compound waveguides')
        raise Exception('error: this function cannot handle compound waveguides (%s)' % waveguide_type)

    # draw the waveguide
    waveguide_length = layout_waveguide2(TECHNOLOGY, layout, cell, [wg['layer'] for wg in params['component']], [wg['width'] for wg in params['component']], [wg['offset'] for wg in params['component']], pts, radius, params['adiabatic'], params['bezier'])

    # Draw the marking layers
    from SiEPIC.utils import angle_vector
    LayerPinRecN = layout.layer(TECHNOLOGY['PinRec'])

    make_pin(cell, 'opt1', pts[0], wg_width, LayerPinRecN, angle_vector(pts[0]-pts[1])%360)
    make_pin(cell, 'opt2', pts[-1], wg_width, LayerPinRecN, angle_vector(pts[-1]-pts[-2])%360)

    from pya import Trans, Text, Path, Point

    '''
    t1 = Trans(angle_vector(pts[0]-pts[1])/90, False, pts[0])
    cell.shapes(LayerPinRecN).insert(Path([Point(-10, 0), Point(10, 0)], wg_width).transformed(t1))
    cell.shapes(LayerPinRecN).insert(Text("opt1", t1, 0.3/dbu, -1))

    t = Trans(angle_vector(pts[-1]-pts[-2])/90, False, pts[-1])
    cell.shapes(LayerPinRecN).insert(Path([Point(-10, 0), Point(10, 0)], wg_width).transformed(t))
    cell.shapes(LayerPinRecN).insert(Text("opt2", t, 0.3/dbu, -1))
    '''

    LayerDevRecN = layout.layer(TECHNOLOGY['DevRec'])

    # Compact model information
    angle_vec = angle_vector(pts[0]-pts[1])/90
    halign = 0 # left
    angle=0
    dpt = Point(0,0)
    if angle_vec == 0: # horizontal
      halign = 2 # right
      angle=0
      dpt = Point(0, 0.2*wg_width)
    if angle_vec == 2: # horizontal
      halign = 0 # left
      angle = 0
      dpt=Point(0, 0.2*wg_width)
    if angle_vec == 1: # vertical
      halign = 2 # right
      angle = 1
      dpt=Point(0.2*wg_width,0)
    if angle_vec == -1: # vertical
      halign = 0 # left
      angle = 1
      dpt=Point(0.2*wg_width,0)
    pt2=pts[0] + dpt
    pt3=pts[0] - dpt
    pt4=pts[0] - 6*dpt
    pt5=pts[0] + 2*dpt

    t = Trans(angle, False, pt3)
    text = Text ('Lumerical_INTERCONNECT_library=Design kits/%s' % CML, t, 0.1*wg_width, -1)
    text.halign=halign
    shape = cell.shapes(LayerDevRecN).insert(text)
    t = Trans(angle, False, pt2)
    text = Text ('Component=%s' % model, t, 0.1*wg_width, -1)
    text.halign=halign
    shape = cell.shapes(LayerDevRecN).insert(text)
    t = Trans(angle, False, pt5)
    text = Text ('cellName=%s' % cellName, t, 0.1*wg_width, -1)
    text.halign=halign
    shape = cell.shapes(LayerDevRecN).insert(text)
    t = Trans(angle, False, pts[0])
    pts_txt = str([ [round(p.to_dtype(dbu).x,3), round(p.to_dtype(dbu).y,3)] for p in pts ]).replace(', ',',')
    text = Text ( \
      'Spice_param:wg_length=%.9f wg_width=%.3g points="%s" radius=%.3g' %\
        (waveguide_length*1e-6, wg_width*1e-9, pts_txt,radius*1e-6 ), t, 0.1*wg_width, -1  )
    text.halign=halign
    shape = cell.shapes(LayerDevRecN).insert(text)
    t = Trans(angle, False, pt4)
    text = Text ( \
      'Length=%.3f (microns)' %(waveguide_length), t, 0.5*wg_width, -1  )
    text.halign=halign
    shape = cell.shapes(LayerDevRecN).insert(text)

    return waveguide_length

'''
Create a waveguide, in a specific technology
inputs
- TECHNOLOGY, layout, cell:
    from SiEPIC.utils import get_layout_variables
    TECHNOLOGY, lv, layout, cell = get_layout_variables()
- layers: list of text names, e.g., ['Waveguide']
- widths: list of floats in units Microns, e.g., [0.50]
- offsets: list of floats in units Microns, e.g., [0]
- pts: a list of pya.Points, e.g.
    L=15/dbu
    pts = [Point(0,0), Point(L,0), Point(L,L)]
- radius: in Microns, e.g., 5
- adiab: 1 = Bezier curve, 0 = radial bend (arc)
- bezier: the bezier parameter, between 0 and 0.45 (almost a radial bend)

Note: bezier parameters need to be simulated and optimized, and will depend on
    wavelength, polarization, width, etc.  TM and rib waveguides don't benefit from bezier curves
    most useful for TE
by Lukas Chrostowski
'''
def layout_waveguide2(TECHNOLOGY, layout, cell, layers, widths, offsets, pts, radius, adiab, bezier):
    from SiEPIC.utils import arc_xy, arc_bezier, angle_vector, angle_b_vectors, inner_angle_b_vectors, translate_from_normal
    from SiEPIC.extend import to_itype
    from pya import Path, Polygon, Trans
    dbu = layout.dbu

    if 'Errors' in TECHNOLOGY:
        error_layer = layout.layer(TECHNOLOGY['Errors'])
    else:
        error_layer = None

    width=widths[0]
    turn=0
    waveguide_length = 0
    for lr in range(0, len(layers)):
        wg_pts = [pts[0]]
        layer = layout.layer(TECHNOLOGY[layers[lr]])
        width = to_itype(widths[lr],dbu)
        offset = to_itype(offsets[lr],dbu)
        for i in range(1,len(pts)-1):
            turn = ((angle_b_vectors(pts[i]-pts[i-1],pts[i+1]-pts[i])+90)%360-90)/90
            dis1 = pts[i].distance(pts[i-1])
            dis2 = pts[i].distance(pts[i+1])
            angle = angle_vector(pts[i]-pts[i-1])/90
            pt_radius = to_itype(radius,dbu)
            error_seg1 = False
            error_seg2 = False
            # determine the radius, based on how much space is available
            if len(pts)==3:
                # simple corner, limit radius by the two edges
                if dis1 < pt_radius:
                    error_seg1 = True
                if dis2 < pt_radius:
                    error_seg2 = True
                pt_radius = min (dis1, dis2, pt_radius)
            else:
                if i==1:
                    # first corner, limit radius by first edge, or 1/2 of second one
                    if dis1 < pt_radius:
                        error_seg1 = True
                    if dis2/2 < pt_radius:
                        error_seg2 = True
                    pt_radius = min (dis1, dis2/2, pt_radius)
                elif i==len(pts)-2:
                    # last corner, limit radius by second edge, or 1/2 of first one
                    if dis1/2 < pt_radius:
                        error_seg1 = True
                    if dis2 < pt_radius:
                        error_seg2 = True
                    pt_radius = min (dis1/2, dis2, pt_radius)
                else:
                    if dis1/2 < pt_radius:
                        error_seg1 = True
                    if dis2/2 < pt_radius:
                        error_seg2 = True
                    pt_radius = min (dis1/2, dis2/2, pt_radius)

            if error_seg1 or error_seg2:
                if not error_layer:
                    # we have an error, but no Error layer
                    print('- SiEPIC:layout_waveguide2: missing Error layer')
                elif layer == layout.layer(TECHNOLOGY['Waveguide']): # and pt_radius < to_itype(radius,dbu):
                    # add an error polygon to flag the incorrect bend
                    if error_seg1:
                        error_pts = pya.Path([pts[i-1], pts[i]], width)
                        cell.shapes(error_layer).insert(error_pts)
                    if error_seg2:
                        error_pts = pya.Path([pts[i], pts[i+1]], width)
                        cell.shapes(error_layer).insert(error_pts)
    #                error_pts = pya.Path([pts[i-1], pts[i], pts[i+1]], width)
    #                cell.shapes(error_layer).insert(error_pts)
            # waveguide bends:
            if abs(turn)==1:
                if(adiab):
                    wg_pts += Path(arc_bezier(pt_radius, 270, 270 + inner_angle_b_vectors(pts[i-1]-pts[i], pts[i+1]-pts[i]), float(bezier), DevRec='DevRec' in layers[lr]), 0).transformed(Trans(angle, turn < 0, pts[i])).get_points()
                else:
                    wg_pts += Path(arc_xy(-pt_radius, pt_radius, pt_radius, 270, 270 + inner_angle_b_vectors(pts[i-1]-pts[i], pts[i+1]-pts[i]),DevRec='DevRec' in layers[lr]), 0).transformed(Trans(angle, turn < 0, pts[i])).get_points()

        wg_pts += [pts[-1]]
        wg_pts = pya.Path(wg_pts, 0).unique_points().get_points()
        wg_polygon = Polygon(translate_from_normal(wg_pts, width/2 + (offset if turn > 0 else - offset))+translate_from_normal(wg_pts, -width/2 + (offset if turn > 0 else - offset))[::-1])
        cell.shapes(layer).insert(wg_polygon)

        if layout.layer(TECHNOLOGY['Waveguide']) == layer:
            waveguide_length = wg_polygon.area() / width * dbu

    return waveguide_length

def layout_waveguide(cell, layer, points_list, width):
    """ Lays out a waveguide (or trace) with a certain width with along given points.

    This is very useful for laying out Bezier curves with or without adiabatic tapers.

    Args:
        cell: cell to place into
        layer: layer to place into. It is done with cell.shapes(layer).insert(pya.Polygon)
        points_list: list of pya.DPoint (at least 2 points)
        width (microns): constant or list. If list, then it has to have the same length as points

    """
    if len(points_list) < 2:
        raise NotImplemented("ERROR: points_list too short")
        return

    if type(width)==type(0.0):
        width_iterator = repeat(width)
        points_iterator = iter(points_list)
    else:
        try:
            if len(width) == len(points_list):
                width_iterator = iter(width)
            else:
                width_iterator = repeat(width[0])
        except TypeError:
            width_iterator = repeat(width)
        finally:
            points_iterator = iter(points_list)

    dbu = cell.layout().dbu

    points_low = list()
    points_high = list()

    def norm(self):
        return sqrt(self.x**2 + self.y**2)

    def cos_angle(point1, point2):
        return point1 * point2 / norm(point1) / norm(point2)

    point_width_list = list(zip(points_iterator, width_iterator))
    N = len(point_width_list)

    first_point, first_width = point_width_list[0]
    next_point, next_width = point_width_list[1]

    delta = next_point - first_point
    theta = np.arctan2(delta.y, delta.x)
    first_high_point = first_point + 0.5 * first_width * \
        pya.DPoint(cos(theta + pi / 2), sin(theta + pi / 2))
    first_low_point = first_point + 0.5 * first_width * \
        pya.DPoint(cos(theta - pi / 2), sin(theta - pi / 2))
    points_high.append(first_high_point)
    points_low.append(first_low_point)

    for i in range(1, N - 1):
        prev_point, prev_width = point_width_list[i - 1]
        point, width = point_width_list[i]
        next_point, next_width = point_width_list[i + 1]

        delta_prev = point - prev_point
        delta_next = next_point - point
        theta_prev = np.arctan2(delta_prev.y, delta_prev.x)
        theta_next = np.arctan2(delta_next.y, delta_next.x)

        next_point_high = (next_point + 0.5 * next_width *
                           pya.DPoint(cos(theta_next + pi / 2), sin(theta_next + pi / 2)))
        next_point_low = (next_point + 0.5 * next_width *
                          pya.DPoint(cos(theta_next - pi / 2), sin(theta_next - pi / 2)))

        forward_point_high = (point + 0.5 * width *
                              pya.DPoint(cos(theta_next + pi / 2), sin(theta_next + pi / 2)))
        forward_point_low = (point + 0.5 * width *
                             pya.DPoint(cos(theta_next - pi / 2), sin(theta_next - pi / 2)))

        prev_point_high = (prev_point + 0.5 * prev_width *
                           pya.DPoint(cos(theta_prev + pi / 2), sin(theta_prev + pi / 2)))
        prev_point_low = (prev_point + 0.5 * prev_width *
                          pya.DPoint(cos(theta_prev - pi / 2), sin(theta_prev - pi / 2)))

        backward_point_high = (point + 0.5 * width *
                               pya.DPoint(cos(theta_prev + pi / 2), sin(theta_prev + pi / 2)))
        backward_point_low = (point + 0.5 * width *
                              pya.DPoint(cos(theta_prev - pi / 2), sin(theta_prev - pi / 2)))

        # High point decision
        next_high_edge = pya.DEdge(forward_point_high, next_point_high)
        prev_high_edge = pya.DEdge(backward_point_high, prev_point_high)

        if next_high_edge.crossed_by(prev_high_edge):
            intersect_point = next_high_edge.crossing_point(prev_high_edge)
            points_high.append(intersect_point)
        else:
            if width * (1 - cos_angle(delta_next, delta_prev)) > dbu:
                points_high.append(backward_point_high)
                points_high.append(forward_point_high)
            else:
                points_high.append((backward_point_high + forward_point_high) * 0.5)

        # Low point decision
        next_low_edge = pya.DEdge(forward_point_low, next_point_low)
        prev_low_edge = pya.DEdge(backward_point_low, prev_point_low)

        if next_low_edge.crossed_by(prev_low_edge):
            intersect_point = next_low_edge.crossing_point(prev_low_edge)
            points_low.append(intersect_point)
        else:
            if width * (1 - cos_angle(delta_next, delta_prev)) > dbu:
                points_low.append(backward_point_low)
                points_low.append(forward_point_low)
            else:
                points_low.append((backward_point_low + forward_point_low) * 0.5)

    last_point, last_width = point_width_list[-1]
    point, width = point_width_list[-2]
    delta = last_point - point
    theta = np.arctan2(delta.y, delta.x)
    final_high_point = last_point + 0.5 * last_width * \
        pya.DPoint(cos(theta + pi / 2), sin(theta + pi / 2))
    final_low_point = last_point + 0.5 * last_width * \
        pya.DPoint(cos(theta - pi / 2), sin(theta - pi / 2))
    if (final_high_point - points_high[-1]) * delta > 0:
        points_high.append(final_high_point)
    if (final_low_point - points_low[-1]) * delta > 0:
        points_low.append(final_low_point)

    # Append point only if change in direction is less than 120 degrees.
    def smooth_append(point_list, point):
        if point_list is None:
            print(point)
        if len(point_list) < 1:
            point_list.append(point)
            return point_list
        elif len(point_list) < 2:
            curr_edge = point - point_list[-1]
            if norm(curr_edge) > dbu:
                point_list.append(point)
                return point_list

        curr_edge = point - point_list[-1]
        if norm(curr_edge) > dbu:
            prev_edge = point_list[-1] - point_list[-2]
            if cos_angle(curr_edge, prev_edge) > cos(120 / 180 * pi):
                point_list.append(point)
        return point_list

    polygon_points = points_low + list(reversed(points_high))
    polygon_points = list(reduce(smooth_append, polygon_points, list()))

    poly = pya.DPolygon(polygon_points)
    cell.shapes(layer).insert(poly)

def layout_ring(cell, layer, center, r, w):
        # function to produce the layout of a ring
        # cell: layout cell to place the layout
        # layer: which layer to use
        # center: origin DPoint
        # r: radius
        # w: waveguide width
        # units in microns

        # example usage.  Places the ring layout in the presently selected cell.
        # cell = pya.Application.instance().main_window().current_view().active_cellview().cell
        # layout_ring(cell, cell.layout().layer(LayerInfo(1, 0)), pya.DPoint(0,0), 10, 0.5)

    layout_arc(cell, layer, center, r, w, 0, 2 * np.pi)

def layout_arc(cell, layer, center, r, w, theta_start, theta_end, ex=None):
    # function to produce the layout of an arc
    # cell: layout cell to place the layout
    # layer: which layer to use
    # center: origin DPoint
    # r: radius
    # w: waveguide width
    # theta_start, theta_end: angle in radians
    # units in microns

    # example usage.  Places the ring layout in the presently selected cell.
    # cell = pya.Application.instance().main_window().current_view().active_cellview().cell
    # layout_arc(cell, layer, pya.DPoint(0,0), 10, 0.5, 0, np.pi/2)

    # fetch the database parameters

    if ex is None:
        ex = pya.DPoint(1, 0)

    delta_theta = np.arctan2(ex.y, ex.x)
    theta_start += delta_theta
    theta_end += delta_theta

    # optimal sampling
    arc_function = lambda t: np.array([center.x + r * np.cos(t), center.y + r * np.sin(t)])
    t, coords = sample_function(arc_function,
                                [theta_start, theta_end], tol=0.002 / r)

    # # This yields a better polygon
    coords = np.insert(coords, 0, arc_function(theta_start - 0.001),
                       axis=1)  # start the waveguide a little bit before
    coords = np.append(coords, np.atleast_2d(arc_function(theta_end + 0.001)).T,
                       axis=1)  # finish the waveguide a little bit after

    layout_waveguide(cell, layer, [pya.DPoint(x, y) for x, y in zip(*coords)], w)

def layout_arc_drc_exclude(cell, drc_layer, center, r, w, theta_start, theta_end, ex=None):
    corner_points = [center + (r + w / 2) * rotate(ex, theta_start),
                     center + (r - w / 2) * rotate(ex, theta_start),
                     center + (r + w / 2) * rotate(ex, theta_end),
                     center + (r - w / 2) * rotate(ex, theta_end)]
    for corner_point in corner_points:
        layout_square(cell, drc_layer, corner_point, 0.1, ex)

def layout_arc_with_drc_exclude(cell, layer, drc_layer, center, r, w, theta_start, theta_end, ex=None):
    layout_arc(cell, layer, center, r, w, theta_start, theta_end, ex)
    layout_arc_drc_exclude(cell, drc_layer, center, r, w, theta_start, theta_end, ex)

def layout_circle(cell, layer, center, r):
    # function to produce the layout of a filled circle
    # cell: layout cell to place the layout
    # layer: which layer to use
    # center: origin DPoint
    # r: radius
    # w: waveguide width
    # theta_start, theta_end: angle in radians
    # units in microns
    # optimal sampling

    arc_function = lambda t: np.array([center.x + r * np.cos(t), center.y + r * np.sin(t)])
    t, coords = sample_function(arc_function,
                                [0, 2 * np.pi - 0.001], tol=0.002 / r)

    dbu = cell.layout().dbu
    dpoly = pya.DPolygon([pya.DPoint(x, y) for x, y in zip(*coords)])
    cell.shapes(layer).insert(dpoly.to_itype(dbu))

def layout_path(cell, layer, point_iterator, w):
    path = pya.DPath(list(point_iterator), w, 0, 0).to_itype(cell.layout().dbu)
    cell.shapes(layer).insert(pya.Path.from_dpath(path))

def layout_path_with_ends(cell, layer, point_iterator, w):
    dpath = pya.DPath(list(point_iterator), w, w / 2, w / 2)
    cell.shapes(layer).insert(dpath)

def box_dpolygon(point1, point3, ex=None):
    # position point2 to the right of point1
    if ex is None:
        ex = pya.DPoint(1, 0)
    ey = rotate90(ex)
    point2 = point1 * ex * ex + point3 * ey * ey
    point4 = point3 * ex * ex + point1 * ey * ey

    return pya.DPolygon([point1, point2, point3, point4])

def square_dpolygon(center, width, ex=None):
    # returns the polygon of a square centered at center,
    # aligned with ex, with width in microns
    if ex is None:
        ex = pya.DPoint(1, 0)
    ey = rotate90(ex)
    quadrant = (width / 2) * (ex + ey)
    point1 = center + quadrant
    quadrant = rotate90(quadrant)
    point2 = center + quadrant
    quadrant = rotate90(quadrant)
    point3 = center + quadrant
    quadrant = rotate90(quadrant)
    point4 = center + quadrant

    return pya.DPolygon([point1, point2, point3, point4])

def layout_square(cell, layer, center, width, ex=None):
    """ Lays out a square in the DRC layer

    Args:
        center: pya.DPoint (um units)
        width: float (um units)
        ex: orientation

    """

    if ex is None:
        ex = pya.DPoint(1, 0)

    square = square_dpolygon(center, width, ex)
    cell.shapes(layer).insert(square)

def layout_taper(cell, layer, trans, w1, w2, length, insert = True):
    """ Lays out a taper

    Args:
        trans: pya.Trans: location and rotation
        w1: width of waveguide, float for DPoint type (microns); int for Point type (nm)
        w2: width of waveguide, float for DPoint type (microns); int for Point type (nm)
        length: length, float
        insert: flag to insert drawn waveguide or return shape, boolean

    """
    import pya
    if type(w1)==type(float()):
        pts = [pya.DPoint(0,-w1/2), pya.DPoint(0,w1/2), pya.DPoint(length,w2/2), pya.DPoint(length,-w2/2)]
        shape_taper = pya.DPolygon(pts).transformed(trans)
    else:
        pts = [pya.Point(0,-w1/2), pya.Point(0,w1/2), pya.Point(length,w2/2), pya.Point(length,-w2/2)]
        shape_taper = pya.Polygon(pts).transformed(trans)

    if insert == True:
        cell.shapes(layer).insert(shape_taper)
    else:
        return shape_taper

def layout_waveguide_sbend_bezier(cell, layer, trans, w=0.5, wo=None, h=2.0, length=15.0, insert = True):
    """ Creates a waveguide s-bend using a bezier curve
    Author: Lukas Chrostowski
    Args:
        trans: pya.Trans: location and rotation
        w: width of input waveguide, float for DPoint type (microns); int for Point type (nm)
        wo (optional): width of output waveguide, float
        h: height
        length: length
        insert: flag to insert drawn waveguide or return shape, boolean
    Usage:
        from SiEPIC.utils import get_layout_variables
        TECHNOLOGY, lv, ly, cell = get_layout_variables()
        layer = cell.layout().layer(TECHNOLOGY['Waveguide'])
        layout_waveguide_sbend_bezier(cell, layer, pya.Trans(), w=0.5, h=2.0, length=15.0, insert = True)
    """

    if wo==None:
        wo = w

    from SiEPIC.utils.geometry import bezier_parallel, translate_from_normal2
    from pya import DPoint, DPolygon, Point, Polygon

    if type(w)==type(int()):
        dbu = cell.layout().dbu
        w=w*dbu
        wo=wo*dbu
        h=h*dbu
        length=length*dbu
        trans=trans.to_dtype(dbu)

    p = bezier_parallel(DPoint(0,0), DPoint(length,h), 0)

    pt1 = translate_from_normal2(p,w/2,wo/2)
    pt2 = translate_from_normal2(p,-w/2, -wo/2)
    pt = pt1+pt2[::-1]

    poly = pya.DPolygon(pt)
    print(poly)
    poly_t = poly.transformed(trans)
    if insert == True:
        cell.shapes(layer).insert(poly_t)
        return poly_t.area()/((w+wo)/2)
    else:
        return poly_t

def layout_waveguide_sbend(cell, layer, trans, w=500, r=25000, h=2000, length=15000, insert = True):
    """ Lays out an s-bend

    Args:
        trans: pya.Trans: location and rotation
        w: width of waveguide, int
        r: radius, int
        h: height, int
        length: length, int
        insert: flag to insert drawn waveguide or return shape, boolean

    """

    from math import pi, cos, sin, log, sqrt, acos
    from SiEPIC.utils import points_per_circle
    import pya

    theta = acos(float(r-abs(h/2))/r)*180/pi
    x = int(2*r*sin(theta/180.0*pi))
    straight_l = int( (length - x)/2 )

    if (straight_l < 0):
        # Problem: target length is too short. increase
        print('SBend, too short: straight_l = %s' % straight_l)
        length = x
        straight_l = 0

    # waveguide_length = (2*pi*r*(2*theta/360.0)+straight_l*2)

    # define the cell origin as the left side of the waveguide sbend

    if (straight_l >= 0):
      circle_fraction = abs(theta) / 360.0
      npoints = int(points_per_circle(r*cell.layout().dbu) * circle_fraction)
      if npoints == 0:
        npoints = 1
      da = 2 * pi / npoints * circle_fraction # increment, in radians
      x1=straight_l
      x2=length-straight_l

      if h>0:
        y1=r
        theta_start1 = 270
        y2=h-r
        theta_start2 = 90
        pts = []
        th1 = theta_start1 / 360.0 * 2 * pi
        th2 = theta_start2 / 360.0 * 2 * pi
        pts.append(pya.Point.from_dpoint(pya.DPoint(0,w/2)))
        pts.append(pya.Point.from_dpoint(pya.DPoint(0,-w/2)))
        for i in range(0, npoints+1): # lower left
          pts.append(pya.Point.from_dpoint(pya.DPoint((x1+(r+w/2)*cos(i*da+th1))/1, (y1+(r+w/2)*sin(i*da+th1))/1)))
        for i in range(npoints, -1, -1): # lower right
          pts.append(pya.Point.from_dpoint(pya.DPoint((x2+(r-w/2)*cos(i*da+th2))/1, (y2+(r-w/2)*sin(i*da+th2))/1)))
        pts.append(pya.Point.from_dpoint(pya.DPoint(length,h-w/2)))
        pts.append(pya.Point.from_dpoint(pya.DPoint(length,h+w/2)))
        for i in range(0, npoints+1): # upper right
         pts.append(pya.Point.from_dpoint(pya.DPoint((x2+(r+w/2)*cos(i*da+th2))/1, (y2+(r+w/2)*sin(i*da+th2))/1)))
        for i in range(npoints, -1, -1): # upper left
          pts.append(pya.Point.from_dpoint(pya.DPoint((x1+(r-w/2)*cos(i*da+th1))/1, (y1+(r-w/2)*sin(i*da+th1))/1)))
      else:
        y1=-r
        theta_start1 = 90-theta
        y2=r+h
        theta_start2 = 270-theta
        pts = []
        th1 = theta_start1 / 360.0 * 2 * pi
        th2 = theta_start2 / 360.0 * 2 * pi
        pts.append(pya.Point.from_dpoint(pya.DPoint(length,h-w/2)))
        pts.append(pya.Point.from_dpoint(pya.DPoint(length,h+w/2)))
        for i in range(npoints, -1, -1): # upper right
          pts.append(pya.Point.from_dpoint(pya.DPoint((x2+(r-w/2)*cos(i*da+th2))/1, (y2+(r-w/2)*sin(i*da+th2))/1)))
        for i in range(0, npoints+1): # upper left
          pts.append(pya.Point.from_dpoint(pya.DPoint((x1+(r+w/2)*cos(i*da+th1))/1, (y1+(r+w/2)*sin(i*da+th1))/1)))
        pts.append(pya.Point.from_dpoint(pya.DPoint(0,w/2)))
        pts.append(pya.Point.from_dpoint(pya.DPoint(0,-w/2)))
        for i in range(npoints, -1, -1): # lower left
          pts.append(pya.Point.from_dpoint(pya.DPoint((x1+(r-w/2)*cos(i*da+th1))/1, (y1+(r-w/2)*sin(i*da+th1))/1)))
        for i in range(0, npoints+1): # lower right
         pts.append(pya.Point.from_dpoint(pya.DPoint((x2+(r+w/2)*cos(i*da+th2))/1, (y2+(r+w/2)*sin(i*da+th2))/1)))

      shape_bend = pya.Polygon(pts).transformed(trans)
      if insert == True:
        cell.shapes(layer).insert(shape_bend)
      else:
        return shape_bend

    print('SBend: theta %s, x %s, straight_l %s, r %s, h %s, length %s' % (theta, x, straight_l, r, h, length) )
    return length

def append_relative(points, *relative_vectors):
    """ Appends to list of points in relative steps """
    try:
        if len(points) > 0:
            origin = points[-1]
    except TypeError:
        raise TypeError("First argument must be a list of points")

    for vector in relative_vectors:
        points.append(origin + vector)
        origin = points[-1]
    return points

def place_cell(parent_cell, cell, placement_origin, params=None, relative_to=None):
    """ Places a cell and return ports
    Args:
        parent_cell: cell to place into
        cell: cell to be placed
        placement_origin: pya.Point object to be used as origin
        relative_to: port name

    Returns:
        ports(dict): key:port.name, value: geometry.Port with positions relative to parent_cell's origin
    """
    layout = parent_cell.layout()
    pcell, ports = cell.pcell(layout, params=params)
    if relative_to is not None:
        offset = next((port.position for port in ports if port.name == relative_to), None)
        placement_origin = placement_origin - offset
    parent_cell.insert(pya.CellInstArray(pcell.cell_index(),
                                         pya.Trans(pya.Trans.R0, placement_origin.to_itype(layout.dbu))))
    for port in ports:
        port.position += placement_origin

    return {port.name: port for port in ports}

def layout_connect_ports(cell, layer, port_from, port_to):

    P0 = port_from.position
    P3 = port_to.position
    angle_from = np.arctan2(port_from.direction.y, port_from.direction.x) * 180 / pi
    angle_to = np.arctan2(-port_to.direction.y, -port_to.direction.x) * 180 / pi

    curve = bezier_optimal(P0, P3, angle_from, angle_to)
    layout_waveguide(cell, layer, curve, [port_from.width, port_to.width])
    return curve_length(curve)

def make_pin(cell, name, center, w, layer, direction, debug=False):

    '''
    Makes a pin that SiEPIC-Tools will recognize
    cell: which cell to draw it in
    name: text label for the pin
    center: location, int [x,y]
    w: pin width
    layer: layout.layer() type
    direction =
        0: right
        90: up
        180: left
        270: down

    Units: intput can be float for microns, or int for nm
    '''

    from SiEPIC.extend import to_itype
    from pya import Point, DPoint
    import numpy
    dbu = cell.layout().dbu
    if type(w)==type(float()):
        w = to_itype(w,dbu)
        if debug:
            print('SiEPIC.utils.layout.make_pin: w converted to %s' %w )
    else:
        if debug:
            print('SiEPIC.utils.layout.make_pin: w %s' %w )
#    print(type(center[0]))
    if type(center) == type(Point()) or type(center) == type(DPoint()):
        center = [center.x, center.y]
    if type(center[0])==type(float()) or type(center[0])==type(numpy.float64()):
        center[0] = to_itype(center[0],dbu)
        center[1] = to_itype(center[1],dbu)
        if debug:
            print('SiEPIC.utils.layout.make_pin: center converted to %s' % (center)  )
    else:
        if debug:
            print('SiEPIC.utils.layout.make_pin: center %s' % (center)  )

    from SiEPIC._globals import PIN_LENGTH as pin_length

    direction = direction % 360
    if direction not in [0, 90, 180, 270]:
        raise('error in make_pin: direction must be one of [0, 90, 180, 270]')

    # text label
    t = pya.Trans(pya.Trans.R0, center[0],center[1])
    text = pya.Text (name, t)
    shape = cell.shapes(layer).insert(text)
    shape.text_dsize = float(w*dbu/2)
    shape.text_valign=1

    if direction == 0:
        p1 = pya.Point(center[0]-pin_length/2, center[1])
        p2 = pya.Point(center[0]+pin_length/2, center[1])
        shape.text_halign=2
    if direction == 90:
        p1 = pya.Point(center[0], center[1]-pin_length/2)
        p2 = pya.Point(center[0], center[1]+pin_length/2)
        shape.text_halign=2
        shape.text_rot=1
    if direction == 180:
        p1 = pya.Point(center[0]+pin_length/2, center[1])
        p2 = pya.Point(center[0]-pin_length/2, center[1])
        shape.text_halign=3
    if direction == 270:
        p1 = pya.Point(center[0], center[1]+pin_length/2)
        p2 = pya.Point(center[0], center[1]-pin_length/2)
        shape.text_halign=3
        shape.text_rot=1

    pin = pya.Path([p1,p2],w)
    cell.shapes(layer).insert(pin)

def y_splitter_tree(cell, tree_depth=4, y_splitter_cell="y_splitter_1310", library="SiEPICfab_Shuksan_PDK", wg_type='Strip TE 1310 nm, w=350 nm', draw_waveguides=True):
    '''
    Create a tree of splitters
    - cell: layout cell to create the structures in
    - tree_depth: Tree depth (2^N outputs)
    - y_splitter_cell: name of the y-splitter cell
    - library: the library containing the y_splitter_cell
    - wg_type: waveguide type from WAVEGUIDES.XML
    - draw_waveguides: True draws the waveguides, False is faster for debugging

    Returns
    - inst_in: instance of the input cell
    - inst_out[]: array of instances of the output cells
    - cell_tree: new cell created
    This is useful for subsequent routing

    Limitations:
    - the design uses regular 90 degree bends, rather than S-bends.
      hence it could be made more compact
    '''

    from SiEPIC.scripts import connect_pins_with_waveguide
    from SiEPIC.extend import to_itype
    from math import floor

    # create a new sub-cell where the tree will go
    ly = cell.layout()
    tech = ly.technology().name
    cell_tree = ly.create_cell("y_splitter_tree")

    # load the y-splitter from the library
    y_splitter = ly.create_cell(y_splitter_cell, library)
    if not y_splitter:
        raise Exception ('Cannot import cell %s:%s' % (library,y_splitter_cell))

    # Load waveguide information
    from SiEPIC.utils import load_Waveguides_by_Tech
    waveguides = load_Waveguides_by_Tech(tech)
    wg = [w for w in waveguides if wg_type in w['name'] ][0]
    if not wg:
        raise Exception("Waveguide type not defined in WAVEGUIDES.XML: %s" %wg_type )
        return
    wg_width = to_itype(float(wg['width']), ly.dbu)
    wg_radius = to_itype(float(wg['radius']), ly.dbu)

    # build the tree, using measurements from the cell and waveguide parameters
    x = 0
    dx = y_splitter.bbox().width() + wg_radius*2
    # calculate the spacing for the y-splitters based on waveguide radius and 90 degree bends
    y_wg_offset = (y_splitter.pinPoint("opt2").y-y_splitter.pinPoint("opt3").y)
    dy = max(y_splitter.bbox().height(), wg_radius*4 + y_wg_offset)
    # intialize loop
    inst_out = []
    y0 = 0
    for i in range(0,tree_depth):
        inst = []
        y = y0
        for j in range(0, int(2**(tree_depth-i-1))):
            t = pya.Trans(pya.Trans.R0, x, y)
            inst.append(cell_tree.insert(pya.CellInstArray(y_splitter.cell_index(), t)))
            # perform waveguide routing
            if (i > 0) and draw_waveguides:
                connect_pins_with_waveguide(
                    inst[j], 'opt2',
                    inst_higher[j*2+1], 'opt1', waveguide_type=wg_type)
                connect_pins_with_waveguide(
                    inst[j], 'opt3',
                    inst_higher[j*2], 'opt1', waveguide_type=wg_type)
            y += dy
        inst_higher = inst
        if i == 0:
            inst_out = inst
        if i == tree_depth-1:
            inst_in = inst[0]
        x += -dx
        y0 = y0 + dy/2
        dy = dy * 2

    return inst_in, inst_out, cell_tree


"""必要的工具"""
import numpy as np
import tensorflow as tf
from PIL import Image
import keras
import numpy as np
import math

def get_new_img_size(width, height, img_min_side=600):
    """把图片最小边resize到600"""
    if width <= height:
        f = float(img_min_side) / width
        resized_height = int(f * height)
        resized_width = int(img_min_side)
    else:
        f = float(img_min_side) / height
        resized_width = int(f * width)
        resized_height = int(img_min_side)

    return resized_width, resized_height

class BBoxUtility(object):
    def __init__(self, priors=None, overlap_threshold=0.7, ignore_threshold=0.3,
                 nms_thresh=0.7, top_k=300):
        self.priors = priors
        self.num_priors = 0 if priors is None else len(priors)
        self.overlap_threshold = overlap_threshold
        self.ignore_threshold = ignore_threshold
        self._nms_thresh = nms_thresh
        self._top_k = top_k
        self.boxes = tf.placeholder(dtype='float32', shape=(None, 4))
        self.scores = tf.placeholder(dtype='float32', shape=(None,))
        self.nms = tf.image.non_max_suppression(self.boxes, self.scores,
                                                self._top_k,
                                                iou_threshold=self._nms_thresh)
        self.sess = tf.Session(config=tf.ConfigProto(device_count={'GPU': 0}))

    @property
    def nms_thresh(self):
        return self._nms_thresh

    @nms_thresh.setter
    def nms_thresh(self, value):
        self._nms_thresh = value
        self.nms = tf.image.non_max_suppression(self.boxes, self.scores,
                                                self._top_k,
                                                iou_threshold=self._nms_thresh)

    @property
    def top_k(self):
        return self._top_k

    @top_k.setter
    def top_k(self, value):
        self._top_k = value
        self.nms = tf.image.non_max_suppression(self.boxes, self.scores,
                                                self._top_k,
                                                iou_threshold=self._nms_thresh)

    def iou(self, box):
        """
        box：真实框
        计算真实框与所有先验框的IOU值
        """
        # 计算出每个真实框与所有的先验框的iou
        # 判断真实框与先验框的重合情况
        inter_upleft = np.maximum(self.priors[:, :2], box[:2])
        inter_botright = np.minimum(self.priors[:, 2:4], box[2:])

        inter_wh = inter_botright - inter_upleft
        inter_wh = np.maximum(inter_wh, 0)
        inter = inter_wh[:, 0] * inter_wh[:, 1]
        # 真实框的面积
        area_true = (box[2] - box[0]) * (box[3] - box[1])
        # 先验框的面积
        area_gt = (self.priors[:, 2] - self.priors[:, 0]) * (self.priors[:, 3] - self.priors[:, 1])
        # 计算iou
        union = area_true + area_gt - inter

        iou = inter / union
        return iou

    def encode_box(self, box, return_iou=True):
        """如果重合度超过0.7，认为可以进行调节成真实框"""
        #box为传入的真实框，计算与先验框的重合度
        iou = self.iou(box)
        encoded_box = np.zeros((self.num_priors, 4 + return_iou))

        # 找到每一个真实框，重合程度较高的先验框
        # 如果iou>0.7，认为可以利用这个先验框回归到真实框
        assign_mask = iou > self.overlap_threshold
        if not assign_mask.any():
            assign_mask[iou.argmax()] = True
        if return_iou:
            encoded_box[:, -1][assign_mask] = iou[assign_mask]

        #对真实框与先验框进行编码，计算出应该有的预测结果，用于网络回归训练
        # 找到对应的先验框
        assigned_priors = self.priors[assign_mask]
        # 逆向编码，将真实框转化为FasterRCNN预测结果的格式
        # 先计算真实框的中心与长宽
        box_center = 0.5 * (box[:2] + box[2:])
        box_wh = box[2:] - box[:2]
        # 再计算重合度较高的先验框的中心与长宽
        assigned_priors_center = 0.5 * (assigned_priors[:, :2] +
                                        assigned_priors[:, 2:4])
        assigned_priors_wh = (assigned_priors[:, 2:4] -
                              assigned_priors[:, :2])

        # 逆向求取FasterRCNN应该有的预测结果
        encoded_box[:, :2][assign_mask] = box_center - assigned_priors_center
        encoded_box[:, :2][assign_mask] /= assigned_priors_wh
        encoded_box[:, :2][assign_mask] *= 4

        encoded_box[:, 2:4][assign_mask] = np.log(box_wh / assigned_priors_wh)
        encoded_box[:, 2:4][assign_mask] *= 4
        return encoded_box.ravel()

    def ignore_box(self, box):
        """
        box:标记中真实的框
        """
        # 获取所有先验框和真实框的重合程度
        iou = self.iou(box)

        ignored_box = np.zeros((self.num_priors, 1))

        # 找到每一个真实框，重合程度较高的先验框
        # 如果重合程度大于0.3或者小于0.7，就应该忽略这个框
        assign_mask = (iou > self.ignore_threshold) & (iou < self.overlap_threshold)

        if not assign_mask.any():
            assign_mask[iou.argmax()] = True

        #找出需要忽略的先验框
        ignored_box[:, 0][assign_mask] = iou[assign_mask]
        return ignored_box.ravel()

    def assign_boxes(self, boxes, anchors):
        """计算真实框对应的先验框，与这个先验框应当有的预测结果"""
        #先求出先验框个数
        self.num_priors = len(anchors)
        self.priors = anchors
        #创建一个全0矩阵，第一维是先验框个数，第二维的前4列是先验框的位置信息，第5列代表是否包含物体
        assignment = np.zeros((self.num_priors, 4 + 1))

        #初始时，认为所有先验框为背景，所以第5列=0
        assignment[:, 4] = 0.0
        if len(boxes) == 0:
            return assignment

        # 对每一个真实框都进行iou计算，找到需要忽略的先验框
        ingored_boxes = np.apply_along_axis(self.ignore_box, 1, boxes[:, :4])
        # 取重合程度最大的先验框，并且获取这个先验框的index
        ingored_boxes = ingored_boxes.reshape(-1, self.num_priors, 1)
        # (num_priors)
        ignore_iou = ingored_boxes[:, :, 0].max(axis=0)
        # (num_priors)
        ignore_iou_mask = ignore_iou > 0

        #代表忽略此先验框，将ignore_iou_mask的序号下的先验框设为忽略
        assignment[:, 4][ignore_iou_mask] = -1

        # (n, num_priors, 5)encode_box将真实框进行编码，计算出训练使用的预测结果：即偏移信息
        encoded_boxes = np.apply_along_axis(self.encode_box, 1, boxes[:, :4])
        # 每一个真实框的编码后的值，和iou
        # (n, num_priors)
        encoded_boxes = encoded_boxes.reshape(-1, self.num_priors, 5)

        #由于某些先验框会与多个真实框重合，所以要找iou最大的真实框进行对应
        # 取重合程度最大的先验框，并且获取这个先验框的index
        # (num_priors)
        best_iou = encoded_boxes[:, :, -1].max(axis=0)
        # (num_priors)
        best_iou_idx = encoded_boxes[:, :, -1].argmax(axis=0)
        # (num_priors)
        best_iou_mask = best_iou > 0
        # 某个先验框它属于哪个真实框
        best_iou_idx = best_iou_idx[best_iou_mask]

        assign_num = len(best_iou_idx)
        # 保留重合程度最大的先验框的应该有的预测结果
        # 哪些先验框存在真实框
        encoded_boxes = encoded_boxes[:, best_iou_mask, :]

        assignment[:, :4][best_iou_mask] = encoded_boxes[best_iou_idx, np.arange(assign_num), :4]
        # 4代表为背景的概率，为0
        # 1为正样本，代表有物体，0为负样本，代表背景，-1代表不要的框
        assignment[:, 4][best_iou_mask] = 1
        # 通过assign_boxes我们就获得了，输入进来的这张图片，应该有的预测结果是什么样子的
        # assignment即为找到的合理先验框
        return assignment

    def decode_boxes(self, mbox_loc, mbox_priorbox):
        """
        使用边框回归公式得到调整后的框
        mbox_loc: RPN输出的先验框的调整参数
        mbox_priorbox: 传入根据特征图生成好的先验框
        """
        # 获得先验框的宽与高
        prior_width = mbox_priorbox[:, 2] - mbox_priorbox[:, 0]
        prior_height = mbox_priorbox[:, 3] - mbox_priorbox[:, 1]

        # 获得先验框的中心点
        prior_center_x = 0.5 * (mbox_priorbox[:, 2] + mbox_priorbox[:, 0])
        prior_center_y = 0.5 * (mbox_priorbox[:, 3] + mbox_priorbox[:, 1])

        # 真实框距离先验框中心的xy轴偏移情况
        decode_bbox_center_x = mbox_loc[:, 0] * prior_width / 4
        decode_bbox_center_x += prior_center_x
        decode_bbox_center_y = mbox_loc[:, 1] * prior_height / 4
        decode_bbox_center_y += prior_center_y

        # 真实框的宽与高的求取
        decode_bbox_width = np.exp(mbox_loc[:, 2] / 4)
        decode_bbox_width *= prior_width
        decode_bbox_height = np.exp(mbox_loc[:, 3] / 4)
        decode_bbox_height *= prior_height

        # 获取真实框的左上角与右下角
        decode_bbox_xmin = decode_bbox_center_x - 0.5 * decode_bbox_width
        decode_bbox_ymin = decode_bbox_center_y - 0.5 * decode_bbox_height
        decode_bbox_xmax = decode_bbox_center_x + 0.5 * decode_bbox_width
        decode_bbox_ymax = decode_bbox_center_y + 0.5 * decode_bbox_height

        # 真实框的左上角与右下角进行堆叠
        decode_bbox = np.concatenate((decode_bbox_xmin[:, None],
                                      decode_bbox_ymin[:, None],
                                      decode_bbox_xmax[:, None],
                                      decode_bbox_ymax[:, None]), axis=-1)
        # 为方便计算，框的坐标值设置在0-1，即防止超出0与1
        decode_bbox = np.minimum(np.maximum(decode_bbox, 0.0), 1.0)
        return decode_bbox

    def detection_out(self, predictions, mbox_priorbox, num_classes, keep_top_k=300,confidence_threshold=0.5):
        """
        初步选出300个建议框
        predictions：RPN输出的结果，对每个先验框都要计算
        mbox_priorbox：传入根据特征图生成好的先验框
        """

        # 网络预测的结果
        # x_class即置信度
        mbox_conf = predictions[0]
        # 先验框的调整参数
        mbox_loc = predictions[1]
        # 先验框
        mbox_priorbox = mbox_priorbox
        results = []
        # 对每一个框进行处理
        for i in range(len(mbox_loc)):
            results.append([])
            #得到调整后的框，这里是所有的框
            decode_bbox = self.decode_boxes(mbox_loc[i], mbox_priorbox)
            #选出含物体的框，并用NMS删除重叠框
            for c in range(num_classes):
                c_confs = mbox_conf[i, :, c]
                #先进行对比，如果x_class的概率大于设置的置信度，认为框里有物体
                c_confs_m = c_confs > confidence_threshold
                if len(c_confs[c_confs_m]) > 0:
                    # 取出得分高于confidence_threshold的框
                    boxes_to_process = decode_bbox[c_confs_m]
                    # 获取我们所选出的框的置信度用于NMS
                    confs_to_process = c_confs[c_confs_m]
                    # 进行iou的非极大抑制，删除重叠框
                    feed_dict = {self.boxes: boxes_to_process,
                                 self.scores: confs_to_process}
                    idx = self.sess.run(self.nms, feed_dict=feed_dict)
                    # 取出在非极大抑制中效果较好的内容
                    good_boxes = boxes_to_process[idx]
                    confs = confs_to_process[idx][:, None]
                    # 将label、置信度、框的位置进行堆叠，label代表是否有物体，作用其实不大，因为置信度也可以反映
                    labels = c * np.ones((len(idx), 1))
                    c_pred = np.concatenate((labels, confs, good_boxes),
                                            axis=1)
                    # 添加进result里
                    results[-1].extend(c_pred)
            #将以上得到的框排序，选出前k个作为RPN的输出框，默认k=300
            if len(results[-1]) > 0:
                # 按照置信度进行排序
                results[-1] = np.array(results[-1])
                argsort = np.argsort(results[-1][:, 1])[::-1]
                results[-1] = results[-1][argsort]
                # 选出置信度最大的keep_top_k个
                results[-1] = results[-1][:keep_top_k]
        # 获得，在所有预测结果里面，置信度比较高的框
        # 另外，利用先验框和RPN的预测结果x_class,x_regr，调整到真实框（预测框）的位置
        return results

    def nms_for_out(self, all_labels, all_confs, all_bboxes, num_classes, nms):
        results = []
        nms_out = tf.image.non_max_suppression(self.boxes, self.scores,
                                               self._top_k,
                                               iou_threshold=nms)
        for c in range(num_classes):
            c_pred = []
            mask = all_labels == c
            if len(all_confs[mask]) > 0:
                # 取出得分高于confidence_threshold的框
                boxes_to_process = all_bboxes[mask]
                confs_to_process = all_confs[mask]
                # 进行iou的非极大抑制
                feed_dict = {self.boxes: boxes_to_process,
                             self.scores: confs_to_process}
                idx = self.sess.run(nms_out, feed_dict=feed_dict)
                # 取出在非极大抑制中效果较好的内容
                good_boxes = boxes_to_process[idx]
                confs = confs_to_process[idx][:, None]
                # 将label、置信度、框的位置进行堆叠。
                labels = c * np.ones((len(idx), 1))
                c_pred = np.concatenate((labels, confs, good_boxes), axis=1)
            results.extend(c_pred)
        return results


import numpy as np
import matplotlib.pyplot as plt
from .GetFieldLine import GetFieldLine

def PlotFieldLineDensity(pos,Params,fig=None,maps=[1,1,0,0],Overplot=False,**kwargs):
	'''
	Simple routine to plot the modelled plasma mass density along a field line.

	Args
	=====
		pos: 3 element ndarray containing the cartesian position to trace from in Rp.
		Params: For power law: 2-element array/list [p_eq,power].
				For Sandhu model: 5-element array/list [n0,alpha,a,beta,mav0] (see GetSandhuParams).
		fig: Pyplot instance, useful if plotting as a subplot on a pre-existing figure, by default will create new figure.
		maps: 4 element list containing subplot information [xmaps,ymaps,xmap,ymap] where xmaps and ymaps represent the number of subplots in the x and y directions and xmap and ymap are the specific x and y indices to plot to.
		Overplot: When set to True, this will inhibit the creation of a new subplot, fig must be supplied with a pre-existing figure, where the routine will plot over.

	Keyword Args
	============
		The following keyword argument form **kwargs - they are completely
		optional and depend on which model is being used.

		Model		kwargs
		T89			iopt,Kp,Vx,Vy,Vz,tilt,CoordIn,CoordOut,Alt,MaxLen,DSMax,FlattenSingleTraces,Verbose
		T96			parmod,Pdyn,SymH,By,Bz,Vx,Vy,Vz,tilt,CoordIn,CoordOut,Alt,MaxLen,DSMax,FlattenSingleTraces,Verbose
		T01			parmod,Pdyn,SymH,By,Bz,Vx,Vy,Vz,tilt,CoordIn,CoordOut,Alt,MaxLen,DSMax,FlattenSingleTraces,Verbose
		TS05		parmod,Pdyn,SymH,By,Bz,Vx,Vy,Vz,tilt,CoordIn,CoordOut,Alt,MaxLen,DSMax,FlattenSingleTraces,Verbose
		KT17

		MaxStepSize: Trace step size maximum in Rp (default = None).
		ModelArgs: 	Tuple containing arguments for the magnetic field models,
					when set to None, a set of default parameters are used.
					'T89'|'T96'|'T01'|'TS05'|'T89c'|'T96c'|'T01c'|'TS05c':
						ModelArgs = (Date,ut,CoordIn,CoordOut,Alt,MaxLen,DSMax,iopt,parmod,tilt,Vx,Vy,Vz)
						****NOTE****
						When using models 'T89'|'T96'|'T01'|'TS05' - the iopt,parmod,tilt,Vx,Vy,Vz parameters need not
						be specified as they will be calculated automatically within the Geopack model from Omni data
						using the Date and ut parameters

						To use those parameters, add 'c' tot he model name string: 'T89c'|'T96c'|'T01c'|'TS05c'
						Then all parameters will be needed

						************
						Date:  Date in format yyyymmdd.
						UT: Time in format hh.hh (hh + mm/60.0 + ss/3600.0).
						CoordIn: Coordinate system of input position, by default 'SM' (Solar-Magnetic), can be set to 'GSM' (Geocentric Solar Magnetospheric).
						CoordOut: Coordinate system of output positions and field vectors, by default 'SM', can be set to 'GSM'.
						Alt: Altitude to stop tracing at - default = 100km
						MaxLen: maximum number of trace steps
						DSMax: Maximum step size
						iopt: integer for controlling T89c model
						parmod: 10-element floating point array to control T96c,T01c and TS05c models,
								for T96:
									parmod[0] = Pdyn (nPa)
									parmod[1] = Dst (nT)
									parmod[2] = IMF By (nT)
									parmod[3] = IMF Bz (nT)
								for T01:
									parmod[0] = Pdyn (nPa)
									parmod[1] = Dst (nT)
									parmod[2] = IMF By (nT)
									parmod[3] = IMF Bz (nT)
									parmod[4] = G1 parameter (See Tsyganenko [2001])
									parmod[5] = G2 parameter (See Tsyganenko [2001])
								for TS05:
									parmod[0] = Pdyn (nPa)
									parmod[1] = Dst (nT)
									parmod[2] = IMF By (nT)
									parmod[3] = IMF Bz (nT)
									parmod[4] = W1 parameter (See Tsyganenko and Sitnov [2005])
									parmod[5] = W2 parameter (See Tsyganenko and Sitnov [2005])
									parmod[6] = W3 parameter (See Tsyganenko and Sitnov [2005])
									parmod[7] = W4 parameter (See Tsyganenko and Sitnov [2005])
									parmod[8] = W5 parameter (See Tsyganenko and Sitnov [2005])
									parmod[9] = W6 parameter (See Tsyganenko and Sitnov [2005])
						tilt: Geodipole tilt angle - if set to NaN then will be calculated using Date and ut
						Vx,Vy,Vz: IMF velocity components
					KT17:
						ModelArgs = (Rsun,DistIndex,MaxLen,InitStep,MaxStepSize,LimType)
						Rsun: radial distance from the Sun in AU
						DistIndex: Disturbance index (0.0 - 100.0)
						MaxLen: Maximum number of steps for trace
						InitStep: Starting step size.
						MaxStepSize: Maximum step size
						LimType: Integer value to define where to stop the field trace (default 0):
							0: Terminate trace at planet surface and magnetopause
							1: Confine to box -6 < x < 2, -4 < y < 4, -4 < z < 4
							2: Confine to box and terminate at planet
							3: Terminate trace at planet
							4: Trace to MP, Planet and stop at 10Rm

					KT14:
						ModelArgs = (MaxLen,InitStep,MaxStep,LimType,Rsm,t1,t2)
						MaxLen: Maximum number of steps for trace
						InitStep: Starting step size.
						MaxStepSize: Maximum step size
						LimType: Integer value to define where to stop the field trace (default 0):
							0: Terminate trace at planet surface and magnetopause
							1: Confine to box -6 < x < 2, -4 < y < 4, -4 < z < 4
							2: Confine to box and terminate at planet
							3: Terminate trace at planet
							4: Trace to MP, Planet and stop at 10Rm
						Rsm: Subsolar magnetopause radius (default=1.42).
						t1: Tail disk current strength (default=7.37).
						t2: Tail quasi-harris current sheet strength (default=2.16).
					Dipole:
						ModelArgs = [Beq]
						Beq: Megnatic field strength at equator in nT (Default=-31200.0).
		Delta: Separation between two traced field lines
		Polarization: 'none'|'toroidal'|'poloidal'
		Core: This only applies to the KT14/KT17 field, as it will include tracing to the core of the planet rather than the surface.

	Returns
	========
		pyplot instance

	'''
	tmp = GetFieldLine(pos,**kwargs)
	Polarization = kwargs.get('Polarization','none')
	if Polarization == 'none':
		T,s = tmp
	else:
		T,s,h = tmp

	Bm = np.sqrt(T.Bx**2.0 + T.By**2.0 + T.Bz**2.0).astype('float32')
	R = np.sqrt(T.x**2.0 + T.y**2.0 + T.z**2.0)
	maxR = np.float32(R.max())

	if np.size(Params) == 2:
		p = Params[0]*(maxR/R)**Params[1]
		label='$\\rho_{eq} = $'+'{:5.1f}'.format(Params[0])+', $m = $'+'{:3.1f}'.format(Params[1])
	else:
		Rnorm = R/maxR
		ne = Params[2]*np.exp(-0.5*((Rnorm-1.0)/0.1)**2) + Params[0]*Rnorm**(-Params[1])
		mav = Params[4]*Rnorm**(-Params[3])
		p = ne*mav
		label='$\\rho_{eq} = $'+'{:5.1f}'.format(Params[0]*Params[4])

	if fig is None:
		fig = plt
		fig.figure()
	if hasattr(fig,'Axes'):
		ax = fig.subplot2grid((maps[1],maps[0]),(maps[3],maps[2]))
	else:
		print('here')
		ax = fig

	ax.plot(s,p,label=label)
	ax.legend()
	ax.set_ylabel('$\\rho$ (amu cm$^{-3}$)')
	ax.set_xlabel('$x$ (km)')
	return ax


#!/usr/bin/env python3
import logging
from enum import Enum

import numpy as np

log = logging.getLogger('LabberDriver')

class PulseType(Enum):
    """Define possible qubit pulse types."""

    XY = 'XY'
    Z = 'Z'
    DELAY = 'Delay'
    READOUT = 'Readout'
    NONE = 'None'

class PulseShape(Enum):
    """Define possible qubit pulses shapes."""

    GAUSSIAN = 'Gaussian'
    SQUARE = 'Square'
    ZBIAS = 'Z-Bias'
    READ_FAST = 'Read_fast'
    CUSTOM = 'Custom'

class Pulse(object):
    """Represents physical pulses played by an AWG.

    Parameters
    ----------
    pulse_type : :obj:`PulseType`
        Pulse type (the default is PulseType.XY).
    shape : :obj:`PulseShape`
        Pulse shape (the default is PulseShape.GAUSSIAN).

    Attributes
    ----------
    length : float
        Pulse length
    amplitude : float
        Pulse amplitude.
    frequency : float
        Carrier frequency of pulse.
    phase : float
        Pulse phase.
    use_drag : bool
        If True, applies DRAG correction.
    drag_coefficient : float
        Drag coefficient.
    drag_detuning : float
        Applies a frequnecy detuning for DRAG pulses.
    gaussian_num_STDs : float
        The truncation range of Gaussian pulses,
        in units of standard deviations.
    square_rise_time : float
        How fast does the pulse edge rise

    """

    def __init__(self,  pulse_type=PulseType.Z, shape=PulseShape.SQUARE):

        # set variables
        self.pulse_type = pulse_type
        self.shape = shape

        self.length = 30e-9
        self.amplitude = 1.0
        self.frequency = 0.0
        self.phase = 0.0
        self.use_drag = False
        self.drag_coefficient = 0.0
        self.drag_detuning = 0.0
        self.gaussian_num_stds = 3.0
        self.square_rise_time = 2e-9

    def calculate_envelope(self, t,): #args will be for flux modulation on the Z-bias
        """Calculate pulse envelope.

        Parameters
        ----------
        t : numpy array
            Array with time values for which to calculate the pulse envelope.

        Returns
        -------
        envelope : numpy array
            Array containing pulse envelope.

        """
        dt = t[1] - t[0]
        len_pulse = len(t)
        envelope = np.zeros(len_pulse) #initialize envelope array

        if self.shape == PulseShape.SQUARE:

            len_edge = np.int(np.ceil(self.square_rise_time/dt)+1) #calculate length of rising edge in samples

            if len_pulse > 2*len_edge: #adding linear rising and falling edges to pulse if length of pulse > 2* length edge
                envelope[:len_edge] = np.linspace(0,1,len_edge)
                envelope[-len_edge:] = np.linspace(1,0,len_edge)
                envelope[len_edge:-len_edge] = 1

            else: #if length of pulse is smaller than 2*length of edge, then make an envelope with only falling and rising edges
                halfway_t = np.int(np.floor(len_pulse/2)) #halfway point of the pulse, in samples
                envelope[:halfway_t] = np.linspace(0,1, halfway_t)
                envelope[halfway_t:] = np.linspace(1,0, len_pulse - halfway_t) #I do len_pulse - halfway_t instead of just halfway_t if the number of samples is odd

        elif self.shape == PulseShape.GAUSSIAN:
            halfway_t = int(np.floor(len_pulse/2)) #calculate midpoint of pulse, in samples
            t0 = t[halfway_t] #midpoint of pulse
            σ = self.length/(2*self.gaussian_num_stds) #calculate σ of gaussian based on pulse length and num of std deviations we want to fit
            envelope = np.exp( -(t-t0)**2/(2*σ**2) )
            envelope = envelope - envelope.min() #make pulse start at 0
            envelope = envelope/envelope.max()

        elif self.shape == PulseShape.READ_FAST: #for readout pulses with a starting high amplitude transient region to ring-up RO resonator faster
            # Below I am hardcoding length of transient here to be 1/4 of pulse length, for it's shape to be Gaussian to minimize spectal leakage, and for it's amplitude to be 2.5 times larger than the steady region
            halfway_t = int(np.floor(len_pulse/8)) #halfway_t is midpoint of transient region, in samples
            t0 = t[halfway_t] #midpoint of transient region
            σ = (self.length/4)/(2*3) #I am harcoding 3 gaussian standard deviations here
            envelope[:halfway_t] = 2.5*np.exp( -(t[:halfway_t]-t0)**2/(2*σ**2) ) #the first half of the transient region is a pure Gaussian
            envelope[halfway_t:2*halfway_t] = 1.5*np.exp( -(t[halfway_t:2*halfway_t]-t0)**2/(2*σ**2) )
            envelope[halfway_t:] = envelope[halfway_t:] + np.ones(len_pulse-halfway_t) #the second region of the transient region is a down-scaled gaussian + the steady part; such that the second half of the transient smoothly converges to the steady part
            envelope = envelope - envelope.min()

        envelope = envelope * self.amplitude

        return envelope

        #if self.shape == PulseShape.ZBIAS: #Z-bias pulse consists of starting from the last Z-bias value and ramping to a new Z-bias value
            #len_edge = np.int(np.ceil(self.square_rise_time/dt)+1) #calculate length of rising edge in samples
            #if len_pulse > len_edge:
            #    envelope[:len_edge] = np.linspace(last_value_Zbias,self.amplitude,len_edge) #make rising edge start from last Z-bias value and end at the new Z-bias level
            #    envelope[len_edge:] = self.amplitude
            #else:
            #    envelope = np.linspace(last_value_Zbias, self.amplitude, len_pulse) #if length of pulse is smaller than rising edge, make the whole Z-bias pulse a rising edge
            #envelope = np.ones(len_pulse) * self.amplitude + args[0] * np.cos(2*pi*args[1] * 1e6* (t-t[0]))
            #return envelope

    def calculate_waveform(self, t, IQ_ratio = 1.0, IQ_skew = 0.0, I_offset = 0.0, Q_offset = 0.0, custom_waveform = None,
                            Z_bias_mod_amp = 0.0, Z_bias_mod_freq = 0.0): #custom waveform input is for Custom pulses
        """Calculate pulse waveform including carrier frequency, phase offsets, IQ mixer compensations, and DRAG implementation

        Parameters
        ----------
        t : numpy array
            Array with time values for which to calculate the pulse waveform.

        Returns
        -------
        waveform : numpy array
            Array containing pulse waveform.

        """
        π = np.pi
        if self.shape == PulseShape.CUSTOM: #make the "envelope" the custom waveform for a "Custom" Pulse
            envelope = custom_waveform
        elif self.shape == PulseShape.ZBIAS:
            envelope = np.ones(len(t)) * self.amplitude + Z_bias_mod_amp * np.cos(2*np.pi*Z_bias_mod_freq * 1e6* (t-t[0]))
        else:
            envelope = self.calculate_envelope(t)
        ω = self.frequency * 2 * π
        dt = t[1]-t[0]

        if self.pulse_type == PulseType.XY: #below I generate I/Q waveforms with the IQ parameters and DRAG
            DRAG_δ = self.drag_detuning * 2 * π
            I_waveform = (envelope * np.cos( (ω + DRAG_δ) * t + self.phase ) +
                          -self.drag_coefficient * np.gradient(envelope)/dt * np.sin( (ω + DRAG_δ) * t + self.phase ) ) + I_offset
            Q_waveform = IQ_ratio*(envelope * np.sin( (ω + DRAG_δ) * t + self.phase + IQ_skew ) -
                          -self.drag_coefficient * np.gradient(envelope)/dt * np.cos( (ω + DRAG_δ) * t + self.phase + IQ_skew ) ) + Q_offset
            waveform = I_waveform + 1j*Q_waveform

        elif self.pulse_type == PulseType.Z:
            t_noPhase = t-t[0] #make sure pulse doesn't have any intrinsic phase regardless of when it starts
            waveform = envelope * np.cos(ω * t_noPhase + self.phase)

        elif self.pulse_type == PulseType.READOUT:
            I_waveform = envelope * np.cos(ω * t + self.phase) + I_offset
            Q_waveform = envelope * IQ_ratio * np.sin(ω * t + self.phase + IQ_skew) + Q_offset
            waveform = I_waveform + 1j*Q_waveform

        elif self.pulse_type == PulseType.NONE or self.pulse_type == PulseType.DELAY:
            waveform = np.zeros(len(t))

        return waveform

if __name__ == '__main__':
    pass


"""
Functions for explaining classifiers that use Multimodal data.
Developed initially for the hateful memes challenge in mmf
"""

import copy
from functools import partial
import numpy as np
import scipy as sp
import sklearn
from sklearn.utils import check_random_state
from tqdm.auto import tqdm

from lime.wrappers.scikit_image import SegmentationAlgorithm
from mmxai.interpretability.classification.lime.lime_base import LimeBase

from lime.lime_text import IndexedString, TextDomainMapper
from lime.exceptions import LimeError

from sklearn.calibration import CalibratedClassifierCV
from PIL import Image
from skimage.segmentation import mark_boundaries

# from object_detection import *

class MultiModalExplanation(object):
    def __init__(
        self,
        image,
        segments,
        domain_mapper,
        n_txt_features,
        n_img_features,
        n_detection_features,
        ratio_txt_img,
        detection_label,
        mode="classification",
        class_names=None,
        random_state=None,
    ):
        """Init function.

        Args:
            image: 3d numpy array
            segments: 2d numpy array, with the output from skimage.segmentation
            domain_mapper: Maps text feature ids to words or word-positions
            n_txt_features:number of words to include in explanation
            n_img_features: number of superpixels to include in explanation
            n_detection_features:number of detected objects to include in explanation
            ratio_txt_img: weight ratio between text and image features
            detection_label: numpy array of names of detected objects
            mode: what kind of model for this object, classification default
            class_names: list of class names, ordered according to whatever the
                classifier is using. If not present, class names will be '0',
                '1', ...
            random_state: an integer or numpy.RandomState that will be used to
                generate random numbers. If None, the random state will be
                initialized using the internal numpy seed.
        """
        self.image = image
        self.segments = segments
        self.random_state = random_state
        self.mode = mode
        self.domain_mapper = domain_mapper
        self.n_txt_features = n_txt_features
        self.n_img_features = n_img_features
        self.n_detection_features = n_detection_features
        self.ratio_txt_img = ratio_txt_img
        self.detection_label = detection_label
        self.local_exp = {}
        self.intercept = {}
        self.score = {}
        self.local_pred = {}
        self.unsorted_weights = {}

        # divide explanations of the two modalities
        self.local_exp_img = {}
        self.local_exp_txt = {}
        self.local_exp_det = {}

        if mode == "classification":
            self.class_names = class_names
            self.top_labels = None
            self.predict_proba = None

    def get_image_and_mask(
        self,
        label,
        positive_only=True,
        negative_only=False,
        hide_rest=False,
        num_features=5,
        min_weight=0.0,
    ):
        """

        Args:
            label: label to explain
            positive_only: if True, only take superpixels that positively contribute to
                the prediction of the label.
            negative_only: if True, only take superpixels that negatively contribute to
                the prediction of the label. If false, and so is positive_only, then both
                negativey and positively contributions will be taken.
                Both can't be True at the same time
            hide_rest: if True, make the non-explanation part of the return
                image gray
            num_features: number of superpixels to include in explanation
            min_weight: minimum weight of the superpixels to include in explanation
        Returns:
            (image, mask), where image is a 3d numpy array and mask is a 2d
            numpy array that can be used with
            skimage.segmentation.mark_boundaries
        """

        if label not in self.local_exp_img:
            raise KeyError("Label not in explanation")
        if positive_only & negative_only:
            raise ValueError(
                "Positive_only and negative_only cannot be true at the same time."
            )
        segments = self.segments
        image = self.image
        exp = self.local_exp_img[label]
        mask = np.zeros(segments.shape, segments.dtype)
        if hide_rest:
            temp = np.zeros(self.image.shape)
        else:
            temp = self.image.copy()
        if positive_only:
            fs = [x[0] for x in exp if x[1] > 0 and x[1] > min_weight][:num_features]
        if negative_only:
            fs = [x[0] for x in exp if x[1] < 0 and abs(x[1]) > min_weight][
                :num_features
            ]
        if positive_only or negative_only:
            for f in fs:
                temp[segments == f] = image[segments == f].copy()
                mask[segments == f] = 1
            return temp, mask
        else:
            for f, w in exp[:num_features]:
                if np.abs(w) >= min_weight:
                    c = 0 if w < 0 else 1
                    mask[segments == f] = -1 if w < 0 else 1
                    temp[segments == f] = image[segments == f].copy()
                    temp[segments == f, c] = np.max(image)
            return temp, mask

    def as_list(self, label=1, **kwargs):
        """
        Returns the explanation as a list.
        Args:
            label: desired label. If you ask for a label for which an
                explanation wasn't computed, will throw an exception.
                Will be ignored for regression explanations.
            kwargs: keyword arguments, passed to domain_mapper
        Returns:
            list of tuples (representation, weight), where representation is
            given by domain_mapper. Weight is a float.
        """
        label_to_use = label
        ans = self.domain_mapper.map_exp_ids(self.local_exp_txt[label_to_use], **kwargs)
        ans = [(x[0], float(x[1])) for x in ans]
        return ans

    def get_explanation(self, label, num_features=10, which_exp="positive"):

        # the explanation to display:
        """
        :param label: label to explain
        :param num_features: how many top features to display
        :param which_exp: want features that encourage or discourage the dicision (label)
        :return:
            text_message: informative message to interpret text features
            img_message: informative message to interpret image features
            txt_exp_list: text part of the explanation, ready to display
            temp, mask: image part of the explanation, ready to display
        """
        this_exp = np.array(self.local_exp[label])

        if which_exp == "positive":
            positives = this_exp[this_exp[:, 1] >= 0]
        else:  # negative
            positives = this_exp[this_exp[:, 1] < 0]

        if positives.shape[0] < num_features:
            num_features = positives.shape[0]
        top_exp = positives[:num_features]
        top_exp_unique, top_idx = np.unique(top_exp[:, 0], return_index=True)

        txt_exp = top_exp[top_exp[:, 0] < self.n_txt_features]
        n_txt_exp = txt_exp.shape[0]
        txt_top_idx = []
        for txt_feature in txt_exp:
            txt_top_idx.append(top_idx[top_exp_unique == txt_feature[0]] + 1)

        img_exp = top_exp[self.n_txt_features <= top_exp[:, 0]]
        img_exp = img_exp[img_exp[:, 0] < (self.n_txt_features + self.n_img_features)]
        n_img_exp = img_exp.shape[0]
        img_top_idx = []
        for img_feature in img_exp:
            img_top_idx.append(top_idx[top_exp_unique == img_feature[0]] + 1)

        # detection features
        det_exp = top_exp[top_exp[:, 0] >= self.n_txt_features + self.n_img_features]
        n_det_exp = det_exp.shape[0]
        det_top_idx = []
        for det_feature in det_exp:
            det_top_idx.append(
                det_feature[0] - n_txt_exp - n_img_exp
            )  # index for retrieving the labels

        if n_det_exp != 0:
            readable_exp_det = (
                f" Also, we have detected {n_det_exp} types "
                f"of objects from the input image that can be the reason for the decision, they are:"
            )
            for i in det_top_idx:
                readable_exp_det += str(self.detection_label[i])
        else:
            readable_exp_det = (
                " No objects in the image contributed to the model decision"
            )

        # explanation of explanations
        readable_exp_txt = ""
        readable_exp_img = ""
        if n_txt_exp > 0:
            readable_exp_txt = f"{n_txt_exp} are from the text (the top"
            for i in txt_top_idx:
                readable_exp_txt += str(i)
            readable_exp_txt += "th), "
        else:
            readable_exp_txt += "none are from the words, "

        if n_img_exp > 0:
            readable_exp_img = f"{n_img_exp} are from the image (the top"
            for j in img_top_idx:
                readable_exp_img += str(j)
            readable_exp_img += (
                "th, some adjacent regions might merge into a larger area)."
            )
        else:
            readable_exp_img += "none are from the image pixel areas."

        txt_exp_list = np.array(self.as_list(label), dtype="object")

        # return explanations upon request
        if which_exp == "positive":
            # txt_list = txt_exp_list[txt_exp_list[:, 1] >= 0]
            temp, mask = self.get_image_and_mask(
                label, num_features=n_img_exp, positive_only=True
            )
        else:
            # txt_list = txt_exp_list[txt_exp_list[:, 1] < 0]
            temp, mask = self.get_image_and_mask(
                label, num_features=n_img_exp, positive_only=False, negative_only=True
            )
        # txt_list = txt_list[:n_txt_exp]

        # image and text hover display message
        label_decision = ""
        if label == 1:
            label_decision = "hateful"
        else:
            label_decision = "not hateful"

        txt_message = (
            f"For this result, the value associated with each word indicates how much it pushes "
            f"the model towards making a {label_decision} decision."
        )

        img_message = (
            f"Your image has been segmented into {self.n_img_features} small pixel areas, "
            f"the ones that most encourage (or discourage) your "
            f"model decision has been marked by the yellow boundaries."
        )

        top_message = (
            f"Each small pixel area in your image input and each distinct word "
            f"in your text input are called an interpretable feature. There are "
            f"{self.n_img_features + self.n_txt_features} features in total ({self.n_img_features} "
            f"pixel areas and {self.n_txt_features} words). Among the top 10 "
            f"such features that encourage (or discourage) your model decision, "
        )
        top_message = top_message + readable_exp_txt + readable_exp_img

        ratio_message = self.get_txt_img_ratio()

        # format new line as html
        img_message = "
" + img_message + "

" + top_message + "
"
        img_message = img_message + ratio_message
        return txt_message, img_message, txt_exp_list, temp, mask

    def get_txt_img_ratio(self):
        """
        Get informative message about the weight ratio between text and image features
        Return:
            words: informative string message to interpret the relative weight of text and image features
        """
        img_percentage = 1 / (1 + self.ratio_txt_img)
        txt_percentage = 1 - img_percentage
        words = (
            f"For this prediction, the relative importance of "
            f"text and image inputs to your model decision are respectively {round(100*txt_percentage, 2)}% "
            f"and {round(100*img_percentage, 2)}%"
        )
        return words

class LimeMultimodalExplainer(object):
    def __init__(
        self,
        image,
        text,
        model,
        kernel_width=0.25,
        kernel=None,
        feature_selection="auto",
        class_names=None,
    ):
        """
        Object to explain predictions on texts and images
        Args:
            image: input 3D numpy array
            text: input string in the meme
            model: multi-modal model to give predictions for given image and text
            kernel_width: kernel width for the exponential kernel.
            If None, defaults to sqrt(number of columns) * 0.75.
            kernel: similarity kernel that takes euclidean distances and kernel
                width as input and outputs weights in (0,1). If None, defaults to
                an exponential kernel.
            feature_selection: feature selection method. can be
                'forward_selection', 'lasso_path', 'none' or 'auto'.
                See function 'explain_instance_with_data' in lime_base.py for
                details on what each of the options does.
                class_names: list of class names, ordered according to whatever the
                classifier is using. If not present, class names will be '0',
                '1', ...
        """
        self.image = image
        self.text = text
        self.pred_model = model
        self.random_state = check_random_state(None)
        if kernel is None:

            def kernel(d, kernel_width):
                return np.sqrt(np.exp(-(d ** 2) / kernel_width ** 2))

        kernel_fn = partial(kernel, kernel_width=kernel_width)
        self.feature_selection = feature_selection
        self.class_names = class_names
        self.base = LimeBase(kernel_fn, verbose=False, random_state=self.random_state)

    def explain_instance(self, classifier_fn, n_samples, top_labels=2):

        """
        Generate explanations for a multi-modal input
        Arguments:
            classifier_fn: classification function to give predictions for given texts and images
            num_samples: size of the neighborhood to learn the linear model
            top_labels: if not None, ignore labels and produce explanations for
                the K labels with highest prediction probabilities, where K is
                this parameter.
            Return:
                ret_exp: A MultiModalExplanation object with the corresponding
            explanations.
        """
        (
            data,
            labels,
            distances,
            n_txt_features,
            n_img_features,
            segments,
            domain_mapper,
            n_detection_features,
            detection_label,
            ratio_txt_img,
        ) = self.data_labels(n_samples, classifier_fn)
        num_features = data.shape[1]

        if self.class_names is None:
            self.class_names = [str(x) for x in range(labels[0].shape[0])]

        ret_exp = MultiModalExplanation(
            self.image,
            segments,
            domain_mapper=domain_mapper,
            n_txt_features=n_txt_features,
            n_img_features=n_img_features,
            n_detection_features=n_detection_features,
            ratio_txt_img=ratio_txt_img,
            detection_label=detection_label,
            class_names=self.class_names,
            random_state=self.random_state,
        )
        ret_exp.predict_proba = labels[0]

        if top_labels:
            top = np.argsort(labels[0])[-top_labels:]
            ret_exp.top_labels = list(top)
            ret_exp.top_labels.reverse()
        for label in top:
            (
                ret_exp.intercept[label],
                ret_exp.unsorted_weights[label],
                ret_exp.local_exp[label],
                ret_exp.local_exp_txt[label],
                ret_exp.local_exp_img[label],
                ret_exp.local_exp_det[label],
                ret_exp.score[label],
                ret_exp.local_pred[label],
            ) = self.base.explain_instance_with_data(
                data,
                labels,
                distances,
                label,
                num_features,
                n_txt_features,
                n_img_features,
                n_detection_features,
                feature_selection=None,
            )

            # split local explanation into text and image features

        return ret_exp

    def data_labels(self, num_samples, classifier_fn, detection=False):
        """
        Steps of this function:
            1. generate perturbed text features and image features
            2. in a loop, 1) using these features to make instances of perturbed (text, image) pairs,
                          2) make predictions on these pairs, store labels into 'labels'
            3. concatenate text and image features, store into 'data',
                also append the original input and prediction of it
            4. calculate distances
            Arguments:
                classifier_fn: classification function to give predictions for given texts and images
                num_samples: size of the neighborhood to learn the linear model
                detection: Whether object detection method is invoked, default to be false
            Return:
            data: dense num_samples * num_superpixels
            labels: prediction probabilities matrix
            distances:distance including text/image distance ratio where
            text and image distance are cosine distances between the original instance and
                    each perturbed instance (computed in the binary 'data'
                    matrix), times 100.
            doc_size: number of words in indexed string, where indexed string is the string with various indexes
            n_img_features: number of superpixels to include in explanation
            segments:2d numpy array, with the output from skimage.segmentation
            domain_mapper:Maps text feature ids to words or word-positions
            num_object_detection:number of detected objects to include in explanation
            ori_label: numpy including deteced objects in the original image
            ratio_txt_img: weight ratio between text and image features
        """

        """ 1. make text features """
        indexed_string = IndexedString(
            self.text, bow=True, split_expression=r"\W+", mask_string=None
        )
        domain_mapper = TextDomainMapper(indexed_string)

        doc_size = indexed_string.num_words()
        sample = self.random_state.randint(
            1, doc_size + 1, num_samples
        )  # num_samples - 1
        data_txt = np.ones((num_samples, doc_size))
        # data[0] = np.ones(doc_size)
        features_range = range(doc_size)
        inverse_data_txt = []

        """ 1. make image features """
        random_seed = self.random_state.randint(0, high=1000)
        segmentation_fn = SegmentationAlgorithm(
            "quickshift",
            kernel_size=4,
            max_dist=200,
            ratio=0.2,
            random_seed=random_seed,
        )

        # segmentation_fn = SegmentationAlgorithm('felzenszwalb', scale=200, sigma=2, min_size=100)
        """segmentation_fn = SegmentationAlgorithm('slic', n_segments=60, compactness=10, sigma=1,
                     start_label=1)"""

        segments = segmentation_fn(self.image)  # get segmentation
        n_img_features = np.unique(segments).shape[0]  # get num of superpixel features
        data_img = self.random_state.randint(
            0, 2, n_img_features * num_samples
        ).reshape((num_samples, n_img_features))
        data_img_rows = tqdm(data_img)
        imgs = []

        """ 1. make object detection features
        if detection:
            predictor, cfg = object_detection_predictor()
            ori_label = object_detection_obtain_label(predictor, cfg, self.image)
            num_object_detection = ori_label.shape[0]
            data_object_detection = np.zeros((num_samples,num_object_detection))"""

        # create fudged_image
        fudged_image = self.image.copy()
        for x in np.unique(segments):
            fudged_image[segments == x] = (
                np.mean(self.image[segments == x][:, 0]),
                np.mean(self.image[segments == x][:, 1]),
                np.mean(self.image[segments == x][:, 2]),
            )

        # img_features[0, :] = 1  # the first sample is the full image                                # num_samples

        """2. create data instances and make predictions"""
        labels = []
        for i, instance in enumerate(zip(sample, data_img_rows)):
            size_txt, row_img = instance

            # make text instance
            inactive = self.random_state.choice(features_range, size_txt, replace=False)
            data_txt[i, inactive] = 0
            inverse_data_txt.append(indexed_string.inverse_removing(inactive))

            # make image instance
            temp = copy.deepcopy(self.image)
            zeros = np.where(row_img == 0)[
                0
            ]  # get segment numbers that are turned off in this instance
            mask = np.zeros(segments.shape).astype(bool)
            for zero in zeros:
                mask[segments == zero] = True
            temp[mask] = fudged_image[mask]

            """if detection:
                label = object_detection_obtain_label(predictor, cfg, temp)
                label_diff = compare_labels(ori_label,label)
                data_object_detection[i] = label_diff"""
            imgs.append(temp)

            # make prediction and append result
            if len(imgs) == 10:
                preds = classifier_fn(self.pred_model, imgs, inverse_data_txt)
                labels.extend(preds)
                imgs = []
                inverse_data_txt = []

        if len(imgs) > 0:
            preds = classifier_fn(self.pred_model, imgs, inverse_data_txt)
            labels.extend(preds)

        """3. concatenate and append features"""
        data = np.concatenate((data_txt, data_img), axis=1)

        # append the original input to the last
        orig_img_f = np.ones((n_img_features,))
        orig_txt_f = np.ones(doc_size)

        """if detection:
            data = np.concatenate((data, data_object_detection),axis=1)
            orig_ot = np.ones(num_object_detection)
            data = np.vstack((data, np.concatenate((np.concatenate((orig_txt_f, orig_img_f)),orig_ot))))
        else:"""
        data = np.vstack((data, np.ones((data.shape[1]))))  ###

        labels.extend(classifier_fn(self.pred_model, [self.image], [self.text]))

        """4. compute distance# distances[:, :(doc_size-1)] *= 100
            use platt scaling t get relative importance of text and image modalities
        """

        labels = np.array(labels, dtype=float)

        # Modify MMF source code to zero out image / text attributes
        # dummy_label_image = np.array(classifier_fn([self.image], [self.text], zero_text=True))  # zero out text
        # dummy_label_text = np.array(classifier_fn([self.image], [self.text], zero_image=True))  # zero out image

        # perform calibration
        try:
            labels_for_calib = np.array(labels[:, 0] < 0.5, dtype=float)
            calibrated = CalibratedClassifierCV(cv=3)
            calibrated.fit(data[:, : doc_size + n_img_features], labels_for_calib)

            calib_data = np.ones((3, doc_size + n_img_features), dtype=float)
            calib_data[0][:doc_size] = 0  # zero out text
            calib_data[1][doc_size:] = 0  # zero out image
            calibrated_labels = calibrated.predict_proba(calib_data)

            delta_txt = abs(calibrated_labels[-1][0] - calibrated_labels[0][0])
            delta_img = abs(calibrated_labels[-1][0] - calibrated_labels[1][0])

            ratio_txt_img = max(min(100, delta_txt / delta_img), 0.01)
        except:
            dummy_text = ""
            dummy_image = np.zeros_like(self.image)
            label_text_out = np.array(
                classifier_fn(
                    self.pred_model, [self.image], [self.text], zero_text=True
                )
            )  # zero out text
            label_image_out = np.array(
                classifier_fn(
                    self.pred_model, [self.image], [self.text], zero_image=True
                )
            )  # zero out image

            delta_txt = abs(labels[-1][0] - label_text_out[0][0])
            delta_img = abs(labels[-1][0] - label_image_out[0][0])
            ratio_txt_img = max(min(10, delta_txt / delta_img), 0.1)

        # calculate distances
        distances_img = sklearn.metrics.pairwise_distances(
            data[:, doc_size:], data[-1, doc_size:].reshape(1, -1), metric="cosine"
        ).ravel()

        def distance_fn(x):
            return sklearn.metrics.pairwise.pairwise_distances(
                x, x[-1], metric="cosine"
            ).ravel()

        distances_txt = distance_fn(sp.sparse.csr_matrix(data[:, :doc_size]))

        distances = (
            1 / (1 + ratio_txt_img) * distances_img
            + (1 - 1 / (1 + ratio_txt_img)) * distances_txt
        )

        # As required by lime_base, make the first element of data, labels, distances the original data point
        data[0] = data[-1]
        labels[0] = labels[-1]
        distances[0] = distances[-1]

        """if not detection:"""
        num_object_detection = 0
        ori_label = None

        return (
            data,
            labels,
            distances,
            doc_size,
            n_img_features,
            segments,
            domain_mapper,
            num_object_detection,
            ori_label,
            ratio_txt_img,
        )


from CHECLabPy.plotting.setup import Plotter
import numpy as np
import pandas as pd
from sstcam_simulation.utils.efficiency import CameraEfficiency
from sstcam_simulation.utils.sipm import SiPMOvervoltage
from sstcam_simulation.utils.window_durham_needle import SSTWindowRun3, Window

MINIMGAMP_GAMMA = 250
MINIMGAMP_PROTON = 480
BTEL1170_PDE = 0.2
BTEL0090_SNR = 0.35
PROD4_OPCT = 0.08
PROD4_MINIMGAMP_GAMMA = 153
PROD4_MINIMGAMP_PROTON = 208

def main():
    sipm_tool = SiPMOvervoltage.lvr3_6mm_50um_uncoated()
    sipm_tool.overvoltage = 6
    pde_at_450nm = sipm_tool.pde
    window_tool = SSTWindowRun3()

    eff = CameraEfficiency.from_sstcam(
        fov_angle=0,
        pde_at_450nm=pde_at_450nm,
        window=window_tool,
    )

    p = Plotter()
    p.ax.plot(eff.wavelength, eff._cherenkov_diff_flux_on_ground, "-", color='blue', alpha=0.3, label="On-ground")
    p.ax.plot(eff.wavelength, eff._cherenkov_diff_flux_at_camera, "--", color='blue', alpha=0.6, label="At-camera")
    p.ax.plot(eff.wavelength, eff._cherenkov_diff_flux_at_pixel, "-.", color='blue', alpha=0.9, label="At-pixel")
    p.ax.plot(eff.wavelength, eff._cherenkov_diff_flux_inside_pixel, ":", color='blue', alpha=1, label="Inside-pixel")
    p.ax.set_ylim(0, 0.45)
    p.ax.set_xlabel("Wavelength [nm]")
    p.ax.set_ylabel("Cherenkov photons [100 * 1 / nm]")
    p.save("foldings/spectra_cherenkov.pdf")

    p = Plotter()
    p.ax.plot(eff.wavelength, eff._nsb_diff_flux_on_ground, "-", color='red', alpha=0.3, label="On-ground")
    p.ax.plot(eff.wavelength, eff._nsb_diff_flux_at_camera, "--", color='red', alpha=0.6, label="At-camera")
    p.ax.plot(eff.wavelength, eff._nsb_diff_flux_at_pixel, "-.", color='red', alpha=0.9, label="At-pixel")
    p.ax.plot(eff.wavelength, eff._nsb_diff_flux_inside_pixel, ":", color='red', alpha=1, label="Inside-pixel")
    p.ax.set_ylim(0, 25)
    p.ax.set_xlabel("Wavelength [nm]")
    p.ax.set_ylabel("NSB photons [ 1 / (nm m2 ns sr) ]")
    p.save("foldings/spectra_nsb.pdf")

    p = Plotter()
    camera_sensitivity = eff.pde * eff.window_transmissivity
    p.ax.plot(eff.wavelength, eff.pde, label="PDE")
    p.ax.plot(eff.wavelength, eff.window_transmissivity, label="Window Transmissivity")
    p.ax.plot(eff.wavelength, camera_sensitivity, label="PDE * Window Transmissivity")
    # p.ax.plot(eff.wavelength, eff.telescope_transmissivity, label="Telescope Transmissivity")
    p.ax.plot(eff.wavelength, eff.mirror_reflectivity, label="Mirror Reflectivity")
    p.ax.set_xlabel("Wavelength [nm]")
    p.ax.set_ylabel("Sensitivity")
    p.add_legend()
    p.save("foldings/sensitivity.pdf")

    p = Plotter()
    def add_window(window_tool):
        angles = np.array([0, 20, 45, 50, 60])
        arrays = np.vstack([
            window_tool.df['M0'],
            window_tool.df['M20'],
            window_tool.df['M45'],
            window_tool.df['M50'],
            window_tool.df['M60'],
        ])
        window_tool = Window(incidence_angles=angles, transmission=arrays)

        eff = CameraEfficiency.from_sstcam(
            fov_angle=0,
            pde_at_450nm=pde_at_450nm,
            window=window_tool,
        )
        label = f"NSB={eff.nominal_nsb_rate.to_value('MHz'):.2f}, PDE={eff.camera_cherenkov_pde:.2f}, S/N={eff.camera_signal_to_noise:.2f}"
        x = eff.wavelength.value
        y = eff.window_transmissivity
        p.ax.plot(x, y, alpha=0.7, label=label)

    add_window(window_tool)
    # from IPython import embed
    # embed()
    # window_tool.df.loc[window_tool.df.index < 300] = 0
    # add_window(window_tool)
    # window_tool.df.loc[window_tool.df.index > 750] = 0.2
    # add_window(window_tool)
    p.ax.set_xlabel("Wavelength [nm]")
    p.ax.set_ylabel("Sensitivity")
    p.add_legend()
    p.save("foldings/window.pdf")

if __name__ == '__main__':
    main()