#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
Created on Wed Mar  3 23:12:47 2021

@author: tim
Simulation algorthims, with parallel edits for 2+ processor case
"""
# from mpi4py import MPI
import numpy as np
import itertools as it
from comparison_fns import compare, inter_bdry, dist, which_seg, mt, mt_to_l
from parameters import v_s, L, xdomain, ydomain, tip_scaled
from event_fns import valid_events, check_region
from zippering import zip_cat
import random as rnd
from time import time
from plotting import plot_snap

def nuc_t(t):
    '''
    Next nucleation time using rate k, from Gillespie

    Parameters
    ----------
    t : Current time
    Returns
    -------
    Next nucleation time and nucleation point

    '''
    lx = xdomain[1]-xdomain[0]
    ly = ydomain[1]-ydomain[0]
    k = lx*ly*L**3*0.1/8#nucleation rate
    # k=1
    r, x, y = rnd.uniform(0,1), rnd.uniform(xdomain[0],xdomain[1]), rnd.uniform(ydomain[0],ydomain[1])
    p = [x,y] #point
    dt = np.log(1/r)/k #new time
    T = t+dt
    return(T, p)

def update_event_list(rank, MT_list,event_list,t,pevent_list,del_mt,last_result,policy,cross_data, rewind):
    '''
    Parameters
    ----------
    MT_list : List of MT objects
    pair_list : List of tuples of pairwise MT indices. Each tuple represents MTs to be compared
    t: current time
    pevet_list: list of previous events which occured
    del_mt:
    cross_data: for when MT crosses into processor
    rewind: if true, need to recalculate event_list from existing MTs

    Returns
    -------
    Collision policy
    BE CAREFUL: THE DISTANCE RETURNED IN BDRY COLLISION IS DISTANCE OF PRESENT TIP TO BDRY, NOT
    SEGMENT LENGTH!!!
    '''
    r = tip_scaled*0.5 #radius
    mt_index = [mt.number for mt in MT_list if mt.exist] #indices of active mts
    n = len(mt_index)
    l = None
    pair_list = None
    if t==0 or rewind: #calculate all possible pairs
        if rewind and t!= 0:
            event_list[:] = [x for x in event_list if x[2]=='nucleate'] #find all nucleation events
        pair_list = list(it.combinations(mt_index,2)) #list of combinations of indices of MTs
        l = len(pair_list)

        for j in range(n): #check for bdry collisions
            MT = MT_list[j]
            if MT.grow == True: #no bdry collision
                bdry_res = inter_bdry(rank,MT, MT_list) #find intersection info
                next_time = bdry_res[0] + MT.update_t[-1] #time of collision
                event_list.append([False,MT.number,bdry_res[2],next_time,bdry_res[1]])
                #regions
                event_list.append([False,MT.number,'new_region',t+r,None])#new when there's a new region
            elif MT.grow ==False and MT.hit_bdry==False: #disappearing events need to be recalculated
                mt_l = np.sum(MT.seg_dist)#total length of shrinking mt, cannot got lower than this
                disap_t = mt_l/v_s+MT.update_t[-1] #collision distances from 1 to 2, growing can only collide w/ shrinking
                event_list.append([False,MT.number,'disap',disap_t,None]) #disappearing event
        if t== 0: #only time we need to new nucleation
            th = rnd.uniform(0,2*np.pi)
            nucleate  = nuc_t(t) #nucleation event
            event_list.append([False,None,'nucleate',nucleate[0],[nucleate[1],th]])
        if policy == 'comm': #if there's been a communication REDUNDANT, ONLY IN CASE REWIND GOES TO 0
            new_policy = cross_data[0] #did an MT shrink or appear?
            mt_cross = cross_data[1] #MT of interest
            if new_policy =='disap': #shrank
                event_list.append([False,None,'shrink_in',mt_cross.disap_t,None]) #shrinking in the future
            else:
                event_list.append([False,None,'cross_in',mt_cross.update_t[-1],None]) #crossing in
    else:
        if policy == 'comm': #if there's been a communication
            new_policy = cross_data[0] #did an MT shrink or appear?
            mt_cross = cross_data[1] #MT of interest
            if new_policy =='disap': #shrank
                event_list.append([False,None,'shrink_in',mt_cross.disap_t,None]) #shrinking in the future
            else:
                event_list.append([False,None,'cross_in',mt_cross.update_t[-1],None]) #crossing in
        else:
            l_idx = mt_to_l(MT_list, last_result)
            if policy =='cross_in': #crossing in from left or right
                event_list.pop(0)
                new_idx = [last_result]
                mt_index= [p for p in mt_index if ((p != last_result) \
                and check_region(last_result,p,MT_list,t,r))] #for comparing new MT
                pair_list = list(it.product(mt_index,new_idx)) #generate new pairs
                l = len(pair_list)
                MT = MT_list[l_idx]
                bdry_res = inter_bdry(rank,MT, MT_list) #find intersection info
                next_time = bdry_res[0] + MT.update_t[-1] #time of collision
                event_list.append([False,last_result,bdry_res[2],next_time,bdry_res[1]])
                #Find region event
                event_list.append([False,last_result,'new_region',t+r,None])
            elif policy in ['left','right', 'top', 'bottom']: #these mts have entered a new processor
                '''PARALLEL EDIT'''
                new_idx = [last_result]
                mt_index= [p for p in mt_index if ((p != last_result) \
                and check_region(last_result,p,MT_list,t,r))] #for comparing new MT
                pair_list = list(it.product(mt_index,new_idx)) #generate new pairs
                l = len(pair_list)
                event_list[:] = [x for x in event_list if valid_events(x,last_result)] #discard all invalid events
                '''PARALLEL EDIT'''
            elif policy == 'deflect': #deflection
                new_idx = [last_result] #index of updated MTs
                mt_index= [p for p in mt_index if p != last_result and check_region(last_result,p,MT_list,t,r)] #indices of new pairs to check over
                n = len(mt_index)
                pair_list = list(it.product(mt_index,new_idx)) #generate new pairs
                l = len(pair_list)
                assert(last_result == event_list[0][1])
                assert(valid_events(event_list[0],last_result)==False)
                event_list[:] = [x for x in event_list if valid_events(x,last_result)] #discard all invalid events
                MT = MT_list[l_idx]
                #bdry events
                bdry_res = inter_bdry(rank,MT,MT_list) #find intersection info
                next_time = bdry_res[0] + MT.update_t[-1] #time of collision
                event_list.append([False,new_idx[0],bdry_res[2],next_time,bdry_res[1]])
                #Find region event
                event_list.append([False,last_result,'new_region',t+r,None])
            elif policy == 'disap': #if last one disappeared
                if MT_list[l_idx].from_bdry and MT_list[l_idx].origin == rank: #and it's an extension from the same rank
                    prev_idx = MT_list[l_idx].prev_mt #index of newly shrinking MT
                    event_list[:] = [x for x in event_list if valid_events(x,prev_idx)] #discard newly shrinking one
                    event_list[:] = [x for x in event_list if valid_events(x,last_result)] #discard all invalid events
                    new_idx = [prev_idx] #new list of indices
                    mt_index= [p for p in mt_index if p != prev_idx and check_region(last_result,p,MT_list,t,r)]
                    pair_list = list(it.product(mt_index,new_idx)) #new pairs
                    l = len(pair_list)
                    assert MT_list[l_idx].exist==False
                    #bdry events
                    l_idxp = mt_to_l(MT_list, prev_idx) #find microtuble
                    mt_l = np.sum(MT_list[l_idxp].seg_dist)#total length of shrinking mt, cannot got lower than this
                    disap_t = mt_l/v_s+t #collision distances from 1 to 2, growing can only collide w/ shrinking
                    event_list.append([False,prev_idx,'disap',disap_t,None])

                    del_mt.append(MT_list[l_idx]) #add deleted mt to list
                    MT_list.pop(l_idx) #remove from existent mt list
                else: #if last one disappeared, no continuation
                    event_list[:] = [x for x in event_list if valid_events(x,last_result)] #discard all invalid events
                    l = 0 #no new comparisons to be made with non-existant MT
                    assert MT_list[l_idx].exist==False
                    del_mt.append(MT_list[l_idx])
                    MT_list.pop(l_idx)
            elif policy == 'shrink_in': #shrinking case
                event_list.pop(0)
                event_list[:] = [x for x in event_list if valid_events(x,last_result)] #discard all invalid events
                new_idx = [last_result] #new list of indices
                mt_index= [p for p in mt_index if p != last_result and check_region(last_result,p,MT_list,t,r)]
                pair_list = list(it.product(mt_index,new_idx)) #new pairs
                l = len(pair_list)
                #bdry events
                mt_l = np.sum(MT_list[l_idx].seg_dist)#total length of shrinking mt, cannot got lower than this
                disap_t = mt_l/v_s+t #collision distances from 1 to 2, growing can only collide w/ shrinking
                event_list.append([False,last_result,'disap',disap_t,None])
            elif policy in ['1hit2','2hit1']:
                if MT_list[l_idx].grow is False: #if updated MT is due to cat. and it didn't disap
                    new_idx = [last_result] #index of updated MTs
                    mt_index= [p for p in mt_index if p != last_result] #indices of new pairs to check over
                    n = len(mt_index)
                    pair_list = list(it.product(mt_index,new_idx)) #generate new pairs
                    l = len(pair_list)
                    event_list[:] = [x for x in event_list if valid_events(x,last_result)] #discard all invalid events

                    mt_l = np.sum(MT_list[l_idx].seg_dist)#total length of shrinking mt, cannot got lower than this
                    disap_t = mt_l/v_s+t #collision distances from 1 to 2, growing can only collide w/ shrinking
                    event_list.append([False,last_result,'disap',disap_t,None]) #disappearing event
                else: #if two dynamics MTs collide, must be growing
                    new_idx = [last_result] #index of updated MTs
                    mt_index= [p for p in mt_index if p != last_result and check_region(last_result,p,MT_list,t,r)] #indices of new pairs to check over
                    n = len(mt_index)
                    pair_list = list(it.product(mt_index,new_idx)) #generate new pairs
                    l = len(pair_list)
                    event_list[:] = [x for x in event_list if valid_events(x,last_result)] #discard all invalid events

                    l_idx = mt_to_l(MT_list, new_idx[0])
                    MT = MT_list[l_idx]
                    #bdry event
                    bdry_res = inter_bdry(rank,MT,MT_list) #find intersection info
                    next_time = bdry_res[0] + MT.update_t[-1] #time of collision
                    event_list.append([False,new_idx[0],bdry_res[2],next_time,bdry_res[1]])
                    #Find region event
                    event_list.append([False,last_result,'new_region',t+r,None])
            elif policy == 'nucleate':#nucleation occurs
                event_list.pop(0)
                new_idx = [last_result] #index of updated MTs
                mt_index= [p for p in mt_index if check_region(last_result,p,MT_list,t,r)] #indices of new pairs to check over
                n = len(mt_index)
                pair_list = list(it.product(mt_index,new_idx)) #generate new pairs
                l = len(pair_list)
                #bdry events
                MT = MT_list[l_idx]
                bdry_res = inter_bdry(rank,MT, MT_list) #find intersection info
                next_time = bdry_res[0] + MT.update_t[-1] #time of collision
                event_list.append([False,new_idx[0],bdry_res[2],next_time,bdry_res[1]])
                #new nucleation event
                nuc_list = [x for x in event_list if x[2]=='nucleate']#check current nucleation queue
                if len(nuc_list) == 0: #if there are no mpre, calculate
                    th = rnd.uniform(0,2*np.pi)
                    nucleate  = nuc_t(t) #nucleation event
                    event_list.append([False,None,'nucleate',nucleate[0],[nucleate[1],th]])
                #Find region event
                event_list.append([False,last_result,'new_region',t+r,None])
            elif policy == 'new_region': #need to check new region
                new_idx = [last_result] #index of updated MTs
                mt_index= [p for p in mt_index if p!= last_result and check_region(last_result,p,MT_list,t,r)] #indices of new pairs to check over
                n = len(mt_index)
                pair_list = list(it.product(mt_index,new_idx)) #generate new pairs
                l = len(pair_list)
                event_list[:] = [x for x in event_list if valid_events(x,last_result)] #discard all invalid events
                #bdry events
                MT = MT_list[l_idx]
                bdry_res = inter_bdry(0,MT, MT_list) #find intersection info
                next_time = bdry_res[0] + MT.update_t[-1] #time of collision
                event_list.append([False,new_idx[0],bdry_res[2],next_time,bdry_res[1]])
                #Find region event
                event_list.append([False,last_result,'new_region',t+r,None])
            for i in range(l): #compare all pairs
                mt1_i = pair_list[i][0] #numbers of mts to compare
                mt2_i = pair_list[i][1]
                l_mt = mt_to_l(MT_list,mt1_i,mt2_i) #convert to list index
                l_mt1, l_mt2 = l_mt[0],l_mt[1]
                result = compare(MT_list[l_mt1], MT_list[l_mt2],t) #retrieve output object
                if result.policy not in ['no_collision']: #if there is collision
                    assert result.dist >=0 and MT_list[l_mt1].exist and MT_list[l_mt2].exist
                    next_time = result.dist+t
                    event_list.append([True,pair_list[i],result.policy,next_time,result.point]) #apend event details
    event_list.sort(key=lambda x: x[3]) #sort according to time
    pevent_list.append(event_list[0])#record event occurance
    return(event_list[0])

def update(rank,N,mt_list,collided,idxs,policy,dists,pt,t,cross_data):
    '''
    Updates MT based on next event
    Parameters
    ----------
    rank: rank of process
    N: number of current+past MTs in the system
    mt_list : list of MTs
    collided : bool of whether MTs collided (True) or it was a bdry event
    policy : what type of collision event
    dists : event time
    pt : point of intersection
    idx : tuple of idices if two MTs collide, single index if bdry event
    t : current time
    cross_data: for when an MT crosses into the processor

    Returns
    -------
    Index of updated MT
    Whether event is a crossing

    '''
    crossing = False #whether there has been a crossing
    if policy =='1hit2' or policy=='2hit1':
        if policy == '1hit2':
            mt_idx1,mt_idx2 = idxs[0],idxs[1] #index of MT to be updated
        else:
            mt_idx1,mt_idx2 = idxs[1],idxs[0]

        l_idx = mt_to_l(mt_list, mt_idx1,mt_idx2) #convert to list indices
        l_idx1, l_idx2 = l_idx[0],l_idx[1]

        mt1 =mt_list[l_idx1]
        mt2 =mt_list[l_idx2]

        seg_idx = which_seg(mt1,mt2,t) #find segment of intersection
        r = rnd.randint(0, 1)
        zip_res = zip_cat(mt1.angle[-1],mt2.angle[seg_idx],pt,r) #zipper or not
        new_pt,new_angle = zip_res[1], zip_res[0]
        mt1.seg.append(new_pt) #update point
        mt1.update_t.append(dists) #update time
        actual_dist = dist(mt1.seg[-2],mt1.seg[-1])
        mt1.seg_dist.append(actual_dist) #append seg distance
        if not zip_res[2]: #no catastrophe
            mt1.angle.append(new_angle) #TODO: CHANGE ANGLE UPDATE
            mt1.region_t.append(dists) #region change
        else: #catastrophe
            mt1.grow = False
        return(mt_idx1,crossing)
    # if collided is False: #if collision w/ bdry or deflection
    elif policy in ['disap']: #one has disappeared
        mt_idx1 = idxs
        l_idx1 = mt_to_l(mt_list, mt_idx1) #need to convert to indices in the list of existing mts
        assert l_idx1 >= 0
        mt1 =mt_list[l_idx1]
        mt1.exist = False #has disappeared
        mt1.disap_t = dists
        mt1.update_t.append(dists)
        '''PARALLEL EDIT START'''
        if mt1.origin != rank: #it's a continuation of another processor's mt
            assert mt1.from_bdry
            crossing = True #will need to update with crossing_update
            '''PARALLEL EDIT END'''
        elif mt1.from_bdry: #if it's an extension, must update previous MT
            mt_prev_idx = mt1.prev_mt
            l_prev_idx = mt_to_l(mt_list, mt_prev_idx) #convert to list indices

            mt_list[l_prev_idx].grow = False #it must shrink now
            mt_list[l_prev_idx].hit_bdry = False #no longer hitting bdry
            mt_list[l_prev_idx].update_t.append(dists) #update time of new bdry point now shrinking

            # print('previous index', prev_idx)
        return(mt_idx1,crossing)
    elif policy == 'shrink_in':
        mt1 = cross_data[1]
        l_prev_idx  = mt_to_l(mt_list,mt1.prev_mt)
        mt_list[l_prev_idx].grow = False #it must shrink now
        mt_list[l_prev_idx].hit_bdry = False #no longer hitting bdry
        mt_list[l_prev_idx].update_t.append(dists) #update time of new bdry point now shrinking
        return(mt1.prev_mt,crossing)
    elif policy == 'nucleate':
        mt_idx = N
        mt_list.append(mt(mt_idx,rank)) #introduce new MT
        mt1 = mt_list[-1]
        mt1.update_t[-1] = dists
        mt1.region_t.append(dists) #region change
        mt1.seg = [pt[0]] #assign points and angles
        mt1.angle = [pt[1]]
        return(mt_idx,crossing)
    elif policy == 'cross_in':
        mt1 = cross_data[1] #corresponding microtuble
        mt1.number = N #give a number
        mt_list.append(mt1)
        return(N, crossing)
    elif policy =='deflect':
        l_idx = mt_to_l(mt_list,idxs)#convert to list index

        mt1 = mt_list[l_idx] #mt to be updated
        assert l_idx>=0
        mt1.seg.append(pt) #update point
        power = rnd.randint(0,1) #to be used to determine deflection +/-
        dth  =np.pi/16 #change in angle
        th_new = mt1.angle[-1] + (-1)**power*dth #deflection
        if th_new> 2*np.pi:
            th_new -= 2*np.pi
        elif th_new <0:
            th_new += 2*np.pi
        mt1.angle.append(th_new)

        segment_dist = dist(mt1.seg[-2],mt1.seg[-1])
        mt1.seg_dist.append(segment_dist) #segment length

        mt1.update_t.append(dists) #update time
        mt1.region_t.append(dists) #region change
        return(mt1.number,crossing)
    elif policy =='new_region': #new region, no MT change
        l_idx = mt_to_l(mt_list,idxs)
        mt_list[l_idx].region_t.append(dists) #new region calculated at this time
        return(idxs,crossing)
    else: #hit bdry, nucleates new MT
        assert policy in ['top','bottom','left','right']
        l_idx = mt_to_l(mt_list, idxs) #convert to list index
        mt1 = mt_list[l_idx] #mt to be updated
        mt1.seg.append(pt) #update point
        th_og = mt1.angle[-1]
        mt1.angle.append(th_og)
        mt1.update_t.append(dists) #update time
        mt1.grow = False #no longer considered growing
        mt1.hit_bdry = True #hit bdry
        mt_idx = None
        segment_dist = dist(mt1.seg[-2],mt1.seg[-1])
        mt1.seg_dist.append(segment_dist) #segment length
        assert mt1.checkd()
        mt_idx = idxs #original MT
        crossing = True
        '''PARALLEL EDIT END'''
        return(mt_idx,crossing)

def undo_update(rank, mt_list, pevent_list, event_list, del_mt):
    '''
    Undoes updates
    Parameters
    ----------
    mt_list : list of MTs
    collided : bool of whether MTs collided (True) or it was a bdry event
    policy : what type of collision event
    dists : event time
    pt : point of intersection
    idx : tuple of idices if two MTs collide, single index if bdry event
    t : current time

    Returns
    -------
    Time of rewinded step

    '''
    event = pevent_list[-1] #previous event
    idxs, policy = event[1], event[2] #event details

    # if collided is True: #if dynamic ends collided
    assert policy not in ['shrink_in','comm','cross_in']
    if policy =='1hit2' or policy=='2hit1':
        if policy == '1hit2':
            mt_idx1,mt_idx2 = idxs[0],idxs[1] #index of MT to be updated
        else:
            mt_idx1,mt_idx2 = idxs[1],idxs[0]
        l_idx = mt_to_l(mt_list, mt_idx1,mt_idx2) #convert to list indices
        l_idx1, l_idx2 = l_idx[0],l_idx[1]
        assert l_idx1>=0 and l_idx2>=0
        mt1 =mt_list[l_idx1]

        mt1.seg.pop(-1) #un-update point
        mt1.update_t.pop(-1) #un-update time
        mt1.seg_dist.pop(-1) #append seg distance
        if mt1.grow: #no catastrophe
            mt1.angle.pop(-1) #undo angle change
            mt1.region_t.pop(-1)
        else: #catastrophe
            mt1.grow = True #undo cat
    elif policy =='disap': #one has disappeared
        mt_idx1 = idxs
        l_idx1 = mt_to_l(mt_list, mt_idx1) #need to convert to indices in the list of existing mts
        ldel_idx = mt_to_l(del_mt, mt_idx1) #might not exist, check non-existing mts
        mt1 = None
        if l_idx1 >=0: #if in this list
            mt1 =mt_list[l_idx1]
        else: #if in deleted mts
            assert ldel_idx >= 0
            mt1 = del_mt[ldel_idx]
            del_mt.pop(ldel_idx) #delete from deleted mt list
            mt_list.append(mt1) #re-introduce to mtlist
        assert mt1.number == idxs
        mt1.exist = True #has un-dissapeared
        mt1.disap_t = None #undo this
        mt1.update_t.pop(-1)
        if mt1.from_bdry and mt1.origin==rank: #if it's an extension, must un-update previous MT
            mt_prev_idx = mt1.prev_mt
            l_prev_idx = mt_to_l(mt_list, mt_prev_idx) #convert to list indices
            mt_list[l_prev_idx].grow = False #it must shrink now
            mt_list[l_prev_idx].hit_bdry = True #hitting bdry
            mt_list[l_prev_idx].update_t.pop(-1) #un-update time of new bdry point
    elif policy == 'nucleate':
        t = event[3] #in a nucleation event, we do not know the MT index ahead of time
        mt_list2 = [mt for mt in mt_list if mt.update_t[-1]==t] #must find the MT this way
        if len(mt_list2) == 0:
            print([mt.number for mt in mt_list])
            print('CANNOT FIND MT IN LIST', rank,t)
            print(event)
            assert False
        mt_idx = mt_list2[0].number #new mt number
        l_idx = mt_to_l(mt_list, mt_idx) #index in list
        assert l_idx >= 0
        mt_list.pop(l_idx) #un-introduce new MT

        event_list.append(event) #reinstate nucleation event
    elif policy =='new_region': #new region, no MT change
        l_idx = mt_to_l(mt_list,idxs)
        mt_list[l_idx].region_t.pop(-1) #new region calculated at this time
    elif policy =='deflect':
        l_idx = mt_to_l(mt_list,idxs)#convert to list index
        mt1 = mt_list[l_idx] #mt to be updated
        mt1.seg.pop(-1) #un-update point
        mt1.angle.pop(-1)
        mt1.seg_dist.pop(-1) #segment length
        mt1.update_t.pop(-1) #un-update time
        mt1.region_t.pop(-1)
    else: #hit bdry, nucleates new MT
        l_idx = mt_to_l(mt_list, idxs) #convert to list index
        mt1 = mt_list[l_idx] #mt to be updated
        mt1.seg.pop(-1) #update point
        mt1.angle.pop(-1)
        mt1.update_t.pop(-1) #un-update time
        mt1.grow = True #un-no longer considered growing
        mt1.hit_bdry = False #un-hit bdry
        mt1.seg_dist.pop(-1) #segment length
    pevent_list.pop(-1) #un-updated the latest event
    if not pevent_list:
        t_new = 0
    else:
        t_new = pevent_list[-1][3]
    return(t_new)

def exit_data(rank, policy, idx, mt_list,t):
    '''
    Given that MT w/ number idx in rank is crossing outside, prepare data to be
    sent to other process

    Parameters
    ----------
    t: current time
    rank : rank of process
    policy: left/right or disap
    idx : ID of MT in this rank
    mt_list: list of MTs

    Returns
    -------
    Data

    '''
    data = [] #data to be sent
    if policy == 'disap':
        data.append(policy)
        l_idx = mt_to_l(mt_list, idx) #convert to list index
        mt1 = mt_list[l_idx] #mt that disappered
        data.append(mt1) #index in the other processor
    else:
        assert policy in['right','left', 'top', 'bottom']
        if policy=='right': #new event in other process
            data.append('new_left')
        elif policy=='left':
            data.append('new_right')
        elif policy=='top':
            data.append('new_bottom')
        else:
            assert policy=='bottom'
            data.append('new_top')
        l_idx = mt_to_l(mt_list, idx) #convert to list index
        mt1 = mt_list[l_idx] #mt to be updated
        pt = mt1.seg[-1][:] #previous point
        # if rank == 0: #if rank 0 (left), find identified edge
        if policy == 'right':
            pt[0] = xdomain[0]
        elif policy == 'left':
            pt[0] = xdomain[1]
        elif policy == 'top':
            pt[1] = ydomain[0]
        else:
            assert policy == 'bottom'
            pt[1] = ydomain[1]
        th = mt1.angle[-1] #angle
        mt2 = mt(None,rank) #continuation mt
        mt2.angle = [th]
        mt2.seg = [pt]
        mt2.from_bdry = True
        mt2.prev_mt = idx
        mt2.update_t[-1] = t#give update time
        mt2.region_t = [t]
        if len(mt1.seg) > 2: #it changed trajectory before hitting bdry
            mt2.prev_t = mt1.update_t[-2]
        else:
            assert len(mt1.seg)==2
            if mt1.origin == rank: #can search native list for other info
                if mt1.from_bdry:
                    l_idx2 = mt_to_l(mt_list,mt1.prev_mt)#may need to consider ext of ext
                    mtp2 = mt_list[l_idx2]
                    assert len(mtp2.update_t) >= 2 #implicit assumption here that previous mt will be more segs
                        # plot_snap([mt1,mtp2],t,i,'./plots2/',False)
                    mt2.prev_t = mtp2.update_t[-2] #time before it hit wall
                else:
                    mt2.prev_t = mt1.update_t[0]
            else:
                mt2.prev_t = mt1.prev_t #last physical update same as parent MT
        data.append(mt2)
    return(data)

#testing below
if __name__ == '__main__':
    rank = 0
    rnd.seed(0)
    N = 3#number of MTs
    mt_list = []
    for i in range(N):
        mt_list.append(mt(i,rank))
        x,y = rnd.uniform(xdomain[0],xdomain[1]), rnd.uniform(ydomain[0],ydomain[1])
        th = rnd.uniform(0,2*np.pi)
        mt_list[i].seg = [[x,y]]
        mt_list[i].angle = [th]
        mt_list[i].region_t.append(0)

    t = 0
    T = 100
    start = time()
    event_list = []
    pevent_list = []#past events
    del_mt = [] #deleted MTs
    cross_data = None
    last_result = None
    policy = None
    rewind = False
    i=0
    # up_t = 0
    # mt_t = 0
    conv = L/0.08 #time conversion factor
    cross_data = None
    # while t*conv< 100:
    j = 0
    start = time()
    ti = 0
    tf = 0
    i1,i2 = 0,0
    print('Simulating on', xdomain,'x',ydomain)
    print('Simulate to time', T)
    print('Length of', L, 'microm \n')
    while t*conv< T:
        i2 = np.floor(t*conv/10)
        if i2 != i1:
            tf = time()
            print('Sim time (s): ', t*conv)
            print('Wall time elapsed (s): ',tf-ti)
            print('Length of previous event list: ', len(pevent_list))
            print('Length of event list: ', len(event_list),'\n')
            # plot_snap(rank,mt_list,t,i,'./plots2/',False)
            ti = time()
            i1 = i2
        # if i >0:
        #     plot_snap(rank,mt_list,t,i,'./plots5/',False)
        Result = update_event_list(rank, mt_list, event_list,t,pevent_list,del_mt, last_result, policy,cross_data,rewind)
        rewind= False
        policy = Result[2]
        if Result[3]*conv > 90:
            plot_snap(0, mt_list, t,i,dest='./', save=True)
        update_return = update(rank,N,mt_list,Result[0],Result[1],Result[2], Result[3], Result[4],t,cross_data)
        last_result = update_return[0]
        t1 = Result[3]
        t = Result[3]
        if update_return[1]: #if there's a crossing
            cross_data = exit_data(rank, policy, last_result, mt_list,t) #package data to be transferred
            update_event_list(rank, mt_list, event_list,t,pevent_list,del_mt, last_result, policy,cross_data,rewind) #update list with previous changed MT
            policy = 'comm' #policy will be used to add the new event to list
            j+=1
        if policy in ['nucleate', 'top','bottom','cross_in']:
            N+=1
        i+=1
    end = time()
    print('Total steps: ', i)
    print('Percent of crossings: ', j)
    print('Total time: ',end-start)


#from __future__ import print_function
"""Train  with PyTorch."""
"""
This code is forked and modified from 'https://github.com/kuangliu/pytorch-cifar'. Thanks to its contribution.
"""

import torch
import torch.nn as nn
import torch.optim as optim
import torch.nn.functional as F
import torch.backends.cudnn as cudnn
from torch_geometric.data import DataListLoader,DataLoader
import statistics
import torchvision
import torchvision.transforms as transforms
from torch_geometric.nn import GCNConv, ChebConv,global_mean_pool,SplineConv,GraphConv, AGNNConv,TopKPooling,DataParallel,GATConv
#from pyts.image import RecurrencePlot
from torch_geometric.datasets import MNISTSuperpixels,Planetoid,TUDataset,PPI,Amazon,Reddit,CoraFull
import torch_geometric.transforms as T
from torch_geometric.nn import global_mean_pool as gap, global_max_pool as gmp
import numpy as np

from sklearn import metrics
import argparse
import os,sys
global resume
global ConCoeff

def ChooseModel(model_name,datasetroot,NetInfo):
    if len(NetInfo)==1:
        width=NetInfo[0]
    elif len(NetInfo)==2:
        width=NetInfo[0]
        weights=NetInfo[1]
    else:
        raise Exception("wrong weight info")

    if model_name=="GCN":
        net=GCN(datasetroot,width)
        if len(NetInfo)==2:
            state_dict = net.state_dict()
            for i in range(len(weights)-1):
                name='layers.{}.weight'.format(i)
                if i==0:
                    state_dict[name]=weights[i]
                else:
                    state_dict[name]=weights[i][:-1,:]
            net.load_state_dict(state_dict)

    elif model_name=="GAT":
        net=GAT(datasetroot,width)
    elif model_name=="SplineNet":
        net=SplineNet(datasetroot,width)
    elif model_name=="ChebConvNet":
        net=ChebConvNet(datasetroot,width)

    elif model_name=="AGNNNet":
        net=AGNNNet(datasetroot,width)

    else:
        raise Exception("model not support, Choose GCN, or SplineNet or ChebConvNet")

    return net

def TrainPart(modelName,datasetroot,start_epoch,num_epochs,trainloader,testloader,OptimizedNet,optimizerNew,criterionNew,NumCutoff,mark,SaveModule,model_to_save):
    best_acc=1
    TrainConvergence=[]
    TestConvergence=[]
    TestAccs=[]
    for epoch in range(start_epoch,num_epochs):
        if epoch%40==0 or epoch==num_epochs-1:
            global SVDOrNot
            SVDOrNot=[NumCutoff,"{}-{}".format(mark,epoch)]
            TrainLoss=train(trainloader,OptimizedNet,optimizerNew,criterionNew)

            """if epoch==40:
                NewNetworksize,NewNetworkWeight=RetainNetworkSize(OptimizedNet,params[2])[0:2]
                NetworkInfo=[NewNetworksize[0:-1],NewNetworkWeight]
                OptimizedNet=ChooseModel(modelName,datasetroot,NetworkInfo)
                OptimizedNet = OptimizedNet.to(device)
                WeightsContate=NewNetworkWeight

            elif epoch>40:
                for layer in range(len(NewNetworkWeight)):
                    NewNetworksize,NewNetworkWeight=RetainNetworkSize(OptimizedNet,params[2])[0:2]
                    WeightsContate[layer]=torch.cat((WeightsContate[layer],NewNetworkWeight[layer]),1)
                    [U,D,V]=torch.svd(WeightsContate[layer])
                    CutoffPoint=FindCutoffPoint(D,params[3])
                    NewNetworksize.append(CutoffPoint)
                    NewNetworkWeight.append(U[:,:CutoffPoint]@V[:CutoffPoint,:CutoffPoint])
                    NetworkInfo=[NewNetworksize[0:-1],NewNetworkWeight]
                    OptimizedNet=ChooseModel(modelName,datasetroot,NetworkInfo)
                    OptimizedNet = OptimizedNet.to(device)"""

        else:
            SVDOrNot=[]
            TrainLoss=train(trainloader,OptimizedNet,optimizerNew,criterionNew)

        [TestLoss,TestAcc]=test(testloader,OptimizedNet,criterionNew)
        print('\n Epoch: {},  tain loss: {:.4f}, test loss: {:.4f},test acc: {:.4f} \n'.format(epoch,TrainLoss,TestLoss,TestAcc))
        TrainConvergence.append(TrainLoss)
        TestConvergence.append(TestLoss)
        TestAccs.append(TestAcc)
               # save model
        if SaveModule and TestAcc< best_acc:
                state = {'net': OptimizedNet.state_dict,
                                'TrainConvergence': TrainConvergence,
                                'TestAcc': TestAcc,
                                'epoch': num_epochs,
                       }
                if not os.path.isdir('checkpoint'):
                    os.mkdir('checkpoint')
                torch.save(state, model_to_save)
                best_acc = TestAcc

                ## save recurrence plots
        """if epoch%20==0:
                save_recurrencePlots_file="../Results/RecurrencePlots/RecurrencePlots_{}_{}_BatchSize{}    \_ConCoeffi{}_epoch{}.png".format(dataset, model_name,params[0],params[1],epoch)
            save_recurrencePlots(net,save_recurrencePlots_file)"""
    del OptimizedNet
    return TrainConvergence, TestConvergence,TestAccs

def ContractionLayerCoefficients(num_features,*args):
    Numlayers,alpha=args
    width=[]
    tmpOld=np.random.randint(num_features*alpha,num_features)
    for k in range(Numlayers):
        tmpNew=np.random.randint(tmpOld*alpha,tmpOld)
        width.append(tmpNew)
        tmpOld=tmpNew
    return width

def ResumeModel(model_to_save):
    # Load checkpoint.
    #print('==> Resuming from checkpoint..')
    checkpoint = torch.load(model_to_save)
    net = checkpoint['net']
    TrainConvergence = checkpoint['TrainConvergence']
    Acc = checkpoint['TestAcc']
    start_epoch = checkpoint['epoch']
    return net,TrainConvergence,Acc,start_epoch

def FindCutoffPoint(DiagValues,coefficient):
    for i in range(DiagValues.shape[0]-1):
        if DiagValues[i]>DiagValues[i+1]*coefficient:
            CutoffPoint=i+1

    try:
        return CutoffPoint
    except:
        return DiagValues.shape[0]

"""def init_weights(m):
    if type(m) == nn.Linear:
        torch.nn.init.orthogonal_(m.weight)
        m.bias.data.fill_(0.01)"""

def weights_init(m):
    classname = m.__class__.__name__
    if classname.find('GCNConv') != -1:
        torch.nn.init.orthogonal_(m.weight.data)
        #torch.nn.init.xavier()
        m.bias.data.fill_(0.01)

def weight_reset(m):
    if isinstance(m, GCNConv) :
        m.reset_parameters()

def logging(message):
    global  print_to_logging
    if print_to_logging:
        print(message)
    else:
        pass

def print_nvidia_useage():
    global print_device_useage
    if print_device_useage:
        os.system('\n echo check gpu;nvidia-smi;echo check done \r')
    else:
        pass

def save_recurrencePlots(net,save_recurrencePlots_file):
    global save_recurrence_plots
    if save_recurrence_plots:
        for name,parameters in net.named_parameters():
            if "fc" in name and parameters.cpu().detach().numpy().ndim==2:
                hiddenState=parameters.cpu().detach().numpy()
                rp = RecurrencePlot()
                X_rp = rp.fit_transform(hiddenState)
                plt.figure(figsize=(6, 6))
                plt.imshow(X_rp[0], cmap='binary', origin='lower')
                plt.savefig(save_recurrencePlots_file,dpi=600)
            else:
                continue
    else:
        pass

      #net.apply(weight_reset)

class GCN(torch.nn.Module):
    def __init__(self,datasetroot,width):
        super(GCN,self).__init__()
        self.NumLayers=len(width)
        self.layers = nn.ModuleList()
        self.layers.append(GCNConv(datasetroot.num_features, width[0]))
        for i in range(self.NumLayers-1):
            layer=GCNConv(width[i],width[i+1])
            nn.init.xavier_uniform_(layer.weight)
            self.layers.append(layer)
        self.layers.append(GCNConv(width[-1], datasetroot.num_classes))

    def forward(self,data):
        x, edge_index,batch = data.x, data.edge_index,data.batch
        DiagElemnt=[]
        for layer in self.layers[:-1]:
            if len(SVDOrNot)==2:
                NumCutoff=SVDOrNot[0]
                mark=SVDOrNot[1]
                [U,D,V]=torch.svd(x)
                DiagElemnt.append(D[0:NumCutoff].detach().tolist())
                np.save(mark,DiagElemnt)
            x=layer(x, edge_index)
            x =x*torch.sigmoid(x)
        #x = F.dropout(x, training=self.training)
        x=self.layers[-1](x,edge_index)
        x = torch.cat([gmp(x, batch), gap(x, batch)], dim=1)
        x = F.log_softmax(x,dim=1)
        return x

class GAT(torch.nn.Module):
    def __init__(self,datasetroot,width):
        super(GAT, self).__init__()
        self.NumLayers=len(width)
        self.layers = nn.ModuleList()
        self.layers.append(GATConv(datasetroot.num_features, width[0]))
        for i in range(self.NumLayers-1):
            layer=GATConv(width[i],width[i+1])
            nn.init.xavier_uniform_(layer.weight)
            self.layers.append(layer)
        self.layers.append(GATConv(width[-1], datasetroot.num_classes))

    def forward(self,data):
        x, edge_index,batch = data.x, data.edge_index,data.batch
        DiagElemnt=[]
        for layer in self.layers[:-1]:
            if len(SVDOrNot)==2:
                NumCutoff=SVDOrNot[0]
                mark=SVDOrNot[1]
                [U,D,V]=torch.svd(x)
                DiagElemnt.append(D[0:NumCutoff].detach().tolist())
                np.save(mark,DiagElemnt)
            x=layer(x, edge_index)
            x =x*torch.sigmoid(x)
        #x = F.dropout(x, training=self.training)
        x=self.layers[-1](x,edge_index)
        x = torch.cat([gmp(x, batch), gap(x, batch)], dim=1)
        x = F.log_softmax(x,dim=1)
        return x

class AGNNNet(torch.nn.Module):
    def __init__(self,datasetroot,width):
        super(AGNNNet, self).__init__()
        self.NumLayers=len(width)
        self.layers = nn.ModuleList()
        self.lin1=torch.nn.Linear(datasetroot.num_features, width[0])
        for i in range(self.NumLayers-1):
            self.layers.append(torch.nn.Linear(width[i],width[i+1]))
        self.layers.append(torch.nn.Linear(width[-1], datasetroot.num_classes))
        self.prop1 = AGNNConv(requires_grad=True)
        self.prop2 = AGNNConv(requires_grad=True)
        self.layers = nn.ModuleList()

    def forward(self,data):
        x,edge_index=data.x,data.edge_index
        x = F.dropout(x, training=self.training)
        x=self.lin1(x)
        x = self.prop1(x, edge_index)
        x = self.prop2(x, edge_index)
        DiagElemnt=[]
        for layer in self.layers[:-1]:

            x=layer(x)
            x =x*torch.sigmoid(x)
            x = F.dropout(x, training=self.training)
        return F.log_softmax(x, dim=1)

class ChebConvNet(torch.nn.Module):
    def __init__(self,datasetroot,width):
        super(ChebConvNet,self).__init__()
        self.NumLayers=len(width)
        self.layers = nn.ModuleList()
        self.layers.append(ChebConv(datasetroot.num_features,width[0],K=1))
        for i in range(self.NumLayers-1):
            layer=ChebConv(width[i],width[i+1],K=1)
            nn.init.xavier_uniform_(layer.weight)
            self.layers.append(layer)
        self.layers.append(ChebConv(width[-1], datasetroot.num_classes,K=1))

    def forward(self,data):
        x, edge_index,batch = data.x, data.edge_index,data.batch
        DiagElemnt=[]
        for layer in self.layers[:-1]:
            if len(SVDOrNot)==2:
                NumCutoff=SVDOrNot[0]
                mark=SVDOrNot[1]
                [U,D,V]=torch.svd(x)
                DiagElemnt.append(D[0:NumCutoff].detach().tolist())
                np.save(mark,DiagElemnt)
            x=layer(x,edge_index)
            x =x*torch.sigmoid(x)
        x=self.layers[-1](x,edge_index)
        x = torch.cat([gmp(x, batch), gap(x, batch)], dim=1)
        x = F.log_softmax(x,dim=1)
        return x

class SplineNet(torch.nn.Module):
    def __init__(self,datasetroot, width):
        super(SplineNet, self).__init__()
        self.NumLayers=len(width)
        self.layers = nn.ModuleList()
        self.layers.append(SplineConv(datasetroot.num_features, width[0], dim=1, kernel_size=2))
        for i in range(self.NumLayers-1):
            layer=SplineConv(width[i],width[i+1], dim=1, kernel_size=2)
            nn.init.xavier_uniform_(layer.weight)
            self.layers.append(layer)
        self.layers.append(SplineConv(width[-1], datasetroot.num_classes, dim=1, kernel_size=2))

    def forward(self,data):
        x, edge_index,edge_attr,batch = data.x, data.edge_index,data.edge_attr,data.batch
        DiagElemnt=[]
        for layer in self.layers[:-1]:
            if len(SVDOrNot)==2:
                NumCutoff=SVDOrNot[0]
                mark=SVDOrNot[1]
                [U,D,V]=torch.svd(x)
                DiagElemnt.append(D[0:NumCutoff].detach().tolist())
                np.save(mark,DiagElemnt)
            x=layer(x, edge_index,edge_attr)
            x =x*torch.sigmoid(x)
        x=self.layers[-1](x,edge_index)
        x = torch.cat([gmp(x, batch), gap(x, batch)], dim=1)
        x = F.log_softmax(x,dim=1)
        return x

class topk_pool_Net(torch.nn.Module):
    def __init__(self,datasetroot, width):
        super(topk_pool_Net, self).__init__()
        self.NumLayers=len(width)
        self.conv1 = GraphConv(datasetroot.num_features, 128)
        self.pool1 = TopKPooling(128, ratio=0.8)
        self.conv2 = GraphConv(128, 128)
        self.pool2 = TopKPooling(128, ratio=0.8)
        self.conv3 = GraphConv(128, 128)
        self.pool3 = TopKPooling(128, ratio=0.8)
        self.layers = nn.ModuleList()
        self.layers.append(torch.nn.Linear(256,width[0]))
        for i in range(self.NumLayers-1):
            self.layers.append(torch.nn.Linear(width[i],width[i+1]))
        self.layers.append(torch.nn.Linear(width[-1], datasetroot.num_classes))

    def forward(self, data):
        x, edge_index, batch = data.x, data.edge_index, data.batch

        x = F.relu(self.conv1(x, edge_index))
        x, edge_index, _, batch, _, _ = self.pool1(x, edge_index, None, batch)
        x1 = torch.cat([gmp(x, batch), gap(x, batch)], dim=1)

        x = F.relu(self.conv2(x, edge_index))
        x, edge_index, _, batch, _, _ = self.pool2(x, edge_index, None, batch)
        x2 = torch.cat([gmp(x, batch), gap(x, batch)], dim=1)

        x = F.relu(self.conv3(x, edge_index))
        x, edge_index, _, batch, _, _ = self.pool3(x, edge_index, None, batch)
        x3 = torch.cat([gmp(x, batch), gap(x, batch)], dim=1)

        x = x1 + x2 + x3
        for layer in self.layers[:-1]:
            x=layer(x)
            x =x*torch.sigmoid(x)
            x = F.dropout(x, p=0.5, training=self.training)
        x = F.log_softmax(self.layers[-1](x),dim=1)
        return x

def ModelAndSave(dataset,modelName,train_dataset,params,num_epochs):
    model_to_save='./checkpoint/{}-{}-param_{}_{}_{}_{}-ckpt.pth'.format(dataset,modelName,params[0],params[1],params[2],params[3])
    if resume=="True" and os.path.exists(model_to_save):
        [net,TrainConvergence,Acc,start_epoch]=ResumeModel(model_to_save)
        if start_epoch>=num_epochs-1:
            pass

    else:
        width=ContractionLayerCoefficients(train_dataset.num_features,*params[1:3])
        NetworkInfo=[width]
        net =ChooseModel(modelName,train_dataset,NetworkInfo)
    return net, model_to_save

# Training

def RetainNetworkSize(net,ConCoeff):
    NewNetworksize=[]
    NewNetworkWeight=[]
    for layer_name, Weight in net.named_parameters():
        if ("weight" in layer_name) and ("layers" in layer_name):
            #print(layer_name)
            if Weight.dim()==3:
                Weight=Weight[0].view(Weight.size()[1],Weight.size()[2])
            [U,D,V]=torch.svd(Weight)
            #NumCutoff=SVDOrNot[0]
            #WeightsDynamics.append(D[:NumCutoff].tolist())
            CutoffPoint=FindCutoffPoint(D,ConCoeff)
            NewNetworksize.append(CutoffPoint)
            NewNetworkWeight.append(U[:,:CutoffPoint]@V[:CutoffPoint,:CutoffPoint])
            #NewNetworkWeight.append(U[:,:CutoffPoint]@torch.diag(D[:CutoffPoint])@V[:CutoffPoint,:CutoffPoint])

            """NewWeight= torch.mm(Weight,V[:,:CutoffPoint])
            net.conv2.weight = torch.nn.Parameter( NewWeight)"""
            #print("Original size is {},After SVD is {}".format(Weight.shape[1],CutoffPoint))
    return NewNetworksize,NewNetworkWeight

def train(trainloader,net,optimizer,criterion):
    net.train()
    train_loss = 0
    optimizer.zero_grad()
    for data in trainloader:
        data=data.to(device)
        optimizer.zero_grad()
        output=net(data)
        loss = criterion(output, data.y)
        loss.backward()
        train_loss+=data.num_graphs * loss.item()
        optimizer.step()

        """for layer_name, parameters in net.named_parameters():
             if "weight" in layer_name:
                net.conv.weight.div_(torch.norm(net.conv.weight, dim=2, keepdim=True)"""
        #test(testloader,net,criterion)
    return train_loss/len(train_dataset)

def test(testloader,net,criterion):
    net.eval()
    test_loss, accs= 0,0
    with torch.no_grad():
        for data in testloader:
            data=data.to(device)
            output= net(data)
            pred= output.max(1)[1]
            loss = criterion(output,data.y)
            test_loss+=loss.item()*data.num_graphs
            accs+=pred.eq(data.y.to(pred.device)).sum().item()

    #print('\n Test set: Average loss: {:.4f} \n'.format(test_loss[-1]))
    return test_loss/len(test_dataset),accs/len(test_dataset)

def TrainingNet(dataset,modelName,params,num_pre_epochs,num_epochs,NumCutoff,optimizerName,MonteSize,savepath):
    Batch_size=int(params[0])
    root='/git/data/GraphData/'+dataset
    TestAccs=[]
    for Monte_iter in range(MonteSize):
        # Data
        start_epoch = 0  # start from epoch 0 or last checkpoint epoch
        NewNetworkSizeAdjust=[]
        WeightsDynamicsEvolution=[]
        # model
        if dataset=='Cora' or dataset =='Citeseer' or dataset =='Pubmed':
            datasetroot= Planetoid(root=root, name=dataset, transform =T.NormalizeFeatures()).shuffle()
            trainloader = DataListLoader(datasetroot, batch_size=Batch_size, shuffle=True)
            [net,model_to_save]=ModelAndSave(dataset,modelName,datasetroot,params,num_epochs)
            criterion = nn.CrossEntropyLoss()

        elif dataset =="CoraFull":
            datasetroot = CoraFull(root=root,transform =T.NormalizeFeatures()).shuffle()
            trainloader = DataListLoader(datasetroot, batch_size=Batch_size, shuffle=True)
            [net,model_to_save]=ModelAndSave(dataset,modelName,datasetroot,params,num_epochs)

        elif dataset=="Amazon":
            datasetroot=Amazon(root, "Photo", transform=None, pre_transform=None)
            trainloader = DataListLoader(datasetroot, batch_size=Batch_size, shuffle=True)
            testloader = DataListLoader(datasetroot, batch_size=100, shuffle=False)
            [net,model_to_save]=ModelAndSave(dataset,modelName,datasetroot,params,num_epochs)

        elif dataset=='ENZYMES' or dataset=='MUTAG':
            datasetroot=TUDataset(root,name=dataset,use_node_attr=True)
            Num=len(datasetroot)// 10
            global train_dataset, test_dataset
            train_dataset = datasetroot[:Num]
            test_dataset = datasetroot[Num:]
            trainloader = DataLoader(train_dataset, batch_size=Batch_size)
            testloader = DataLoader(test_dataset, batch_size=60)
            [net,model_to_save]=ModelAndSave(dataset,modelName,datasetroot,params,num_epochs)

        elif dataset =="PPI":
            train_dataset = PPI(root, split='train')
            test_dataset = PPI(root, split='test')
            trainloader = DataLoader(train_dataset, batch_size=Batch_size, shuffle=True)
            testloader = DataLoader(test_dataset, batch_size=1, shuffle=False)
            [net,model_to_save]=ModelAndSave(dataset,modelName,train_dataset,params,num_epochs)
            criterion = torch.nn.BCEWithLogitsLoss()

        elif dataset =="Reddit":
            datasetroot=Reddit(root)
            trainloader = DataListLoader(datasetroot, batch_size=1, shuffle=True)
            testloader = DataListLoader(datasetroot, batch_size=2, shuffle=False)
            [net,model_to_save]=ModelAndSave(dataset,modelName,datasetroot,params,num_epochs)
            criterion = torch.nn.BCEWithLogitsLoss()

        elif dataset=='MNIST':
            datasetroot = MNISTSuperpixels(root=root, transform=T.Cartesian())
            trainloader = DataListLoader(datasetroot, batch_size=Batch_size, shuffle=True)
            testloader = DataListLoader(datasetroot, batch_size=100, shuffle=False)
            [net,model_to_save]=ModelAndSave(dataset,modelName,datasetroot,params,num_epochs)

        elif dataset=='CIFAR10':
            pass
        else:
            raise Exception("Input wrong datatset!!")

        FileName="{}-{}-param_{}_{}_{}_{}-monte_{}".format(dataset,modelName,params[0],params[1],params[2],params[3],Monte_iter)

        print('Let\'s use', torch.cuda.device_count(), 'GPUs!')
        global device
        device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
        optimizer =optim.Adam(net.parameters(), lr=params[3], betas=(0.9, 0.999), eps=1e-08, weight_decay=0, amsgrad=False)
        criterion = nn.CrossEntropyLoss()
        net = net.to(device)

                #cudnn.benchmark = True
        logging('dataset:{}, Batch size: {}, Number of layers:{} ConCoeff: {}, LR:{}, MonteSize:{}'.format(dataset, params[0], params[1],params[2],params[3],Monte_iter))
        mark="{}/{}Convergence/DiagElement-{}".format(savepath,dataset,FileName)

        PreTrainConvergence,PreTestConvergence,PreTestAcc=TrainPart(modelName,datasetroot,start_epoch,num_pre_epochs,trainloader,testloader,net,optimizer,criterion,NumCutoff,mark,False,model_to_save)
        print('dataset: {}, model name: {}, Number epoches: {},  Pre-train error is: {}, Pre-test error is: {}, test acc is {}'.format(dataset,modelName,num_pre_epochs,PreTrainConvergence[-1],PreTestConvergence[-1],PreTestAcc[-1]))

        NewNetworksize,NewNetworkWeight=RetainNetworkSize(net,params[2])[0:2]
        NetworkInfo=[NewNetworksize[0:-1],NewNetworkWeight]
        OptimizedNet=ChooseModel(modelName,datasetroot,NetworkInfo)
        NewNetworksize.insert(0,datasetroot.num_features)
        NewNetworkSizeAdjust.append(NewNetworksize[0:-1])
        print(NewNetworkSizeAdjust)

        #OptimizedNet.apply(init_weights)
        #OptimizedNet = DataParallel(OptimizedNet)
        OptimizedNet = OptimizedNet.to(device)
        cudnn.benchmark = True
        criterionNew = nn.CrossEntropyLoss()
        if optimizerName =="SGD":
            optimizerNew = getattr(optim,optimizerName)(OptimizedNet.parameters(), lr=params[3], momentum=0.9, weight_decay=5e-4)
        elif optimizerName =="Adam":
            optimizerNew = getattr(optim,optimizerName)(OptimizedNet.parameters(), lr=params[3], betas=(0.9, 0.999), eps=1e-08, weight_decay=5e-4, amsgrad=False)

        TrainConvergence,TestConvergence,TestAcc=TrainPart(modelName,datasetroot,start_epoch,num_epochs,trainloader,testloader,OptimizedNet,optimizerNew,criterionNew,NumCutoff,mark,True,model_to_save)
        np.save("{}/{}Convergence/TrainConvergence-{}".format(savepath,dataset,FileName),TrainConvergence)
        np.save("{}/{}Convergence/TestConvergence-{}".format(savepath,dataset,FileName),TestConvergence)

        np.save("{}/{}Convergence/NewNetworkSizeAdjust-{}".format(savepath,dataset,FileName),NewNetworkSizeAdjust)

        #np.save(savepath+'TestConvergence-'+FileName,TestConvergence)
        #torch.cuda.empty_cache()

        print('dataset: {}, model name:{}, resized network size is {},  Number epoches:{},  Train error is: {}, Test error is: {}, test acc is {}\n'.format(dataset,modelName,NewNetworksize[0:-1],num_epochs,TrainConvergence[-1],TestConvergence[-1],TestAcc[-1]))
        TestAccs.append(TestAcc)
        np.save("{}/{}Convergence/MeanTestAccs-{}".format(savepath,dataset,FileName),TestAccs)
    print("The change of test error is:{}".format(TestAccs))
    print_nvidia_useage()

if __name__=="__main__":
    parser = argparse.ArgumentParser(description='PyTorch Training')
    parser.add_argument('--dataset',default='ENZYMES',type=str, help='dataset to train')
    parser.add_argument('--modelName',default='GCN',type=str, help='model to use')
    parser.add_argument('--LR', default=0.2, type=float, help='learning rate')
    parser.add_argument('--ConCoeff', default=0.99, type=float, help='contraction coefficients')
    parser.add_argument('--CutoffCoeff', default=0.1, type=float, help='contraction coefficients')
    parser.add_argument('--NumCutoff', default=5, type=float, help='contraction coefficients')
    parser.add_argument('--rho', type=float, default=1e-2, metavar='R',
                        help='cardinality weight (default: 1e-2)')
    parser.add_argument('--optimizer',default='SGD',type=str, help='optimizer to train')

    parser.add_argument('--num_pre_epochs', type=int, default=30, metavar='P',
                        help='number of epochs to pretrain (default: 3)')
    parser.add_argument('--num_epochs', type=int, default=100, metavar='N',
                        help='number of epochs to train (default: 10)')
    parser.add_argument('--MonteSize', default=1, type=int, help=' Monte Carlos size')
    parser.add_argument('--gpus', default="0", type=str, help="gpu devices")
    parser.add_argument('--BatchSize', default=512, type=int, help='batch size')
    parser.add_argument('--NumLayers', default=2, type=int, help='Number of layers')
    parser.add_argument('--PruningTimes', default=2, type=int, help='Pruning times')
    parser.add_argument('--savepath', type=str, default='Results', help='Path to save results')
    parser.add_argument('--return_output', type=str, default=False, help='Whether output')
    parser.add_argument('--resume', '-r', type=str,default=False, help='resume from checkpoint')
    parser.add_argument('--print_device_useage', type=str, default=False, help='Whether print gpu useage')
    parser.add_argument('--print_to_logging', type=str, default=True, help='Whether print')
    parser.add_argument('--save_recurrence_plots', type=str, default=False, help='Whether print')

    args = parser.parse_args()
    os.environ['CUDA_VISIBLE_DEVICES'] = args.gpus
    print_to_logging=args.print_to_logging
    print_device_useage=args.print_device_useage
    return_output=args.return_output
    resume=args.resume
    save_recurrence_plots=args.save_recurrence_plots
    #params=[args.BatchSize,args.NumLayers,args.args.ConCoeff,args.CutoffCoeff]
    params=[args.BatchSize,args.NumLayers,args.ConCoeff,args.LR]

    TrainingNet(args.dataset,args.modelName,params,args.num_pre_epochs,args.num_epochs,args.NumCutoff,args.optimizer,args.MonteSize,args.savepath)


#!/usr/bin/env python

import roslib
import rospy
import time
import math
import cv2
from sensor_msgs.msg import Image
from cv_bridge import CvBridge
import numpy as np

PATH = rospy.get_param("path")
PAUSE = rospy.get_param("pause")

class NodeRawImage():
	def __init__(self):
		self.seq = 1
		rospy.init_node('image_raw_publisher')
		self.publisher = rospy.Publisher('/image_raw', Image, queue_size = 1)
		self.bridge = CvBridge()

	def waitInput(self):
		if PAUSE:
			t = raw_input()
			if t == 'q':
				self.seq -= 1
			elif t == 'qq':
				self.seq -= 2
		self.publish()

	def publish(self):
		img = cv2.imread(PATH + "images/" + str(self.seq) + ".jpg", cv2.CV_LOAD_IMAGE_COLOR)
		# msg = self.bridge.cv2_to_imgmsg(img, encoding="passthrough")
		msg = self.bridge.cv2_to_imgmsg(img, encoding="bgr8")
		msg.header.stamp = rospy.Time.now()
		msg.header.frame_id = 'image_raw_publisher'
		msg.header.seq = self.seq
		self.seq += 1
		self.publisher.publish(msg)

if __name__ == '__main__':
	try:
		n = NodeRawImage()
		#rospy.spin()
		while not rospy.is_shutdown():
			n.waitInput()
	except rospy.ROSInterruptException:
		pass


import gym

from ray.rllib.agents.dqn import DQNTrainer

from smarts.core.controllers import ActionSpaceType
from smarts.core.agent_interface import RGB, Waypoints, NeighborhoodVehicles

from utils import common

NAME = "DQN"
TRAINER = DQNTrainer

def make_config(**kwargs):
    use_stacked_observation = kwargs.get("use_stacked_observation", False)
    use_rgb = kwargs.get("use_rgb", False)
    closest_neighbor_num = kwargs.get("max_observed_neighbors", 8)

    img_resolution = 24
    observe_lane_num = 3
    look_ahead = 10
    stack_size = 3 if use_stacked_observation else 1
    subscribed_features = dict(
        goal_relative_pos=(stack_size, 2),
        distance_to_center=(stack_size, 1),
        speed=(stack_size, 1),
        steering=(stack_size, 1),
        heading_errors=(stack_size, look_ahead),
        neighbor=(stack_size, closest_neighbor_num * 5),  # dist, speed, ttc
        img_gray=(stack_size, img_resolution, img_resolution) if use_rgb else False,
    )

    action_space = common.ActionSpace.from_type(1)
    observation_space = gym.spaces.Dict(
        common.subscribe_features(**subscribed_features)
    )

    policy_config = (
        None,
        observation_space,
        action_space,
        dict(model=dict(custom_model_config=dict(obs_space_dict=observation_space),)),
    )

    interface_config = dict(
        obs_stack_size=stack_size,
        neighborhood_vehicles=NeighborhoodVehicles(radius=100),
        waypoints=Waypoints(lookahead=10),
        rgb=RGB(width=256, height=256, resolution=img_resolution / 256)
        if use_rgb
        else None,
        action=ActionSpaceType.Lane,
        road_waypoints=None,
        drivable_area_grid_map=None,
        ogm=None,
        lidar=None,
        debug=False,
    )

    observation_adapter = common.get_observation_adapter(
        observation_space,
        look_ahead=look_ahead,
        observe_lane_num=observe_lane_num,
        resize=(img_resolution, img_resolution),
        closest_neighbor_num=closest_neighbor_num,
    )

    agent_config = dict(
        observation_adapter=observation_adapter,
        reward_adapter=get_reward_adapter(observation_adapter),
        action_adapter=common.ActionAdapter.from_type(1),
        info_adapter=common.default_info_adapter,
    )

    learning_config = dict()

    other_config = dict(
        stop={"time_total_s": 2 * 60 * 60},
        checkpoint_freq=40,
        checkpoint_at_end=True,
        max_failures=1000,
    )

    return common.Config(
        name=NAME,
        agent=agent_config,
        interface=interface_config,
        policy=policy_config,
        learning=learning_config,
        other=other_config,
        trainer=DQNTrainer,
        spec={"obs": observation_space, "act": action_space},
    )

def get_reward_adapter(observation_adapter, **kwargs):
    def reward_adapter(env_obs, env_reward):
        return 0.05 * env_reward

    return reward_adapter


import re, datetime

from django.forms import ModelForm, HiddenInput, BaseModelFormSet
from django.contrib.auth.forms import UserCreationForm
from django.contrib.auth.models import User
from django import forms

from accounting.models import Product, Order
from .models import Company, CompanyAddress, CompanyCustomer, CompanyCustomerResponsible, CompanyContact, CompanyCustomerComment, Group

class CompanyForm(ModelForm):
	class Meta:
		model = Company
		fields = ['name', 'identity_number', 'website', 'type', 'ths_customer_id', 'invoice_name', 'invoice_address_line_1', 'invoice_address_line_2', 'invoice_address_line_3', 'invoice_zip_code', 'invoice_city', 'invoice_country', 'invoice_reference', 'invoice_email_address']

class CompanyAddressForm(ModelForm):
	class Meta:
		model = CompanyAddress
		fields = "__all__"

class GroupForm(ModelForm):
	class Meta:
		model = Group
		fields = "__all__"

	def is_valid(self, group):
		valid = super(GroupForm, self).is_valid()

		if not valid:
			return valid

		parent = self.cleaned_data.get("parent")
		allow_companies = self.cleaned_data.get("allow_companies")
		allow_registration = self.cleaned_data.get("allow_registration")

		if parent is not None and parent == group:
			self.add_error("parent", "The group cannot be its own parent.")
			valid = False

		if not allow_companies and allow_registration:
			self.add_error("allow_registration", "Companies must be allowed in order to allow registration.")
			valid = False

		return valid

	def __init__(self, fair, *args, **kwargs):
		super(GroupForm, self).__init__(*args, **kwargs)

		self.initial["fair"] = fair.id
		self.fields["fair"].disabled = True
		self.fields["fair"].widget = HiddenInput()

class BaseCompanyAddressFormSet(BaseModelFormSet):
	def clean(self):
		if any(self.errors):
			return

		for form in self.forms:
			pass

class BaseCompanyContactFormSet(BaseModelFormSet):
	def clean(self):
		if any(self.errors):
			return

		for form in self.forms:
			pass

class CompanyCustomerResponsibleForm(ModelForm):
	group = forms.ModelChoiceField(queryset = Group.objects.all(), widget = forms.RadioSelect(), required = True)
	users = forms.ModelMultipleChoiceField(queryset = User.objects.all(), widget = forms.CheckboxSelectMultiple(), required = True)

	def __init__(self, company, *args, **kwargs):
		super(CompanyCustomerResponsibleForm, self).__init__(*args, **kwargs)
		self.initial["company"] = company.id
		self.fields["company"].disabled = True
		self.fields["company"].widget = HiddenInput()

	class Meta:
		model = CompanyCustomerResponsible
		fields = "__all__"

	def is_valid(self):
		valid = super(CompanyCustomerResponsibleForm, self).is_valid()

		if not valid:
			return valid

		group = self.cleaned_data.get("group")

		if not group.allow_responsibilities:
			self.add_error("group", "This group cannot be used for responsibilities.")
			valid = False

		return valid

class CompanyCustomerCommentForm(ModelForm):
	# TODO: this needs to be Fair-specific
	groups = forms.ModelMultipleChoiceField(queryset = Group.objects.filter(allow_comments = True), widget = forms.CheckboxSelectMultiple(), required = False)

	class Meta:
		model = CompanyCustomerComment
		fields = ("groups", "comment",)

def fix_phone_number(n):
	if n is None: return None

	n = n.replace(' ', '')
	n = n.replace('-', '')

	if n.startswith("00"): n = "+" + n[2:]
	if n.startswith("0"): n = "+46" + n[1:]

	return n

class CompanyContactForm(ModelForm):
	def clean(self):
		super(CompanyContactForm, self).clean()

		if "mobile_phone_number" in self.cleaned_data:
			self.cleaned_data["mobile_phone_number"] = fix_phone_number(self.cleaned_data["mobile_phone_number"])

		if "work_phone_number" in self.cleaned_data:
			self.cleaned_data["work_phone_number"] = fix_phone_number(self.cleaned_data["work_phone_number"])

		if "email_address" in self.cleaned_data and self.cleaned_data["email_address"] is not None:
			self.cleaned_data["email_address"] = self.cleaned_data["email_address"].lower()

		return self.cleaned_data

	def is_valid(self):
		valid = super(CompanyContactForm, self).is_valid()

		if not valid:
			return valid

		mobile_phone_number = self.cleaned_data.get("mobile_phone_number")
		work_phone_number = self.cleaned_data.get("work_phone_number")

		if mobile_phone_number is not None and not re.match(r'\+[0-9]+$', mobile_phone_number):
			self.add_error("mobile_phone_number", "Must only contain numbers and a leading plus.")
			valid = False

		if work_phone_number is not None and not re.match(r'\+[0-9]+$', work_phone_number):
			self.add_error("work_phone_number", "Must only contain numbers and a leading plus.")
			valid = False

		return valid

	class Meta:
		model = CompanyContact
		fields = "__all__"
		exclude = ("user", "company",)

class UserForm(UserCreationForm):
	class Meta:
		model = User
		fields = ["password1", "password2"]

class CreateCompanyContactForm(ModelForm):
	def clean(self):
		super(CreateCompanyContactForm, self).clean()

		if "mobile_phone_number" in self.cleaned_data:
			self.cleaned_data["mobile_phone_number"] = fix_phone_number(self.cleaned_data["mobile_phone_number"])

		if "work_phone_number" in self.cleaned_data:
			self.cleaned_data["work_phone_number"] = fix_phone_number(self.cleaned_data["work_phone_number"])

		if "email_address" in self.cleaned_data and self.cleaned_data["email_address"] is not None:
			self.cleaned_data["email_address"] = self.cleaned_data["email_address"].lower()

		return self.cleaned_data

	def __init__(self, *args, **kwargs):
		super(CreateCompanyContactForm, self).__init__(*args, **kwargs)
		self.fields["company"].label = "Company"

	class Meta:
		model = CompanyContact
		fields = "__all__"
		exclude = ("user", "active","confirmed")

	def is_valid(self):
		valid = super(CreateCompanyContactForm, self).is_valid()

		if not valid:
			return valid

		email_address = self.cleaned_data["email_address"]
		mobile_phone_number = self.cleaned_data.get("mobile_phone_number")
		work_phone_number = self.cleaned_data.get("work_phone_number")

		if User.objects.filter(username = email_address).first() != None:
			self.add_error("email_address", "Account already exists")
			return False

		if mobile_phone_number is not None and not re.match(r'\+[0-9]+$', mobile_phone_number):
			self.add_error("mobile_phone_number", "Must only contain numbers and a leading plus.")
			valid = False

		if work_phone_number is not None and not re.match(r'\+[0-9]+$', work_phone_number):
			self.add_error("work_phone_number", "Must only contain numbers and a leading plus.")
			valid = False

		return valid

class CreateCompanyContactNoCompanyForm(CreateCompanyContactForm):
	def __init__(self, *args, **kwargs):
		super(CreateCompanyContactForm, self).__init__(*args, **kwargs)

	class Meta:
		model = CompanyContact
		fields = "__all__"
		exclude = ("user", "active", "confirmed", "company")

class CompanyCustomerStatusForm(forms.Form):
	status = forms.ModelChoiceField(queryset = Group.objects.all(), required = False, label = "Status")

class DateInput(forms.DateInput):
	input_type = 'date'

class StatisticsForm(forms.Form):
	date_from = forms.DateField(widget = DateInput(), )
	date_to = forms.DateField(initial = datetime.date.today, widget = DateInput())

class CompanyNewOrderForm(ModelForm):
	class Meta:
		model = Order
		fields = ('product', 'quantity')

	def is_valid(self):
		valid = super(CompanyNewOrderForm, self).is_valid()

		if not valid:
			return valid

		product = self.cleaned_data['product']
		quantity = self.cleaned_data['quantity']

		if quantity < 1:
			self.add_error('quantity', 'Must be a positive integer.')
			return False

		if product.max_quantity is not None and quantity > product.max_quantity:
			self.add_error('quantity', 'Must not exceed ' + str(product.max_quantity) + '.')
			return False

		return valid

class CompanyEditOrderForm(ModelForm):
	class Meta:
		model = Order
		fields = '__all__'
		exclude = ('product', 'purchasing_company', 'purchasing_user')

	def is_valid(self):
		valid = super(CompanyEditOrderForm, self).is_valid()

		if not valid:
			return valid

		quantity = self.cleaned_data['quantity']

		if quantity < 1:
			self.add_error('quantity', 'Must be a positive integer.')
			return False

		if self.instance.product.max_quantity is not None and quantity > self.instance.product.max_quantity:
			self.add_error('quantity', 'Must not exceed ' + str(product.max_quantity) + '.')
			return False

		return valid


#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
Created on Tue Aug  6 16:19:52 2019

@author: MeganT
"""

#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
Created on Mon Aug  5 14:47:06 2019

@author: MeganT
"""

from astropy.io import fits
import numpy.linalg as li
from extreme_deconvolution import extreme_deconvolution
import matplotlib.pyplot as plt
import shutil
import os
import numpy as np
import math
import sys

with fits.open(os.getcwd()+'/simulated_buzzard_data.fits') as hdul:
#    hdul.info()
    data = hdul[1].data
    header = hdul[1].header
    cols = hdul[1].columns

os.chdir(os.getcwd()+"/xdc_t4")
num_cols = 4
num_rows = data.field(0).shape[1]

ndata = np.zeros((num_rows, num_cols))
errors = np.zeros((num_rows, num_cols))
for i in range(num_cols):
    ndata[:, i] = data.field(i + 4)
    errors[:, i] = data.field(i + 8)

for i in range(num_rows):
    total_flux = sum(ndata[i, :])
    ndata[i, :] = ndata[i, :] / total_flux
    errors[i, :] = errors[i, :] / total_flux

# We fill ndata with the true values of the flux
true_data = np.zeros((num_rows, num_cols))

for i in range(num_cols):
    true_data[:, i] = data.field(i)

# Normalize the true_data
for i in range(num_rows):
    total_flux = sum(true_data[i, :])
    true_data[i, :] = true_data[i, :] / total_flux

neurons = []
with open('neurons.txt', 'r') as filehandle:
    lines = filehandle.readlines()
    for l in lines:
        strings = l.split(' ')
        a = np.asarray(strings[:len(strings)-1])
        neurons.append(a.astype(np.float))
neurons = np.array(neurons)
#print(neurons)

bin_values = []
with open('bin_values.txt', 'r') as filehandle:
    lines = filehandle.read().splitlines()
    for l in lines:
        bin_values.append(float(l))
#print(bin_values)

xmean = []
with open('new_xmeans.txt', 'r') as filehandle:
    lines = filehandle.readlines()
    for l in lines:
        strings = l.split(' ')
        a = np.asarray(strings[:len(strings)-1])
        xmean.append(a.astype(np.float))
xmean = np.array(xmean)
#print(xmean)

dx = 4
xcovar = []
with open('new_xcovars.txt', 'r') as filehandle:
    lines = filehandle.readlines()
    for l in lines:
        strings = l.split(' ')
        a = np.asarray(strings[:len(strings)-1])
        matrix = np.zeros([dx, dx])
        k = 0
        for i in range(dx):
            for j in range(dx):
                matrix[i][j] = a[k]
                k += 1
        xcovar.append(matrix.astype(np.float))
xcovar = np.array(xcovar)

xamp1 = []
with open('new_xamps.txt', 'r') as filehandle:
    lines = filehandle.read().splitlines()
    for l in lines:
        xamp1.append(float(l))
xamp1 = np.array(xamp1)
#print(xamp1)

# Bins true distribution into neurons
samples = 2
data = np.zeros([num_rows*samples,4],dtype=np.float)
#sigma is 0 for all dimensions
for i in range(num_rows):
    x = true_data[i,:]
    data[i*samples:(i+1)*samples, :] = np.asarray(
            [np.random.normal(x[0], 0,samples),
             np.random.normal(x[1], 0,samples),
             np.random.normal(x[2], 0, samples),
             np.random.normal(x[3], 0, samples)]).T

# Apply MCMC to sample the XDC probability
tot_samp = 100000 * 100
dim = xmean.shape[1]

#figure out the number of gaussians in the model
Ngauss = len(xamp1)

samples = np.zeros([tot_samp,dim],dtype=np.double)

#loop over the gaussians in the model
sample_count = 0  #counter for number of samples
count_samples = np.zeros(Ngauss)
for g in range(0,Ngauss):

    #set the number of samples proportional to the amplitude of this gaussian
    Gsamp = int(tot_samp*xamp1[g])
    count_samples[g] = Gsamp

    samples[sample_count:(sample_count+Gsamp),:] = np.random.multivariate_normal(xmean[g], xcovar[g], Gsamp)
    sample_count+=Gsamp

count_samples = np.asarray(count_samples, np.int32)

# Plot the contours
#V = xcovar
#xamp = xamp1
old_samples = samples.copy()
#g_factor = 50 / np.mean(samples[:, 0])
#r_factor = 50 / np.mean(samples[:, 1])
#i_factor = 50 / np.mean(samples[:, 2])
#z_factor = 50 / np.mean(samples[:, 3])
#samples[:, 0] *= np.repeat([g_factor], samples.shape[0], axis = 0)
#samples[:, 1] *= np.repeat([r_factor], samples.shape[0], axis = 0)
#samples[:, 2] *= np.repeat([i_factor], samples.shape[0], axis = 0)
#samples[:, 3] *= np.repeat([z_factor], samples.shape[0], axis = 0)
#print(samples[:,0])
#g_diff = np.amax(samples[:,0]) - np.amin(samples[:,0])
#r_diff = np.amax(samples[:,1]) - np.amin(samples[:,1])
#i_diff = np.amax(samples[:,2]) - np.amin(samples[:,2])
#z_diff = np.amax(samples[:,3]) - np.amin(samples[:,3])
g_min = np.amin(samples[:,0])
r_min = np.amin(samples[:,1])
i_min = np.amin(samples[:,2])
z_min = np.amin(samples[:,3])
g_max = np.amax(samples[:,0])
r_max = np.amax(samples[:,1])
i_max = np.amax(samples[:,2])
z_max = np.amax(samples[:,3])

g_min_u = np.amin(old_samples[:,0])
r_min_u = np.amin(old_samples[:,1])
i_min_u = np.amin(old_samples[:,2])
z_min_u = np.amin(old_samples[:,3])
g_max_u = np.amax(old_samples[:,0])
r_max_u = np.amax(old_samples[:,1])
i_max_u = np.amax(old_samples[:,2])
z_max_u = np.amax(old_samples[:,3])

#print(g_min,r_min,i_min,z_min,g_max,r_max,i_max,z_max)
#g_min = 0.0
#r_min = 0.0
#i_min = 0.15
#z_min = 0.15
#g_max = 0.4
#r_max = 0.3
#i_max = 0.4
#z_max = 0.75

g_diff = g_max-g_min
r_diff = r_max-r_min
i_diff = i_max-i_min
z_diff = z_max-z_min

pixel_num=301

samples[:, 0] =  (samples[:, 0] - np.repeat([g_min], samples.shape[0])) * \
 np.repeat([pixel_num/g_diff], samples.shape[0])
samples[:, 1] = (samples[:, 1] - np.repeat([r_min], samples.shape[0])) * \
 np.repeat([pixel_num/r_diff], samples.shape[0])
samples[:, 2] = (samples[:, 2] - np.repeat([i_min], samples.shape[0])) * \
 np.repeat([pixel_num/i_diff], samples.shape[0])
samples[:, 3] = (samples[:, 3] - np.repeat([z_min], samples.shape[0])) * \
 np.repeat([pixel_num/z_diff], samples.shape[0])

image_gr = np.zeros([pixel_num, pixel_num],dtype=np.double)
image_ri = np.zeros([pixel_num, pixel_num],dtype=np.double)
image_iz = np.zeros([pixel_num, pixel_num],dtype=np.double)
image_gz = np.zeros([pixel_num, pixel_num],dtype=np.double)

for sample in samples:
    g = int(sample[0])
    r = int(sample[1])
    i = int(sample[2])
    z = int(sample[3])
    if g > 0 and g < pixel_num and r > 0 and r < pixel_num:
        image_gr[g, r] += 1
    if r > 0 and r < pixel_num and i > 0 and i < pixel_num:
        image_ri[r, i] += 1
    if i > 0 and i < pixel_num and z > 0 and z < pixel_num:
        image_iz[i, z] += 1
    if g > 0 and g < pixel_num and z > 0 and z < pixel_num:
        image_gz[g, z] += 1

for i in range(image_gr.shape[0]):
    for j in range(image_gr.shape[1]):
        image_gr[i][j] /= samples.shape[0]

for i in range(image_ri.shape[0]):
    for j in range(image_ri.shape[1]):
        image_ri[i][j] /= samples.shape[0]

for i in range(image_iz.shape[0]):
    for j in range(image_iz.shape[1]):
        image_iz[i][j] /= samples.shape[0]

for i in range(image_gz.shape[0]):
    for j in range(image_gz.shape[1]):
        image_gz[i][j] /= samples.shape[0]

#image_gr[:,:] /= np.repeat(np.array([samples.shape[0]]), np.array(image_gr.shape))
#image_ri[:,:] /= np.repeat(np.array([samples.shape[0]]), np.array(image_ri.shape))
#image_iz[:,:] /= np.repeat(np.array([samples.shape[0]]), np.array(image_iz.shape))
#image_gz[:,:] /= np.repeat(np.array([samples.shape[0]]), np.array(image_gz.shape))
image_gr_q = image_gr.copy()
image_ri_q = image_ri.copy()
image_iz_q = image_iz.copy()
image_gz_q = image_gz.copy()

from scipy.ndimage.filters import gaussian_filter

image_gr = gaussian_filter(image_gr, sigma=3)
fits.writeto('gr_model.fits',image_gr,overwrite=True)
image_ri = gaussian_filter(image_ri, sigma=3)
fits.writeto('ri_model.fits',image_ri,overwrite=True)
image_iz = gaussian_filter(image_iz, sigma=3)
fits.writeto('iz_model.fits',image_iz,overwrite=True)
image_gz = gaussian_filter(image_gz, sigma=3)
fits.writeto('gz_model.fits',image_gz,overwrite=True)

g_grmin = pixel_num
g_grmax = -1
r_grmin = pixel_num
r_grmax = -1
for i in range(image_gr.shape[0]):
    for j in range(image_gr.shape[1]):
        if image_gr[i][j] != 0:
            if i < g_grmin:
                g_grmin = i
            if j < r_grmin:
                r_grmin = j
            if i > g_grmax:
                g_grmax = i
            if j > r_grmax:
                r_grmax = j
r_rimin = pixel_num
r_rimax = -1
i_rimin = pixel_num
i_rimax = -1
for i in range(image_ri.shape[0]):
    for j in range(image_ri.shape[1]):
        if image_ri[i][j] != 0:
            if i < r_rimin:
                r_rimin = i
            if j < i_rimin:
                i_rimin = j
            if i > r_rimax:
                r_rimax = i
            if j > i_rimax:
                i_rimax = j
i_izmin = pixel_num
i_izmax = -1
z_izmin = pixel_num
z_izmax = -1
for i in range(image_iz.shape[0]):
    for j in range(image_iz.shape[1]):
        if image_iz[i][j] != 0:
            if i < i_izmin:
                i_izmin = i
            if j < z_izmin:
                z_izmin = j
            if i > i_izmax:
                i_izmax = i
            if j > z_izmax:
                z_izmax = j
g_gzmin = pixel_num
g_gzmax = -1
z_gzmin = pixel_num
z_gzmax = -1
for i in range(image_gz.shape[0]):
    for j in range(image_gz.shape[1]):
        if image_gz[i][j] != 0:
            if i < g_gzmin:
                g_gzmin = i
            if j < z_gzmin:
                z_gzmin = j
            if i > g_gzmax:
                g_gzmax = i
            if j > z_gzmax:
                z_gzmax = j

plt.close()
fig, [(ax1, ax2), (ax3, ax4)] = plt.subplots(2, 2)
plt.subplots_adjust(left=0.125, bottom=0.1, right=0.9, top=0.9, wspace=0.5, hspace=0.5)
levels_gr=[1e-8, 1.16e-6,3.65e-6,9.0e-6,2.06e-5,4.54e-5,9.9e-5,2.14e-4,4.63e-4,1e-3]
levels_ri=[2.33043e-06,7.3512e-06,1.81681e-05,4.14725e-05,9.16801e-05, \
           0.000199849,0.000432893,0.00093497,0.00201666]
levels_iz=[2.06444e-06,6.51214e-06,1.60944e-05,3.67388e-05,8.12159e-05, \
           0.000177039,0.000383483,0.000828253,0.00178648]
levels_gz=[2.21448e-06,6.98544e-06,1.72642e-05,3.9409e-05,8.71186e-05, \
           0.000189906,0.000411354,0.00088845,0.00191632]
#fig.tight_layout()
ax1.contour(image_gr, levels_gr)
interval = 50

#i_ticks = np.around(np.linspace(i_min_u, i_max_u, len(ind)),2)
#z_ticks = np.around(np.linspace(z_min_u, z_max_u, len(ind)),2)

#ax1.set_xlim([0.018689261934454145, 0.30570522434917147])
#ax1.set_ylim([-0.02498731864164616, 0.37520237750109764])
ax1.set_xlabel("R")
ax1.set_ylabel("G")
x_ind = np.arange(start = r_grmin, stop = r_grmax, step=interval)
y_ind = np.arange(start = g_grmin, stop = g_grmax, step=interval)
r_ticks = np.around(np.linspace(r_min_u, r_max_u, len(x_ind)),2)
g_ticks = np.around(np.linspace(g_min_u, g_max_u, len(y_ind)),2)
ax1.set_xticks(x_ind)
ax1.set_yticks(y_ind)
ax1.set_xticklabels(r_ticks)
ax1.set_yticklabels(g_ticks)
#ax1.set_xlim(r_grmin, r_grmax)
#ax1.set_ylim(g_grmin, g_grmax)
ax1.set_xlim(18, 225)
ax1.set_ylim(0, 270)
ax1.set_title("Normalized G vs. Normalized R")

ax2.contour(image_ri, levels_ri)
#ax2.set_xlim([0.168327331025541, 0.3818603775513799])
#ax2.set_ylim([0.01591015407155072, 0.3084843322120745])
ax2.set_xlabel("I")
ax2.set_ylabel("R")
x_ind = np.arange(start = i_rimin, stop = i_rimax, step=interval)
y_ind = np.arange(start = r_rimin, stop = r_rimax, step=interval)
i_ticks = np.around(np.linspace(i_min_u, i_max_u, len(x_ind)),2)
r_ticks = np.around(np.linspace(r_min_u, r_max_u, len(y_ind)),2)
ax2.set_xticks(x_ind)
ax2.set_yticks(y_ind)
ax2.set_xticklabels(i_ticks)
ax2.set_yticklabels(r_ticks)
#ax2.set_xlim(i_rimin, i_rimax)
#ax2.set_ylim(r_rimin, r_rimax)
ax2.set_xlim(150, 270)
ax2.set_ylim(30, 218)
ax2.set_title("Normalized R vs. Normalized I")

ax3.contour(image_iz, levels_iz)
#ax3.set_xlim([0.10870964061586308, 0.7513928898988901])
#ax3.set_ylim([0.16554822316263784, 0.38463948541428344])
ax3.set_xlabel("Z")
ax3.set_ylabel("I")
x_ind = np.arange(start = z_izmin, stop = z_izmax, step=interval)
y_ind = np.arange(start = i_izmin, stop = i_izmax, step=interval)
z_ticks = np.around(np.linspace(z_min_u, z_max_u, len(x_ind)),2)
i_ticks = np.around(np.linspace(i_min_u, i_max_u, len(y_ind)),2)
ax3.set_xticks(x_ind)
ax3.set_yticks(y_ind)
ax3.set_xticklabels(z_ticks)
ax3.set_yticklabels(i_ticks)
#ax3.set_xlim(z_izmin, z_izmax)
#ax3.set_ylim(i_izmin, i_izmax)
ax3.set_xlim(90, 300)
ax3.set_ylim(150, 280)
ax3.set_title("Normalized I vs. Normalized Z")

ax4.contour(image_gz,levels_gz)
#ax4.set_xlim([0.10870964061586297, 0.75139288989889])
#ax4.set_ylim([-0.024987318641646213, 0.3752023775010975])
ax4.set_xlabel("Z")
ax4.set_ylabel("G")
x_ind = np.arange(start = z_gzmin, stop = z_gzmax, step=interval)
y_ind = np.arange(start = g_gzmin, stop = g_gzmax, step=interval)
z_ticks = np.around(np.linspace(z_min_u, z_max_u, len(x_ind)),2)
g_ticks = np.around(np.linspace(g_min_u, g_max_u, len(y_ind)),2)
ax4.set_xticks(x_ind)
ax4.set_yticks(y_ind)
ax4.set_xticklabels(z_ticks)
ax4.set_yticklabels(g_ticks)
#ax4.set_xlim(z_gzmin, z_gzmax)
#ax4.set_ylim(g_gzmin, g_gzmax)
ax4.set_xlim(100, 300)
ax4.set_ylim(0, 230)
ax4.set_title("Normalized G vs. Normalized Z")
fig.suptitle("Scaled and Pixelated")
plt.savefig("Scaled_And_Pixelated.png")

# Plot contours of true fluxes for comparison
true_copy = true_data.copy()
#g_min = np.amin(true_copy[:,0])
#r_min = np.amin(true_copy[:,1])
#i_min = np.amin(true_copy[:,2])
#z_min = np.amin(true_copy[:,3])
#g_max = np.amax(true_copy[:,0])
#r_max = np.amax(true_copy[:,1])
#i_max = np.amax(true_copy[:,2])
#z_max = np.amax(true_copy[:,3])

#g_min_u = np.amin(true_copy[:,0])
#r_min_u = np.amin(true_copy[:,1])
#i_min_u = np.amin(true_copy[:,2])
#z_min_u = np.amin(true_copy[:,3])
#g_max_u = np.amax(true_copy[:,0])
#r_max_u = np.amax(true_copy[:,1])
#i_max_u = np.amax(true_copy[:,2])
#z_max_u = np.amax(true_copy[:,3])

#print(g_min,r_min,i_min,z_min,g_max,r_max,i_max,z_max)
#g_min = 0.0
#r_min = 0.0
#i_min = 0.15
#z_min = 0.15
#g_max = 0.4
#r_max = 0.3
#i_max = 0.4
#z_max = 0.75

#g_diff = g_max-g_min
#r_diff = r_max-r_min
#i_diff = i_max-i_min
#z_diff = z_max-z_min

pixel_num=301

true_copy[:, 0] =  (true_copy[:, 0] - np.repeat([g_min], true_copy.shape[0])) * \
 np.repeat([pixel_num/g_diff], true_copy.shape[0])
true_copy[:, 1] = (true_copy[:, 1] - np.repeat([r_min], true_copy.shape[0])) * \
 np.repeat([pixel_num/r_diff], true_copy.shape[0])
true_copy[:, 2] = (true_copy[:, 2] - np.repeat([i_min], true_copy.shape[0])) * \
 np.repeat([pixel_num/i_diff], true_copy.shape[0])
true_copy[:, 3] = (true_copy[:, 3] - np.repeat([z_min], true_copy.shape[0])) * \
 np.repeat([pixel_num/z_diff], true_copy.shape[0])

image_gr = np.zeros([pixel_num, pixel_num],dtype=np.double)
image_ri = np.zeros([pixel_num, pixel_num],dtype=np.double)
image_iz = np.zeros([pixel_num, pixel_num],dtype=np.double)
image_gz = np.zeros([pixel_num, pixel_num],dtype=np.double)

for sample in true_copy:
    g = int(sample[0])
    r = int(sample[1])
    i = int(sample[2])
    z = int(sample[3])
    if g > 0 and g < pixel_num and r > 0 and r < pixel_num:
        image_gr[g, r] += 1
    if r > 0 and r < pixel_num and i > 0 and i < pixel_num:
        image_ri[r, i] += 1
    if i > 0 and i < pixel_num and z > 0 and z < pixel_num:
        image_iz[i, z] += 1
    if g > 0 and g < pixel_num and z > 0 and z < pixel_num:
        image_gz[g, z] += 1

for i in range(image_gr.shape[0]):
    for j in range(image_gr.shape[1]):
        image_gr[i][j] /= true_copy.shape[0]

for i in range(image_ri.shape[0]):
    for j in range(image_ri.shape[1]):
        image_ri[i][j] /= true_copy.shape[0]

for i in range(image_iz.shape[0]):
    for j in range(image_iz.shape[1]):
        image_iz[i][j] /= true_copy.shape[0]

for i in range(image_gz.shape[0]):
    for j in range(image_gz.shape[1]):
        image_gz[i][j] /= true_copy.shape[0]

#image_gr[:,:] /= np.repeat(np.array([samples.shape[0]]), np.array(image_gr.shape))
#image_ri[:,:] /= np.repeat(np.array([samples.shape[0]]), np.array(image_ri.shape))
#image_iz[:,:] /= np.repeat(np.array([samples.shape[0]]), np.array(image_iz.shape))
#image_gz[:,:] /= np.repeat(np.array([samples.shape[0]]), np.array(image_gz.shape))
image_gr_p = image_gr.copy()
image_ri_p = image_ri.copy()
image_iz_p = image_iz.copy()
image_gz_p = image_gz.copy()

image_gr = gaussian_filter(image_gr, sigma=3)
fits.writeto('gr_orig.fits',image_gr,overwrite=True)
image_ri = gaussian_filter(image_ri, sigma=3)
fits.writeto('ri_orig.fits',image_ri,overwrite=True)
image_iz = gaussian_filter(image_iz, sigma=3)
fits.writeto('iz_orig.fits',image_iz,overwrite=True)
image_gz = gaussian_filter(image_gz, sigma=3)
fits.writeto('gz_orig.fits',image_gz,overwrite=True)

plt.close()
fig, [(ax1, ax2), (ax3, ax4)] = plt.subplots(2, 2)
plt.subplots_adjust(left=0.125, bottom=0.1, right=0.9, top=0.9, wspace=0.5, hspace=0.5)
levels=[1e-8, 1.16e-6,3.65e-6,9.0e-6,2.06e-5,4.54e-5,9.9e-5,2.14e-4,4.63e-4,1e-3]
#fig.tight_layout()
ax1.contour(image_gr, levels_gr)
interval = 50
#ind = np.arange(pixel_num, step=interval)

#ax1.set_xlim([0.018689261934454145, 0.30570522434917147])
#ax1.set_ylim([-0.02498731864164616, 0.37520237750109764])
ax1.set_xlabel("R")
ax1.set_ylabel("G")
#ax1.set_xticks(ind)
#ax1.set_yticks(ind)
x_ind = np.arange(start = r_grmin, stop = r_grmax, step=interval)
y_ind = np.arange(start = g_grmin, stop = g_grmax, step=interval)
r_ticks = np.around(np.linspace(r_min_u, r_max_u, len(x_ind)),2)
g_ticks = np.around(np.linspace(g_min_u, g_max_u, len(y_ind)),2)
ax1.set_xticks(x_ind)
ax1.set_yticks(y_ind)
#ax1.set_xlim(r_grmin, r_grmax)
#ax1.set_ylim(g_grmin, g_grmax)
ax1.set_xlim(18, 225)
ax1.set_ylim(0, 270)
ax1.set_xticklabels(r_ticks)
ax1.set_yticklabels(g_ticks)
ax1.set_title("Normalized G vs. Normalized R")

ax2.contour(image_ri, levels_ri)
#ax2.set_xlim([0.168327331025541, 0.3818603775513799])
#ax2.set_ylim([0.01591015407155072, 0.3084843322120745])
ax2.set_xlabel("I")
ax2.set_ylabel("R")
#ax2.set_xticks(ind)
#ax2.set_yticks(ind)
x_ind = np.arange(start = i_rimin, stop = i_rimax, step=interval)
y_ind = np.arange(start = r_rimin, stop = r_rimax, step=interval)
i_ticks = np.around(np.linspace(i_min_u, i_max_u, len(x_ind)),2)
r_ticks = np.around(np.linspace(r_min_u, r_max_u, len(y_ind)),2)
ax2.set_xticks(x_ind)
ax2.set_yticks(y_ind)
ax2.set_xticklabels(i_ticks)
ax2.set_yticklabels(r_ticks)
#ax2.set_xlim(i_rimin, i_rimax)
#ax2.set_ylim(r_rimin, r_rimax)
ax2.set_xlim(150, 270)
ax2.set_ylim(30, 218)
ax2.set_title("Normalized R vs. Normalized I")

ax3.contour(image_iz, levels_iz)
#ax3.set_xlim([0.10870964061586308, 0.7513928898988901])
#ax3.set_ylim([0.16554822316263784, 0.38463948541428344])
ax3.set_xlabel("Z")
ax3.set_ylabel("I")
#ax3.set_xticks(ind)
#ax3.set_yticks(ind)
x_ind = np.arange(start = z_izmin, stop = z_izmax, step=interval)
y_ind = np.arange(start = i_izmin, stop = i_izmax, step=interval)
z_ticks = np.around(np.linspace(z_min_u, z_max_u, len(x_ind)),2)
i_ticks = np.around(np.linspace(i_min_u, i_max_u, len(y_ind)),2)
ax3.set_xticks(x_ind)
ax3.set_yticks(y_ind)
ax3.set_xticklabels(z_ticks)
ax3.set_yticklabels(i_ticks)
#ax3.set_xlim(z_izmin, z_izmax)
#ax3.set_ylim(i_izmin, i_izmax)
ax3.set_xlim(90, 300)
ax3.set_ylim(150, 280)
ax3.set_title("Normalized I vs. Normalized Z")

ax4.contour(image_gz,levels_gz)
#ax4.set_xlim([0.10870964061586297, 0.75139288989889])
#ax4.set_ylim([-0.024987318641646213, 0.3752023775010975])
ax4.set_xlabel("Z")
ax4.set_ylabel("G")
#ax4.set_xticks(ind)
#ax4.set_yticks(ind)
x_ind = np.arange(start = z_gzmin, stop = z_gzmax, step=interval)
y_ind = np.arange(start = g_gzmin, stop = g_gzmax, step=interval)
z_ticks = np.around(np.linspace(z_min_u, z_max_u, len(x_ind)),2)
g_ticks = np.around(np.linspace(g_min_u, g_max_u, len(y_ind)),2)
ax4.set_xticks(x_ind)
ax4.set_yticks(y_ind)
ax4.set_xticklabels(z_ticks)
ax4.set_yticklabels(g_ticks)
#ax4.set_xlim(z_gzmin, z_gzmax)
#ax4.set_ylim(g_gzmin, g_gzmax)
ax4.set_xlim(100, 300)
ax4.set_ylim(0, 230)
ax4.set_title("Normalized G vs. Normalized Z")
fig.suptitle("True Fluxes Scaled and Pixelated")
plt.savefig("Scaled_And_Pixelated_True.png")

# Compute KL-Divergence using two scaled images

#image_gr_q = np.where(image_gr_q < np.amin(image_gr_p[np.where(image_gr_p != 0)]), 0, image_gr_q)
image_gr_div = np.where(image_gr_q != 0, image_gr_p/image_gr_q, 1)
image_gr_diff = np.where(image_gr_p != 0, image_gr_p * np.log(image_gr_div), 0)
print(np.sum(image_gr_diff))

image_ri_div = np.where(image_ri_q != 0, image_ri_p/image_ri_q, 1)
image_ri_diff = np.where(image_ri_p != 0, image_ri_p * np.log(image_ri_div), 0)
print(np.sum(image_ri_diff))

image_iz_div = np.where(image_iz_q != 0, image_iz_p/image_iz_q, 1)
image_iz_diff = np.where(image_iz_p != 0, image_iz_p * np.log(image_iz_div), 0)
print(np.sum(image_iz_diff))

image_gz_div = np.where(image_gz_q != 0, image_gz_p/image_gz_q, 1)
image_gz_diff = np.where(image_gz_p != 0, image_gz_p * np.log(image_gz_div), 0)
print(np.sum(image_gz_diff))

with open('fits_kl_values.txt', 'w') as filehandle:
    filehandle.write('kl value for gr: ' + str(np.sum(image_gr_diff)) + '\n')
    filehandle.write('kl value for ri: ' + str(np.sum(image_ri_diff)) + '\n')
    filehandle.write('kl value for iz: ' + str(np.sum(image_iz_diff)) + '\n')
    filehandle.write('kl value for gz: ' + str(np.sum(image_gz_diff)) + '\n')

#image_gr = np.zeros([g_diff, r_diff],dtype=np.double)
#image_ri = np.zeros([r_diff, i_diff],dtype=np.double)
#image_iz = np.zeros([i_diff, z_diff],dtype=np.double)
#image_gz = np.zeros([g_diff, z_diff,],dtype=np.double)

#d=4
##factor_arr = np.repeat([20], d, axis = 0)
#count = 0
##print(count_samples)
#for i in range(len(count_samples)):
#    num_samples = count_samples[i]
#    for j in range(count, count + num_samples):
#        #diff = bin_xmean[i][j]-xmean[i]
#        #print(samples[j])
#        diff = old_samples[j] - xmean[i]
#        scalar = np.matmul(np.matmul(diff.T, li.inv(V[i])), diff)
#        #gaussian =np.power((2*np.pi),(-d/2)) * np.power(li.det(V[j]),-1/2)* np.exp((-1/2)*scalar)
#        gaussian = (2*np.pi)**(-d/2) * (li.det(V[i]))**(-1/2)* np.exp((-1/2)*scalar)
#        (g, r, i, z) = np.asarray(samples[j, :], np.int32)
##        print(g)
#        image_gr[g - g_min, r - r_min] += xamp[i] * gaussian
#        image_ri[r - r_min, i - i_min] += xamp[i] * gaussian
#        image_iz[i - i_min, z - z_min] += xamp[i] * gaussian
#        image_gz[g - g_min, z - z_min] += xamp[i] * gaussian
#    count += num_samples


import logging
import time
from contextlib import contextmanager

from .slack import slack_notify

@contextmanager
def timer(name: str, log: bool = True):
    t0 = time.time()
    msg = f"[{name}] start"
    print(msg)
    # slack_notify(msg)

    if log:
        logging.info(msg)
    yield

    msg = f"[{name}] done in {time.time() - t0:.2f} s"
    print(msg)
    # slack_notify(msg)

    if log:
        logging.info(msg)


import pytest

from environments.identities.traits.models import Trait

@pytest.mark.parametrize(
    "value, expected_data",
    (
        (1, {"value_type": "int", "integer_value": 1}),
        (0, {"value_type": "int", "integer_value": 0}),
        ("my_string", {"value_type": "unicode", "string_value": "my_string"}),
        (True, {"value_type": "bool", "boolean_value": True}),
        (False, {"value_type": "bool", "boolean_value": False}),
        (123.4, {"value_type": "float", "float_value": 123.4}),
    ),
)
def test_generate_trait_value_data_for_value(value, expected_data):
    assert Trait.generate_trait_value_data(value) == expected_data

@pytest.mark.parametrize(
    "deserialized_data, expected_data",
    (
        ({"type": "int", "value": 1}, {"value_type": "int", "integer_value": 1}),
        ({"type": "int", "value": 0}, {"value_type": "int", "integer_value": 0}),
        (
            {"type": "unicode", "value": "my_string"},
            {"value_type": "unicode", "string_value": "my_string"},
        ),
        (
            {"type": "bool", "value": True},
            {"value_type": "bool", "boolean_value": True},
        ),
        (
            {"type": "bool", "value": False},
            {"value_type": "bool", "boolean_value": False},
        ),
        (
            {"type": "float", "value": 123.4},
            {"value_type": "float", "float_value": 123.4},
        ),
    ),
)
def test_generate_trait_value_data_for_deserialized_data(
    deserialized_data, expected_data
):
    assert Trait.generate_trait_value_data(deserialized_data) == expected_data


import requests
from nksama import bot
from pyrogram import filters
from pyrogram.types.bots_and_keyboards.inline_keyboard_button import InlineKeyboardButton
from pyrogram.types.bots_and_keyboards.inline_keyboard_markup import InlineKeyboardMarkup
from nksama.plugins.helpers import call_back_in_filter

@bot.on_callback_query(call_back_in_filter('quotek'))
def callback_quotek(_,query):
    if query.data.split(":")[1] == "change":
#         query.message.delete()
        kk = requests.get('https://animechan.vercel.app/api/random').json()
        anime = kk['anime']
        quote = kk['quote']
        character = kk['character']
        caption = f"""
**Anime:** `{anime}`
**Character:** `{character}`
**Quote:** `{quote}`"""
        query.message.edit(caption , reply_markup=InlineKeyboardMarkup([
            [InlineKeyboardButton("Change" , callback_data="quotek:change")],
        ]))

@bot.on_message(filters.command('quote'))
def quote(_,message):
    kk = requests.get('https://animechan.vercel.app/api/random').json()
    anime = kk['anime']
    quote = kk['quote']
    character = kk['character']
    caption = f"""
**Anime:** `{anime}`
**Character:** `{character}`
**Quote:** `{quote}`"""
    bot.send_message(message.chat.id , caption , reply_markup=InlineKeyboardMarkup([
        [InlineKeyboardButton("Change" , callback_data="quotek:change")]
    ]))


#%%
# Etem Kaya 12-Apr-2019
# Iris Data Set Project.

# Histogram of Iris dataset Sepal length and width

# import pandas, seaborn and matplotlib libraries
import pandas as pd
import seaborn as sns
import matplotlib.pyplot as plt

# import iris dataset with Pandas
f = pd.read_csv("iris.csv")

# split 'Setosa, Versicolor and Virginica' dataset from the entire Iris dataset
Setosa = f.loc[f["species"]=="setosa"]
Versicolor = f.loc[f["species"]=="versicolor"]
Virginica = f.loc[f["species"]=="virginica"]

# set style and background colour using seaborn library
sns.set(style="darkgrid")

# Histogram of 'Setosa, Versicolor and Virginica' Sepal Length
sns.distplot(Setosa["sepal_length"], color= "green")
sns.distplot(Versicolor["sepal_length"], color= "blue")
sns.distplot(Virginica["sepal_length"], color= "orange")
plt.title("Histogram - Sepal Length of 'Setosa, Versicolor and Virginica'")
plt.xlabel('Sepal_Length' + "\nSetosa, Versicolor and Virginica")
plt.ylabel('count')
plt.show()

# Histogram of 'Setosa, Versicolor and Virginica' Sepal Width
sns.distplot(Setosa["sepal_width"], color= "green")
sns.distplot(Versicolor["sepal_width"], color= "blue")
sns.distplot(Virginica["sepal_width"], color= "orange")
plt.title("Histogram - Sepal Width of 'Setosa, Versicolor and Virginica'")
plt.xlabel('Sepal_Width' + "\nSetosa, Versicolor and Virginica")
plt.ylabel('count')
plt.show()