## STANDARD IMPORTS
import numpy as np
import pandas as pd
import matplotlib.pyplot as plt
import matplotlib as mpl
from mpl_toolkits.mplot3d import Axes3D
import os

## MISC IMPORTS
import re

# os.system("clear")
np.set_printoptions(suppress=True)
mpl.rcParams['pdf.fonttype'] = 42
mpl.rcParams['ps.fonttype'] = 42

class DataFile:
    def __init__(self,dataPath,fileName,dataType='SIM'):

        ## ORGANIZE FILEPATH AND CREATE TRIAL DATAFRAME
        self.fileName = fileName
        self.dataPath = dataPath
        filepath = os.path.join(self.dataPath, self.fileName)

        self.trial_df = pd.read_csv(filepath,low_memory=False,comment="#")

        ## CLEAN UP TRIAL DATAFRAME
        # Drop row with "Note: ____________"
        self.dataType = dataType

        # Remove rows past final complete rollout
        final_valid_index = self.trial_df[self.trial_df['Error']=='Impact Data'].index.values[-1]
        self.trial_df = self.trial_df.iloc[:final_valid_index+1]

        # Round values to prevent floating point precision issues
        self.trial_df = self.trial_df.round(3)

        # Convert string bools to actual bools
        self.trial_df = self.trial_df.replace({'False':False,'True':True})
        self.trial_df = self.trial_df.replace({'--':np.nan,'nan':np.nan})

        ## CREATE BASIC DF OF K_EP,K_RUN, & REWARD
        self.k_df = self.trial_df.iloc[:][['k_ep','k_run']]
        self.k_df = self.k_df.drop_duplicates()

        if self.dataType=='EXP':
            self.remove_FailedRuns()

        ## COLLECT BASIC TRIAL INFO
        self.n_rollouts = int(self.trial_df.iloc[-3]['Policy_Trg_Action'])
        self.k_epMax = int(self.trial_df.iloc[-3]['k_ep'])
        self.k_runMax = int(self.trial_df.iloc[-3]['k_run'])

    def remove_FailedRuns(self):
        """Drop all rollouts in dataframe that were labelled as 'Failed Rollout'
        """

        failedRuns_list = self.trial_df.query(f"Error=='Failed Rollout'").index.tolist()

        for failed_run in failedRuns_list[::-1]: # Work backwards down list

            idx_prev_run = self.k_df.query(f"index == {failed_run}").index.tolist()[0]-1
            if idx_prev_run < 0: # Catch for very first run
                idx_prev_run = 0

            idx_start = self.k_df.iloc[idx_prev_run]['index'] + 1 # True index of trial_df to start drop
            idx_end = failed_run + 1 # Make sure to include endpoint

            self.trial_df.drop(np.arange(idx_start,idx_end,dtype='int'),inplace=True)

        self.k_df = self.trial_df.iloc[:][['k_ep','k_run']].dropna()
        self.k_df.reset_index(inplace=True)

    def select_run(self,k_ep,k_run): ## Create run dataframe from k_ep and k_run
        """Create run dataframe from selected k_ep and k_run

        Args:
            k_ep (int): Episode number
            k_run (int): Run number

        Returns:
            Experiment:
                tuple: run_df,IC_df,__,__
            Sim:
                tuple: run_df,IC_df,flip_df,impact_df
        """

        ## CREATE RUN DATAFRAME
        run_df = self.trial_df[(self.trial_df['k_ep']==k_ep) & (self.trial_df['k_run']==k_run)]

        IC_df = run_df.iloc[[-3]]       # Create DF of initial conditions
        flip_df = run_df.iloc[[-2]]     # Create DF of flip conditions
        impact_df = run_df.iloc[[-1]]   # Create DF of impact conditions
        run_df = run_df[:-3]            # Drop special rows from dataframe

        return run_df,IC_df,flip_df,impact_df

    def grab_rewardData(self):
        """Returns summary of rewards

        Data Type: Sim/Exp

        Returns:
            tuple: k_ep_r,rewards,k_ep_ravg,rewards_avg

            ===========================
            k_ep_r: reward episodes (np.array)
            rewards: array of rewards (np.array)
            k_ep_ravg: episode of averaged reward (np.array)
            rewards_avg: array of averaged rewards per episode (np.array)
        """
        ## CREATE ARRAYS FOR REWARD, K_EP
        reward_df = self.trial_df.iloc[:][['k_ep','mu','flip_flag']].dropna() # Create df from k_ep/rewards and drop blank reward rows
        reward_df = reward_df.iloc[:][["k_ep","flip_flag"]].astype('float')
        rewards_arr = reward_df.to_numpy()
        rewards = rewards_arr[:,1]
        k_ep_r = rewards_arr[:,0]

        ## CREATE ARRAYS FOR REWARD_AVG, K_EP
        rewardAvg_df = reward_df.groupby('k_ep').mean() # Group rows by k_ep value and take their mean
        rewards_avg = rewardAvg_df.to_numpy()
        k_ep_ravg = reward_df.k_ep.unique() # Drops duplicate values so it matches dimension of rewards_avg (1 avg per ep and n ep)

        return k_ep_r,rewards,k_ep_ravg,rewards_avg

    def plot_rewardData(self,ymax=300):
        """Plot rewards for entire trial
        Data Type: Sim/Exp
        """
        k_ep_r,rewards,k_ep_ravg,rewards_avg = self.grab_rewardData()

        fig = plt.figure()
        ax = fig.add_subplot(111)
        ax.scatter(k_ep_r,rewards,marker='_',color='black',alpha=0.5,label='Reward')
        ax.scatter(k_ep_ravg,rewards_avg,marker='o',color='red',label='Average Reward')

        ax.set_ylabel("Reward")
        ax.set_xlabel("k_ep")
        ax.set_xlim(-2,self.k_epMax+5)
        ax.set_ylim(-2,ymax)
        ax.set_title(f"Reward vs Episode | Rollouts per Episode: {self.n_rollouts}")
        ax.legend()
        ax.grid()

        fig.tight_layout()
        plt.show()

    ## POLICY FUNCTIONS
    def grab_policy(self,k_ep=0,k_run=0):
        """Returns policy from specific run

        Data Type: Sim/Exp

        Args:
            k_ep (int): Episode number
            k_run (int): Run number

        Returns:
            np.array: [RREV,G1,G2,...]
        """
        run_df,IC_df,_,_ = self.select_run(k_ep,k_run)

        ## SELECT POLICY
        policy = IC_df.iloc[0]['policy']
        policy = np.fromstring(policy[1:-1], dtype=float, sep=' ')  # Convert str to np array

        return policy

    def grab_convg_data(self,k_ep=0,k_run=0):
        """Returns series of arrays for episodes, mu values, and sigma values

        Data Type: Sim/Exp

        Returns:
           k_ep_arr [np.array]: Array of episodes
           mu_arr [np.array]: Array of mu values across trial
           sigma_arr [np.array]: Array of sigma values across trial

        """
        ## CLEAN AND GRAB DATA FOR MU & SIGMA
        policy_df = self.trial_df.iloc[:][['k_ep','mu','sigma']]
        policy_df = policy_df.dropna().drop_duplicates()
        k_ep_arr = policy_df.iloc[:]['k_ep'].to_numpy()

        ## CREATE ARRAYS FOR MU & SIGMA OVER TRIAL
        mu_arr = []
        for mu in policy_df.iloc[:]['mu']:
            mu = np.fromstring(mu[1:-1], dtype=float, sep=' ')
            mu_arr.append(mu)
        mu_arr = np.array(mu_arr)

        sigma_arr = []
        for sigma in policy_df.iloc[:]['sigma']:
            sigma = np.fromstring(sigma[1:-1], dtype=float, sep=' ')
            sigma_arr.append(sigma)
        sigma_arr = np.array(sigma_arr)

        return k_ep_arr,mu_arr,sigma_arr

    def plot_policy_convg(self): ## NEEDS UPDATED
        """Creates subplots to show convergence for policy gains
        """

        ## GRAB CONVERGENCE DATA
        k_ep_arr,mu_arr,sigma_arr = self.grab_convg_data()

        num_col = mu_arr.shape[1] # Number of policy gains in mu [Currently 3]
        G_Labels = ['Tau_trigger','My','G2','G3','G4','G5'] # List of policy gain names
        colors = ['tab:blue','tab:orange','tab:green']
        # vx,vy,vz = self.grab_vel_d()

        # Vel = np.sqrt(vx**2 + vz**2)
        # phi = np.arctan2(vz,vx)

        ## CREATE SUBPLOT FOR MU
        fig = plt.figure()
        ax = fig.add_subplot(111)
        for jj in range(num_col): # Iterate through gains and plot each
            ax.plot(k_ep_arr,mu_arr[:,jj], color=colors[jj],linestyle='--',label=G_Labels[jj])
            ax.plot(k_ep_arr,mu_arr[:,jj] + sigma_arr[:,jj], color=colors[jj])
            ax.plot(k_ep_arr,mu_arr[:,jj] - sigma_arr[:,jj], color=colors[jj])
            ax.fill_between(k_ep_arr,mu_arr[:,jj] + sigma_arr[:,jj],mu_arr[:,jj] - sigma_arr[:,jj],alpha=0.5)

        ax.set_ylabel('Policy Values')
        ax.set_xlabel('K_ep')
        ax.set_ylim(0,10) # Set lower ylim
        # ax.set_title(f'Policy Value vs Episode ($Vel_d$ = {Vel:.2f} | $\phi$ = {np.rad2deg(phi):.2f}$^{{\circ}}$)')
        ax.legend(ncol=num_col)
        ax.grid()
        fig.tight_layout()
        plt.show()

    def grab_finalPolicy(self):
        """Returns the final policy for a trial

        Data Type: Sim/Exp

        Returns:
           mu [np.array]: Final policy average
           sigma [np.array]: Final policy standard deviation

        """

        ## FIND FINAL RUN DF
        k_ep = self.trial_df.iloc[-1]['k_ep']
        k_run = self.trial_df.iloc[-1]['k_run']
        run_df,IC_df,_,_ = self.select_run(k_ep,k_run)

        ## SELECT MU & SIGMA
        mu = IC_df.iloc[0]['mu']
        mu = np.fromstring(mu[1:-1], dtype=float, sep=' ')  # Convert str to np array

        sigma = IC_df.iloc[0]['sigma']
        sigma = np.fromstring(sigma[1:-1], dtype=float, sep=' ')

        return mu,sigma

    def grab_RLParams(self):
        """Returns initial RL Parameters

        Data Type: Sim/Exp

        Returns:
            mu [float]: Mean value
            sigma [float]: Standard deviation
        """

        # CREATE INITIAL PARAMETER DATAFRAME
        param_df = self.trial_df.iloc[:][['alpha_mu','alpha_sig','mu','sigma']].dropna()

        mu = param_df.iloc[0]['mu']
        mu = np.fromstring(mu[1:-1], dtype=float, sep=' ')

        sigma = param_df.iloc[0]['sigma']
        sigma = np.fromstring(sigma[1:-1], dtype=float, sep=' ')

        return mu,sigma

    ## STATE FUNCTIONS
    def grab_stateData(self,k_ep,k_run,stateName: list):
        """Returns np.array of specified state

        Data Type: Sim/Exp

        Args:
            k_ep (int): Episode number
            k_run (int): Run number
            stateName (str): State name

        Returns:
            state [np.array]: Returned state values
        """
        run_df,_,_,_ = self.select_run(k_ep,k_run)

        ## GRAB STATE DATA AND CONVERT TO NUMPY ARRAY
        state = run_df.iloc[:][stateName]
        state = state.to_numpy(dtype=np.float64)

        return state

    def plot_state(self,k_ep,k_run,stateName: list):
        """Plot state data from run vs time

        Data Type: Sim/Exp

        Args:
            k_ep (int): Ep. Num
            k_run (int): Run. Num
            stateName list(str): State name from .csv
        """

        ## PLOT STATE/TIME DATA
        fig = plt.figure()
        ax = fig.add_subplot(111)

        t = self.grab_stateData(k_ep,k_run,'t')
        t_norm = t - np.min(t) # Normalize time array

        ## GRAB STATE/FLIP DATA
        for state in stateName:

            color = next(ax._get_lines.prop_cycler)['color']

            ## PLOT STATE DATA
            state_vals = self.grab_stateData(k_ep,k_run,state)
            ax.plot(t_norm,state_vals,label=f"{state}",zorder=1,color=color)

            # ## MARK STATE DATA AT FLIP TIME
            # t_flip,t_flip_norm = self.grab_flip_time(k_ep,k_run)
            # state_flip = self.grab_flip_state(k_ep,k_run,state)
            # ax.scatter(t_flip_norm,state_flip,label=f"{state}-Flip",zorder=2,
            #             marker='x',
            #             color=color,
            #             s=100)

            # # ## MARK STATE DATA AT IMPACT TIME
            # t_impact,t_impact_norm = self.grab_impact_time(k_ep,k_run)
            # state_flip = self.grab_impact_state(k_ep,k_run,state)
            # ax.scatter(t_impact_norm,state_flip,label=f"{state}-Impact",zorder=2,
            #             marker='s',
            #             color=color,
            #             s=50)

            ## MARK STATE DATA AT TRAJ START
            # if self.dataType == 'EXP':
            #     t_traj,t_traj_norm = self.grab_traj_start(k_ep,k_run)
            #     state_traj = self.grab_traj_state(k_ep,k_run,state)
            #     ax.scatter(t_traj_norm,state_traj,label=f"{state}-Traj",zorder=2,
            #                 marker='o',
            #                 color=color,
            #                 s=25)

        ax.set_xlabel("time [s]")
        ax.legend()
        ax.grid()

        plt.show()

    def grab_time(self,k_ep,k_run):
        run_df,_,_,_ = self.select_run(k_ep,k_run)

        ## GRAB STATE DATA AND CONVERT TO NUMPY ARRAY
        state = run_df.iloc[:]['t']
        t = state.to_numpy(dtype=np.float64)

        t_normalized = t - np.min(t)

        return t,t_normalized

    ## DESIRED IC FUNCTIONS
    def grab_vel_IC(self):

        """Return IC velocities

        Returns:
            vx_IC (float): Desired flip x-velocity
            vy_IC (float): Desired flip y-velocity
            vz_IC (float): Desired flip z-velocity
        """
        vel_df = self.trial_df[['mu','vx','vy','vz']].dropna()
        vx_IC,vy_IC,vz_IC = vel_df.iloc[0][['vx','vy','vz']]

        return vx_IC,vy_IC,vz_IC

    def grab_vel_IC_2D_angle(self):
        """Return IC velocities

        Returns:
            Vel_IC (float): Desired flight velocity
            phi_IC (float): Desired flight angle
        """
        ## SELECT X,Y,Z LAUNCH VELOCITIES
        Vel_IC,phi_IC = self.trial_df.iloc[-3][['vx','vy']].astype('float').tolist()

        return Vel_IC,phi_IC

    ## FLIP FUNCTIONS
    def grab_flip_time(self,k_ep,k_run):
        """Returns time of flip

        Data Type: Sim/Exp

        Args:
            k_ep (int): Episode number
            k_run (int): Run number

        Returns:
            float: t_flip,t_flip_normalized
        """

        run_df,_,flip_df,_ = self.select_run(k_ep,k_run)
        if flip_df.iloc[0]['flip_flag'] == True: # If Impact Detected
            t_flip = float(flip_df.iloc[0]['t'])
            t_flip_norm = t_flip - float(run_df.iloc[0]['t']) # Normalize flip time
        else:
            t_flip = np.nan
            t_flip_norm = np.nan

        return t_flip,t_flip_norm

    def grab_flip_state(self,k_ep,k_run,stateName: list):
        """Returns desired state at time of flip

        Data Type: Sim/Exp

        Args:
            k_ep (int): Episode number
            k_run (int): Run number
            stateName (str): State name

        Returns:
            float: state_flip
        """

        _,_,flip_df,_ = self.select_run(k_ep,k_run)
        if flip_df.iloc[0]['flip_flag'] == True: # If Impact Detected

            state_flip = float(flip_df.iloc[0][stateName])
        else:
            state_flip = np.nan

        return state_flip

    ## IMPACT FUNCTIONS
    def grab_impact_time(self,k_ep,k_run):
        """Return time at impact

        Args:
            k_ep (int): Episode number
            k_run (int): Run number
            accel_threshold (float, optional): Threshold at which to declare impact. Defaults to -4.0.

        Returns:
            t_impact [float]: Time at impact
            t_impact_norm [float]: Normalized impact time
        """

        run_df,_,_,impact_df = self.select_run(k_ep,k_run)
        if impact_df.iloc[0]['impact_flag'] == True: # If Impact Detected
            t_impact = float(impact_df.iloc[0]['t'])
            t_impact_norm = t_impact - float(run_df.iloc[0]['t']) # Normalize flip time
        else:
            t_impact = np.nan
            t_impact_norm = np.nan

        return t_impact,t_impact_norm

    def grab_impact_state(self,k_ep,k_run,stateName):
        """Returns desired state at time of impact

        Data Type: Sim/Exp

        Args:
            k_ep (int): Episode number
            k_run (int): Run number
            stateName (str): State Name

        Returns:
            state_impact [float]: State at impact
        """

        _,_,_,impact_df = self.select_run(k_ep,k_run)
        if impact_df.iloc[0]['impact_flag'] == True: # If Impact Detected

            state_impact = float(impact_df.iloc[0][stateName])
        else:
            state_impact = np.nan

        return state_impact

    def trigger2impact(self,k_ep,k_run):
        t_flip,_ = self.grab_flip_time(k_ep,k_run)
        t_impact,_ = self.grab_impact_time(k_ep,k_run)

        t_delta = t_impact-t_flip
        return t_delta

    def landing_conditions(self,k_ep,k_run):
        """Returns landing data for number of leg contacts, the impact leg, contact list, and if body impacted

        Args:
            k_ep (int): Episode number
            k_run (int): Run number

        Returns:
            leg_contacts (int): Number legs that joined to the ceiling
            contact_list (list): List of the order legs impacted the ceiling
            body_impact (bool): True if body impacted the ceiling
        """
        _,_,_,impact_df = self.select_run(k_ep,k_run)
        body_impact = impact_df.iloc[0]['flip_flag']
        leg_contacts = int(impact_df.iloc[0]['F_thrust'])
        contact_list = impact_df.iloc[0][['Theta_x','Theta_x_est','Theta_y','Theta_y_est']].to_numpy(dtype=np.int8)

        return leg_contacts,contact_list,body_impact

    def landing_rate(self,N:int=3,reward_cutoff:float=0.01):
        """Returns the landing rate for the final [N] episodes

        Args:
            N (int, optional): Final [N] episodes to analyze. Defaults to 3.
            reward_cutoff (float, optional): Reward values below this cutoff come from sim errors and not performance. Defaults to 3.00.

        Returns:
            landing_rate_4_leg (float): Four leg landing success percentage
            landing_rate_2_leg (float): Two leg landing success percentage
            contact_rate (float): Combined landing success percentage for either two leg or four leg landing
            contact_list (list,int): List of each conctact case for each run recorded
        """

        ## CREATE ARRAY OF ALL EP/RUN COMBINATIONS FROM LAST 3 ROLLOUTS
        # Use reward to extract only the valid attempts and not simulation mishaps
        ep_df = self.trial_df.iloc[:][['k_ep','k_run','mu','flip_flag']].dropna().drop(columns='mu')
        ep_df = ep_df.astype('float').query(f'flip_flag >= {reward_cutoff}')
        ep_arr = ep_df.iloc[-self.n_rollouts*N:].to_numpy() # Grab episode/run listing from past N rollouts

        ## ITERATE THROUGH ALL RUNS AND RECORD VALID LANDINGS
        four_leg_landing = 0
        two_leg_landing = 0
        one_leg_landing = 0
        failed_landing = 0
        contact_list = []

        for k_ep,k_run in ep_arr[:,:2]:

            ## RECORD LANDING CONDITIONS
            leg_contacts,_,body_impact = self.landing_conditions(k_ep, k_run)

            if leg_contacts >= 3 and not body_impact:
                four_leg_landing += 1

            elif leg_contacts == 2 and not body_impact:
                two_leg_landing += 1

            contact_list.append(leg_contacts)

        ## CALC LANDING PERCENTAGE
        landing_rate_4leg = np.round(four_leg_landing/(N*self.n_rollouts),2)
        landing_rate_2leg = np.round(two_leg_landing/(N*self.n_rollouts),2)
        contact_rate = np.round((four_leg_landing + two_leg_landing)/(N*self.n_rollouts),2)

        return landing_rate_4leg,landing_rate_2leg,contact_rate,np.array(contact_list)

    def grab_trial_data(self,func,*args,N:int=3,reward_cutoff:float=0.05,landing_cutoff:int=3,**kwargs):

        ## CREATE ARRAY OF ALL EP/RUN COMBINATIONS FROM LAST 3 ROLLOUTS
        # Use reward to extract only the valid attempts and not simulation mishaps
        ep_df = self.trial_df.iloc[:][['k_ep','k_run','mu','flip_flag']].dropna().drop(columns='mu')
        ep_df = ep_df.astype('float').query(f'flip_flag >= {reward_cutoff}')
        ep_arr = ep_df.iloc[-self.n_rollouts*N:].to_numpy() # Grab episode/run listing from past N rollouts

        ## ITERATE THROUGH ALL RUNS AND FINDING IMPACT ANGLE
        var_list = []
        for k_ep,k_run in ep_arr[:,:2]:

            leg_contacts,_,body_contact = self.landing_conditions(k_ep, k_run)
            if leg_contacts >= landing_cutoff and body_contact == False: # IGNORE FAILED LANDINGS
                var_list.append(func(k_ep,k_run,*args,**kwargs))

        ## RETURN MEAN AND STD OF STATE
        trial_mean = np.nanmean(var_list)
        trial_std = np.nanstd(var_list)
        trial_arr = var_list

        return np.round(trial_mean,3),trial_std,trial_arr

if __name__ == "__main__":

    dataPath = f"/home/bhabas/catkin_ws/src/sar_simulation/crazyflie_logging/local_logs/"

    fileName = "Control_Playground--trial_24--NL.csv"
    trial = DataFile(dataPath,fileName)

    trial.plot_policy_convg()


import astar
import cv2
from queue import PriorityQueue
import numpy as np
from threading import Timer

NODE_Obstacle = "obstacle"
NODE_Free = "free"
NODE_InOpenList = 'open'
NODE_InCloseList = 'close'
NODE_Unknown = 4

adjacent=[]
adjacent.extend([(1,0),(1,1),(1,-1),(0,1),(0,-1),(-1,0),(-1,1),(-1,-1)])
print(adjacent)
Inf = float('inf')
class Node(object):
    def __init__(self,x=0,y=0,g=Inf,f=Inf,next_=None):
        self.x = x
        self.y = y
        self.h = 0
        self.g = g
        self.f = f
        self.status = NODE_Free
        self.ParentNode=next_
    def point(self):
        return (self.x,self.y)

    def __lt__(self,other):
        if self.f < other.f:
            return True
        else:
            return False
    def __repr__(self):
        return str((self.x,self.y,self.f,self.status))

class mAstar(object):
    def __init__(self):
        self.path=[]
        self.panel = "Img"
        self.img=None
        self.graph = {}
        self.start=None
        self.end=None

    def refresh(self):
        self.show()
        self.mTimer = Timer(20,self.refresh)
        self.mTimer.start()

    def load_img(self,imgpath):
        self.img = cv2.imread(imgpath,cv2.IMREAD_GRAYSCALE)
        self.img_raw = self.img.copy()

    def mousecallback(self,event,x,y,flags,param):
        if event == cv2.EVENT_LBUTTONUP:
            if self.start!=None:
                self.end = Node(x,y)
                cv2.circle(self.img,(x,y),1,0,1)
                cv2.imshow(self.panel,self.img)
                self.search(self.start,self.end)
                self.start=None
                return
            self.img = self.img_raw.copy()
            self.start = Node(x,y)
            cv2.circle(self.img,(x,y),1,0,1)
            cv2.imshow(self.panel,self.img)

    def show(self):
        cv2.namedWindow(self.panel,cv2.WINDOW_NORMAL)
        cv2.setMouseCallback(self.panel,self.mousecallback)
        cv2.imshow(self.panel,self.img)
        cv2.waitKey()

    def search(self,start,end):
        self.img[start.y][start.x]=50
        self.img[end.y][end.x]=50
        openlist = PriorityQueue()
        closelist= []
        width = self.img.shape[1]
        height=self.img.shape[0]

        # input()
        t = cv2.getTickCount()
        print(t)
        for x in range(0,width):
            for y in range(0,height):
                tnode = Node(x,y)
                if self.img[y,x]>250:
                    tnode.status = NODE_Free
                else:
                    tnode.status= NODE_Obstacle
                    # print('obstacle',tnode)
                if x not in self.graph.keys():
                    self.graph[x]={}
                self.graph[x][y]=tnode
        print(cv2.getTickCount()-t)
        tnode = Node(start.x,start.y,0,0)
        tnode.status = NODE_Free
        self.graph[start.x][start.y]=tnode
        tnode = Node(end.x,end.y)
        tnode.status = NODE_Free
        self.graph[end.x][end.y]=tnode

        openlist.put(self.graph[start.x][start.y])

        t=cv2.getTickCount()
        print('start:',t)
        while not openlist.empty():
            # input()
            current = openlist.get()
            # print('\nclose:',current)
            if current.x == end.x and current.y == end.y:
                end=current
                print("Found.\n")

                break
            current.status = NODE_InCloseList
            self.img[current.y,current.x]=128
            self.img[start.y][start.x]=50
            for i in range(0,8):
                # input()
                xnow = current.x + adjacent[i][0]
                ynow = current.y + adjacent[i][1]
                if 0<=xnow<width and 0<=ynow<height:
                    neighbor = self.graph[xnow][ynow]
                    if neighbor.status == NODE_Obstacle or neighbor.status == NODE_InCloseList:
                        # print('ignore:',neighbor)
                        continue
                    cost = self.movementcost(current,neighbor) + current.g
                    if cost >= neighbor.g:
                        # print('cost {0} bigger than {1}:'.format(cost,neighbor.g))
                        continue
                    neighbor.ParentNode = current
                    neighbor.g = cost
                    neighbor.f = neighbor.g+self.heuristic(neighbor,end)
                    # print("put:",neighbor)
                    openlist.put(neighbor)
                    # cv2.namedWindow(self.panel,cv2.WINDOW_NORMAL)
            # cv2.imshow(self.panel,self.img)

        print('finished')
        node=end.ParentNode
        print(node)
        print(cv2.getTickCount()-t)
        while node.x!=start.x or node.y!=start.y:
            cv2.circle(self.img,node.point(),1,0,1)
            self.img[node.y][node.x]=80
            node = node.ParentNode
            print(node)
            cv2.imshow(self.panel,self.img)
            cv2.waitKey(1)
        cv2.waitKey(2)

    def movementcost(self,A,B):
        stepX = abs(A.x - B.x)
        stepY = abs(A.y - B.y)
        step = stepX + stepY
        if step == 1:
            cost = 10
        else:
            cost = 14
        return cost

    def heuristic(self,A,B):
        h = 10*((abs(A.x-B.x)+abs(A.y-B.y)))
        return h

if __name__ == "__main__":

    pf=mAstar()
    pf.load_img("test2.jpg")
    pf.show()


from django.contrib import admin
from django.utils.html import format_html
from django.utils.safestring import mark_safe

from .models import Patient, PatientExamination, PatientHistory, Testmodel

# Register your models here.

class PatientAdmin(admin.ModelAdmin):

    TEXT = "Click on to 'hide' to hide this section."

    list_display = ['name', 'surname', 'middlename', 'patient_ID', 'created_at', 'updated_at']
    list_filter = ['name', 'surname', 'patient_ID']
    fieldsets = (
        ('Basic Info', {
            'fields': [('name', 'middlename', 'surname',), 'birthdate', 'patient_ID'],
        }),
        ('Contact Info', {
            'fields': [('phone', 'mobile', 'email')],
            'classes': ('collapse',),
            'description': '%s' % TEXT,
        }),
        ('Address Info', {
            'fields': [('address', 'city', 'country')],
            'classes': ('collapse',),
            'description': '%s' % TEXT,
        }),
    )

class PatientExaminationAdmin(admin.ModelAdmin):
    list_display = ['patient', 'patient_status', 'prescription', 'doctor', 'examined', 'updated_at']
    list_filter = ['patient', 'patient_status', 'doctor']

class PatientHistoryAdmin(admin.ModelAdmin):
    list_display = ['patient', 'doctor', 'history']
    list_filter = ['patient', 'doctor']

class TestmodelAdmin(admin.ModelAdmin):
    list_display = ['name', 'surname']

admin.site.register(Patient, PatientAdmin)
admin.site.register(PatientExamination, PatientExaminationAdmin)
admin.site.register(PatientHistory, PatientHistoryAdmin)
admin.site.register(Testmodel, TestmodelAdmin)


from django.urls import path
from .views import *
from django.views.decorators.cache import cache_page

urlpatterns = [
    path('ad/<int:pk>', AdDetail.as_view(), name = 'ad_detail'),
    path('ad/<int:pk>/update', UpdateAd.as_view(), name = 'update_ad'),
    path('ad/create', CreateAd.as_view(), name = 'create_ad'),
    path('', AdList.as_view(), name = 'ad_list'),
    path('ad/<int:pk>/create_answer', CreateAnswer.as_view(), name = 'create_answer'),
    path('answer/tomyad', AnswersToMyAd.as_view(), name = 'answer_to_my_ad'),
    path('answer/my', MyAnswers.as_view(), name = 'my_answer'),
    path('answer/<int:pk>', AnswerDetail.as_view(), name = 'answer_detail'),
    path('answer/<int:pk>/confirm_answer', ConfirmAnswer.as_view(), name = 'confirm_answer'),
    path('answer/<int:pk>/delete_answer', DeleteAnswer.as_view(), name = 'delete_answer'),
]


# coding=utf-8
from __future__ import absolute_import
from __future__ import division
from __future__ import print_function

import copy
import json
import logging
import math
import torch
import numpy as np
import torch.nn as nn
from torch.nn import CrossEntropyLoss, Dropout, Embedding, Softmax

logger = logging.getLogger(__name__)

def Linear(i_dim, o_dim, bias=True, std=0.02):
    m = nn.Linear(i_dim, o_dim, bias)
    nn.init.normal_(m.weight, std=std)
    if bias:
        nn.init.constant_(m.bias, 0.)
    return m

def gelu(x):
    return x * 0.5 * (1.0 + torch.erf(x / 1.41421))

def swish(x):
    return x * torch.sigmoid(x)

ACT2FN = {"gelu": gelu, "relu": torch.nn.functional.relu, "swish": swish}

class Attention(nn.Module):
    def __init__(self, config):
        super(Attention, self).__init__()
        self.hidden_size = config.hidden_size
        self.num_attention_heads = config.num_heads
        self.attention_head_size = int(config.hidden_size / config.num_heads)
        self.all_head_size = self.num_attention_heads * self.attention_head_size

        self.query = Linear(self.hidden_size, self.all_head_size, std=config.initializer_range)
        self.key = Linear(self.hidden_size, self.all_head_size, std=config.initializer_range)
        self.value = Linear(self.hidden_size, self.all_head_size, std=config.initializer_range)

        self.dropout = Dropout(config.dropout_prob)

        self.fc = Linear(self.hidden_size, self.hidden_size, std=config.initializer_range)

    def transpose_for_scores(self, x):
        new_x_shape = x.size()[:-1] + (self.num_attention_heads, self.attention_head_size)
        x = x.view(*new_x_shape)
        return x.permute(0, 2, 1, 3)

    def forward(self, q, k, v, attention_mask):
        mixed_query_layer = self.query(q)
        mixed_key_layer = self.key(k)
        mixed_value_layer = self.value(v)

        query_layer = self.transpose_for_scores(mixed_query_layer)
        key_layer = self.transpose_for_scores(mixed_key_layer)
        value_layer = self.transpose_for_scores(mixed_value_layer)

        # Take the dot product between "query" and "key" to get the raw attention scores.
        attention_scores = torch.matmul(query_layer, key_layer.transpose(-1, -2))
        attention_scores = attention_scores / math.sqrt(self.attention_head_size)
        # Apply the attention mask is (precomputed for all layers in BertModel forward() function)
        attention_scores = attention_scores + attention_mask

        # Normalize the attention scores to probabilities.
        attention_probs = nn.Softmax(dim=-1)(attention_scores)

        # This is actually dropping out entire tokens to attend to, which might
        # seem a bit unusual, but is taken from the original Transformer paper.
        attention_probs = self.dropout(attention_probs)

        context_layer = torch.matmul(attention_probs, value_layer)
        context_layer = context_layer.permute(0, 2, 1, 3).contiguous()
        new_context_layer_shape = context_layer.size()[:-2] + (self.all_head_size,)
        context_layer = context_layer.view(*new_context_layer_shape)
        attention_output = self.fc(context_layer)

        return attention_output

class Positionwise_ff(nn.Module):
    def __init__(self, config):
        super(Positionwise_ff, self).__init__()
        self.hidden_size = config.hidden_size
        self.ff_dim = config.ff_dim
        self.dropout_prob = config.dropout_prob

        self.fc_1 = Linear(self.hidden_size, self.ff_dim, std=config.initializer_range)
        self.fc_2 = Linear(self.ff_dim, self.hidden_size, std=config.initializer_range)

        self.dropout = Dropout(self.dropout_prob)
        self.act_fn = ACT2FN[config.act_fn] \
            if isinstance(config.act_fn, str) else config.act_fn

    def forward(self, x):
        x = self.dropout(self.act_fn(self.fc_1(x)))
        x = self.fc_2(x)

        return x

try:
    from apex.normalization.fused_layer_norm import FusedLayerNorm as LayerNorm
except ImportError:
    print("import basic LayerNorm version"
          "Better speed can be achieved with apex installed from https://www.github.com/nvidia/apex.")
    class LayerNorm(nn.Module):
        def __init__(self, hidden_size, eps=1e-12):
            """Construct a layernorm module in the TF style (epsilon inside the square root).
            """
            super(LayerNorm, self).__init__()
            self.weight = nn.Parameter(torch.ones(hidden_size))
            self.bias = nn.Parameter(torch.zeros(hidden_size))
            self.variance_epsilon = eps

            self.init_weights()

        def init_weights(self):
            self.weight.data.fill_(1.0)
            self.bias.data.zero_()

        def forward(self, x):
            u = x.mean(-1, keepdim=True)
            s = (x - u).pow(2).mean(-1, keepdim=True)
            x = (x - u) / torch.sqrt(s + self.variance_epsilon)
            return self.weight * x + self.bias

class Config(object):
    def __init__(self,
                 vocab_size_or_config_json_file,
                 act_fn="gelu",
                 hidden_size=768,
                 num_hidden_layers=12,
                 ff_dim=3072,
                 num_heads=12,
                 dropout_prob=0.1,
                 max_position_embeddings=512,
                 type_vocab_size=2,
                 initializer_range=0.02,
                 LN_eps=1e-12,
                 ln_pos=1
                 ):
        if isinstance(vocab_size_or_config_json_file, str):
            with open(vocab_size_or_config_json_file, "r", encoding='utf-8') as reader:
                json_config = json.loads(reader.read())
            for key, value in json_config.items():
                self.__dict__[key] = value
        elif isinstance(vocab_size_or_config_json_file, int):
            self.vocab_size = vocab_size_or_config_json_file
            self.act_fn = act_fn
            self.hidden_size = hidden_size
            self.num_hidden_layers = num_hidden_layers
            self.ff_dim = ff_dim
            self.num_heads = num_heads
            self.dropout_prob = dropout_prob
            self.max_position_embeddings = max_position_embeddings
            self.type_vocab_size = type_vocab_size
            self.LN_eps = LN_eps
            self.initializer_range = initializer_range
            self.ln_pos = ln_pos
        else:
            raise ValueError("First argument must be either a vocabulary size (int)"
                             "or the path to a pretrained model config file (str)")

    @classmethod
    def from_dict(cls, json_object):
        config = Config(vocab_size_or_config_json_file=-1)
        for key, value in json_object.items():
            config.__dict__[key] = value
        return config

    @classmethod
    def from_json_file(cls, json_file):
        with open(json_file, "r", encoding='utf-8') as reader:
            text = reader.read()
        return cls.from_dict(json.loads(text))

    def __repr__(self):
        return str(self.to_json_string())

    def to_dict(self):
        output = copy.deepcopy(self.__dict__)
        return output

    def to_json_string(self):
        return json.dumps(self.to_dict(), indent=2, sort_keys=True) + "\n"

class Embeddings(nn.Module):
    def __init__(self, config):
        super(Embeddings, self).__init__()
        self.word_embeddings = Embedding(config.vocab_size, config.hidden_size, padding_idx=0)
        self.position_embeddings = Embedding(config.max_position_embeddings, config.hidden_size)

        self.LayerNorm = LayerNorm(config.hidden_size, eps=config.LN_eps)
        self.dropout = Dropout(config.dropout_prob)

        self.init_weights(config)

    def init_weights(self, config):
        self.word_embeddings.weight.data.normal_(mean=0.0, std=config.initializer_range)
        self.position_embeddings.weight.data.normal_(mean=0.0, std=config.initializer_range)

    def forward(self, input_ids):
        seq_length = input_ids.size(1)
        position_ids = torch.arange(seq_length, dtype=torch.long, device=input_ids.device)
        position_ids = position_ids.unsqueeze(0).expand_as(input_ids)

        words_embeddings = self.word_embeddings(input_ids)
        position_embeddings = self.position_embeddings(position_ids)

        embeddings = words_embeddings + position_embeddings
        embeddings = self.LayerNorm(embeddings)
        embeddings = self.dropout(embeddings)

        return embeddings

class Encoder(nn.Module):
    def __init__(self, config):
        super(Encoder, self).__init__()
        layer = EncoderLayer(config)
        self.layers = nn.ModuleList([copy.deepcopy(layer) for _ in range(config.num_hidden_layers)])

    def forward(self, hidden_states, src_mask):
        for layer in self.layers:
            hidden_states = layer(hidden_states, src_mask)

        return hidden_states

class Decoder(nn.Module):
    def __init__(self, config, embedding_weights):
        super(Decoder, self).__init__()
        self.vocab_size = config.vocab_size
        layer = DecoderLayer(config)
        self.layers = nn.ModuleList([copy.deepcopy(layer) for _ in range(config.num_hidden_layers)])

        self.decoder = Linear(embedding_weights.size(1),
                              embedding_weights.size(0),
                              bias=False)
        self.decoder.weight = embedding_weights
        self.bias = nn.Parameter(torch.zeros(embedding_weights.size(0)))

    def forward(self, hidden_states, src, trg_mask, src_trg_mask):
        for layer in self.layers:
            hidden_states = layer(hidden_states, src, trg_mask, src_trg_mask)

        logits = self.decoder(hidden_states) + self.bias

        return logits

class EncoderLayer(nn.Module):
    def __init__(self, config):
        super(EncoderLayer, self).__init__()
        self.attn_norm = LayerNorm(config.hidden_size, eps=config.LN_eps)
        self.ffn_norm = LayerNorm(config.hidden_size, eps=config.LN_eps)

        self.self_attention = Attention(config)
        self.ffn = Positionwise_ff(config)
        self.ln_pos = config.ln_pos

    def forward(self, src, src_mask):
        # Self-Attention
        h = src
        if not self.ln_pos:
            src = self.attn_norm(src)
        src = self.self_attention(src, src, src, src_mask)
        src = h + src
        if self.ln_pos:
            src = self.attn_norm(src)

        # FFN
        h = src
        if not self.ln_pos:
            src = self.ffn_norm(src)
        src = self.ffn(src)
        src = h + src
        if self.ln_pos:
            src = self.ffn_norm(src)

        return src

class DecoderLayer(nn.Module):
    def __init__(self, config):
        super(DecoderLayer, self).__init__()
        self.attn_norm = LayerNorm(config.hidden_size, eps=config.LN_eps)
        self.m_attn_norm = LayerNorm(config.hidden_size, eps=config.LN_eps)
        self.ffn_norm = LayerNorm(config.hidden_size, eps=config.LN_eps)

        self.self_attention = Attention(config)
        self.masked_self_attention = Attention(config)

        self.ffn = Positionwise_ff(config)
        self.ln_pos = config.ln_pos

    def forward(self, trg, src, trg_mask, src_mask):
        # Masked Self-Attention
        h = trg
        if not self.ln_pos:
            trg = self.m_attn_norm(trg)
        trg = self.masked_self_attention(trg, trg, trg, trg_mask)
        trg = h + trg
        if self.ln_pos:
            trg = self.m_attn_norm(trg)

        # Self-Attention
        h = trg
        if not self.ln_pos:
            trg = self.attn_norm(trg)
        trg = self.self_attention(trg, src, src, src_mask)
        trg = h + trg
        if self.ln_pos:
            trg = self.attn_norm(trg)

        # FFN
        h = trg
        if not self.ln_pos:
            trg = self.ffn_norm(trg)
        trg = self.ffn(trg)
        trg = h + trg
        if self.ln_pos:
            trg = self.ffn_norm(trg)

        return trg

class Transformer(nn.Module):
    def __init__(self, config):
        super(Transformer, self).__init__()
        self.embedding = Embeddings(config)
        self.encoder = Encoder(config)
        self.decoder = Decoder(config, self.embedding.word_embeddings.weight)
        self.loss_fn = CrossEntropyLoss(ignore_index=0)

    def forward(self, src, trg, src_mask, trg_mask, gt=None):
        # Attention 3D Mask
        extended_src_mask = src_mask[:, None, None, :]
        #extended_src_mask = extended_src_mask.to(dtype=next(self.parameters()).dtype)  # fp16 compatibility
        extended_src_mask = (1.0 - extended_src_mask) * -10000.0

        extended_trg_mask = trg_mask[:, None, :, :]
        #extended_trg_mask = extended_trg_mask.to(dtype=next(self.parameters()).dtype) # fp16 compatibility
        extended_trg_mask = (1.0 - extended_trg_mask) * -10000.0

        # Source Embedding & Encoding
        src_emd = self.embedding(src)
        enc_src = self.encoder(src_emd, extended_src_mask)

        # Target Embedding & Decoding
        trg_emd = self.embedding(trg)
        logits = self.decoder(trg_emd, enc_src, extended_trg_mask, extended_src_mask)

        if gt is not None:
            logits = logits.contiguous().view(-1, logits.size(-1))
            gt = gt.contiguous().view(-1)
            loss = self.loss_fn(logits, gt)
            return loss
        else:
            return logits


import seaborn as sns
import matplotlib.pyplot as plt
import pandas as pd
import numpy as np

def plot_roofline_background(system, unit, max_x):
    op_intensity = system.flops/system.offchip_mem_bw
    flops = unit.raw_to_unit(system.flops, type='C')
    max_x = max(max_x, op_intensity*2.5)
    turning_points = [[0, 0], [op_intensity, flops], [max_x, flops]]
    turning_points = np.array(turning_points)
    plt.plot(turning_points[:,0], turning_points[:,1], c='grey')

    op_intensity = system.flops/system.onchip_mem_bw
    flops = unit.raw_to_unit(system.flops, type='C')
    turning_points = [[0, 0], [op_intensity, flops], [max_x, flops]]
    turning_points = np.array(turning_points)
    plt.plot(turning_points[:,0], turning_points[:,1], '--', c='grey')

    plt.xlabel('Op Intensity (FLOPs/Byte)')
    plt.ylabel(f'{unit.unit_compute.upper()}')

def dot_roofline(df, system, unit = None):
    max_x = max(df['Op Intensity'])
    plot_roofline_background(system, unit, max_x)
    for i in range(len(df)):
        op_intensity = df.loc[i, 'Op Intensity']
        thrpt = df.loc[i, 'Throughput (Tflops)']
        plt.scatter(op_intensity, thrpt, marker='o', label='Frame-'+str(i))
        if 'Scale Sim Thrpt (Tflops)' in df.columns:
           ss_thrpt = df.loc[i, 'Scale Sim Thrpt (Tflops)']
           plt.scatter(op_intensity, ss_thrpt, marker='^', label='Scale-Sim-'+str(i))
    plt.legend(loc='lower right')

def color_bound_type(value):
    if value == 'M':
        color = 'red'
    elif value == 'C':
        color = 'green'
    else:
        return
    return 'color: %s' % color

def highlight_max_cycles(s):
    '''
    highlight the maximum in a Series green.
    '''
    is_max = s == s.max()
    return ['background-color: green' if v else '' for v in is_max]

def display_df(df):
    ## Adding % of total time for each operation
    try:
        total_cycles = np.sum(df['Cycles'])
        for i in range(len(df)):
            df.loc[i, '% of total time'] = 100*df.loc[i,'Cycles']/total_cycles
    except:
        pass
    ## reducing display precision
    pd.set_option("display.precision", 3)

    ## Reordering columns
    first_cols = ['Op Type','Dimension','Op Intensity','Num ops (MFLOP)','Input_a (MB)','Input_w (MB)','Output (MB)','Total Data (MB)','Bound','Compute Cycles','Memory Cycles','C/M ratio','Cycles', '% of total time','Throughput (Tflops)']
    last_cols = [col for col in df.columns if col not in first_cols]
    df = df.loc[:,first_cols+last_cols]

    ## Applying colors and gradients to colmns
    df = df.style.background_gradient(cmap='Blues',axis=0,subset=["Cycles","Total Data (MB)"])\
        .background_gradient(cmap='Spectral_r',axis=1,subset=['Input_a (MB)','Input_w (MB)','Output (MB)'])\
        .background_gradient(cmap='Oranges',axis=0,subset=["Op Intensity"])\
        .applymap(color_bound_type, subset=['Bound'])\
        .apply(highlight_max_cycles,axis=1,subset=pd.IndexSlice[:, ['Compute Cycles','Memory Cycles']])

    display(df)
    # return df


from django import template

from questions.models import QuestionStatus

register = template.Library()

@register.filter
def question_status_class(question):
    return {
        QuestionStatus.DRAFT: 'primary',
        QuestionStatus.DEFERRED: 'warning',
        QuestionStatus.PENDING: 'info',
        QuestionStatus.APPROVED: 'success',
        QuestionStatus.REJECTED: 'danger',
        QuestionStatus.PUBLISHED: 'primary',
        QuestionStatus.ANSWERED: 'primary',
        QuestionStatus.CLOSED: 'secondary',
    }[question.status]

@register.filter
def question_status_icon(question):
    return {
        QuestionStatus.DRAFT: 'icons/clock.png',
        QuestionStatus.DEFERRED: 'icons/band.png',
        QuestionStatus.PENDING: 'icons/clock.png',
        QuestionStatus.APPROVED: 'icons/checkmark.png',
        QuestionStatus.REJECTED: 'icons/cross.png',
        QuestionStatus.PUBLISHED: 'icons/clock.png',
        QuestionStatus.ANSWERED: 'icons/clock.png',
    }.get(question.status, 'icons/question.png')


import requests
import os
import json
from dotenv import load_dotenv
from pprint import pprint

from datetime import datetime
from time import mktime
from dateutil.relativedelta import relativedelta

from pgAccess import myConnection
from htmlParser import htmlParser

class Scraper:
    def __init__(self, params: dict) -> None:
        load_dotenv()
        self.params = params
        self.ids = [33779505]
        self.connection = myConnection()

    def scrapeCars(self) -> None:
        '''
        The main method wich make all of the work - search for good cars
        and return info to google table
        '''
        # get all cars that might help us from API
        # self.get_needed_cars()

        person_dict = {} # id, user_name, phone
        ad_dict = {} # id, add_date, owners_num, last_operation, operation_description, ad_link, user_id

        for id in self.ids:
            # get from ip info about ad for bd
            info = self.get_ad_info(id)

            person_dict['id'] = info[1]

            ad_dict['id'] = id
            ad_dict['user_id'] = info[1]
            # ad_dict['add_date'] = mktime(datetime.strptime(info[0],
            #                                  '%Y-%m-%d %H:%M:%S').timetuple())

            # ad_dict['add_date'] = datetime.strptime(info[0], '%Y-%m-%d %H:%M:%S')
            ad_dict['add_date'] = info[0]
            ad_dict['ad_link'] = f"https://auto.ria.com/uk{info[2]}"

            # scrape website to get info for bd
            webpageParser = htmlParser(f"https://auto.ria.com/uk{info[2]}")

            ad_i = webpageParser.get_ad_info()

            person_dict['user_name'] = ad_i['name']
            person_dict['phone'] = webpageParser.get_phone_numbers()
            ad_dict['owners_num'] = int(ad_i['owners_num'])
            ad_dict['last_operation'] = relativedelta(datetime.now(), datetime.strptime(ad_i["last_operation"], '%d.%m.%Y')).years # difference in years
            ad_dict['operation_description'] = ad_i["operation_description"]

            print(person_dict)
            print(ad_dict)

            # load it to the db
            self.connection.insertUser(person_dict)
            self.connection.insertAd(ad_dict)

        # sort in db by addDate

        # remove unsuitable ad from bd

        # save the rest in google table

    def get_needed_cars(self):
        '''
        Function make GET request to API url and return cars corresponding to the first criteria.
        Parameters are dictionary:
            {name:value}
            which would be converted into string with &

        Return list of id`s of good cars
        '''
        page = 0

        while True:
            # add page to the request
            self.params["page"] = page
            PARAMETERS = self.param_to_string()

            # make a request
            api_url = f"https://developers.ria.com/auto/search?api_key={os.getenv('AUTORIA_API_KEY')}&{PARAMETERS}"
            response = requests.get(api_url)

            if response.ok:
                page_ids = response.json()["result"]["search_result"]
                with open("sample.json", "w") as outfile:
                    outfile.write(json.dumps(response.json(), indent=4))

                # count max of pages
                max_page = page_ids["count"] // 100

                # add all ids on all of pages
                self.ids = self.ids + page_ids["ids"]
            else:
                print("Search: Response code is not 200 OK")
            print(f"Page {page} is processed")

            if page == 0:#max_page:
                break

            page += 1

        # save id in the file (ONLY FOR TESTING)

    def param_to_string(self) -> str:
        """
        Get parameters dictionary and merge to string:
            name=value&name=value
        Return: string
        """
        answ = ''
        for key in self.params:
            answ += str(key) + '=' + str(self.params[key]) + '&'
        return answ[:-1]

    def get_ad_info(self, id) -> tuple((datetime, int, str)):
        '''
        Get info from API about add. Fill the technical_info with useful info.
        Return (addDate, user_id, linkToView)
        '''

        api_url = f"https://developers.ria.com/auto/info?api_key={os.getenv('AUTORIA_API_KEY')}&auto_id={id}"
        response = requests.get(api_url)

        if response.ok:
            ad = response.json()

            return (ad["addDate"], ad["userId"], ad["linkToView"])
        else:
            print("Ad: Response code is not 200 OK")
            return None

if __name__ == "__main__":
    params = {
        "s_yers[0]": 2018,          # рік випуску від
        "abroad": 2,                # прибрати авто, які не в Україні
        "custom": 1,                # лише розмитнені
        "matched_country": 840,     # пригнали з США
        "under_credit": 0,          # не взято в кредит
        "confiscated_car": 0,       # не конфісковані
        "sellerType": 1,            # продацель - приватна особа
        "top": 0,                   # період подачі - всі
        "saledParam": 2,            # не відображати продані
        "countpage": 100            # скільки оголошень на сторінці
    }

    s = Scraper(params)
    s.scrapeCars()

    # h = htmlParser('https://auto.ria.com/uk/auto_mercedes_benz_s_class_33779505.html')
    # h.get_ad_info()


#find leaves
# Definition for a binary tree node.
from collections import List, defaultdict

class TreeNode:
    def __init__(self, val=0, left=None, right=None):
        self.val = val
        self.left = left
        self.right = right

class Solution:
    def findLeaves(self, root: TreeNode) -> List[List[int]]:
        def order(root, dic):
            if not root:
                return 0
            left = order(root.left, dic)
            right = order(root.right, dic)
            level = max(left, right) + 1
            dic[level].append(root.val)
            return level
        dic = defaultdict(list)
        order(root, dic)
        idx = sorted(dic)
        return [dic[i] for i in idx]


# valor do dolar em relação ao real

import requests
from bs4 import BeautifulSoup

headers = {'user-agent': 'user-agent: Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, '
                         'like Gecko) Chrome/99.0.4844.82 Safari/537.36'}
page = requests.get(
    'https://www.google.com/search?q=cota%C3%A7%C3%A3o+dolar&oq=cot&aqs=chrome.0'
    '.69i59j0i131i433i512j69i57j69i59j69i65j69i60l2j69i61.1077j0j7&sourceid=chrome&ie=UTF-8', headers=headers)

soup = BeautifulSoup(page.content, 'html.parser')
atributos = {'class': 'g'}

#class="DFlfde SwHCTb"

valor_dolar = soup.find_all('span', class_ = 'DFlfde SwHCTb')[0]

cotacao_dollar = float(valor_dolar['data-value'])

#print(cotacao_dollar)


from flask import Flask, session
from pyhtml import html, title, body, p, table, tr, td, strong, head

app = Flask(__name__)

app.config['SECRET_KEY'] = 'Og5WDE6A8OidZG-fqsLgo2afWs18VIhmOVdMaYOCr3E'

@app.route('/')
def main():

    if 'visits' not in session:
        session['visits'] = 0
    session['visits'] += 1

    code = html(
        body(
            p(f"Welcome to your {session['visits']} visit.")
        )
    )
    return str(code)

if __name__ == "__main__":
    app.run(debug=True)


from network import *
from frequency import Frequency
from media import RectangularWaveguide

import os
import pylab as plb
import numpy as npy
from scipy.constants import mil
from datetime import datetime

# pre-initialized classes

#frequency bands
f_wr10 	= Frequency(75,110,201, 'ghz')
f_wr3  	= Frequency(220,325,201, 'ghz')
f_wr2p2 = Frequency(330,500,201, 'ghz')
f_wr1p5 = Frequency(500,750,201, 'ghz')
f_wr1	= Frequency(750,1100,201, 'ghz')

# rectangular waveguides
wr10 	= RectangularWaveguide(Frequency(75,110,201, 'ghz'), 100*mil,z0=50)
wr3 	= RectangularWaveguide(Frequency(220,325,201, 'ghz'), 30*mil,z0=50)
wr2p2	= RectangularWaveguide(Frequency(330,500,201, 'ghz'), 22*mil,z0=50)
wr1p5 	= RectangularWaveguide(Frequency(500,750,201, 'ghz'), 15*mil,z0=50)
wr1 	= RectangularWaveguide(Frequency(750,1100,201, 'ghz'), 10*mil,z0=50)

## Functions
# Ploting
def save_all_figs(dir = './', format=['eps','pdf','png']):
	if dir[-1] != '/':
		dir = dir + '/'
	for fignum in plb.get_fignums():
		plb.figure(fignum)
		fileName = plb.gca().get_title()
		if fileName == '':
				fileName = 'unamedPlot'
		for fmt in format:
			plb.savefig(dir+fileName+'.'+fmt, format=fmt)
			print (dir+fileName+'.'+fmt)

def legend_off(ax=None):
	'''
	turn off the legend for a given axes. if no axes is given then
	it will use current axes.
	'''
	if ax is None:
		plb.gca().legend_.set_visible(0)
	else:
		ax.lengend_.set_visible(0)

def plot_complex(z,*args, **kwargs):
	'''
	plots a complex array or list in real vs imaginary.
	'''
	plb.plot(npy.array(z).real,npy.array(z).imag, *args, **kwargs)

# other
def now_string():
	return datetime.now().__str__().replace('-','.').replace(':','.').replace(' ','.')

def find_nearest(array,value):
	'''
	find nearest value in array.
	taken from  http://stackoverflow.com/questions/2566412/find-nearest-value-in-numpy-array

	'''
	idx=(npy.abs(array-value)).argmin()
	return array[idx]

def find_nearest_index(array,value):
	'''
	find nearest value in array.
	taken from  http://stackoverflow.com/questions/2566412/find-nearest-value-in-numpy-array

	'''
	return (npy.abs(array-value)).argmin()


import sys
sys.path.append('../projects/')

import numpy as np
import pandas as pd
import scipy
from scipy import stats
from scipy.special import factorial
from scipy.stats import binom_test, wilcoxon

import os
import pickle
from datetime import datetime
import tabulate
import wandb
from collections import namedtuple, defaultdict, OrderedDict
import json
from ipypb import ipb

# from metalearning import cnnmlp

API = wandb.Api()
MAX_HISTORY_SAMPLES = 4000

DATASET_CORESET_SIZE = 22500
ACCURACY_THRESHOLD = 0.95
TASK_ACC_COLS = [f'Test Accuracy, Query #{i}' for i in range(1, 11)]

QUERY_NAMES = ['blue', 'brown', 'cyan', 'gray', 'green',
               'orange', 'pink', 'purple', 'red', 'yellow',
               'cone', 'cube', 'cylinder', 'dodecahedron', 'ellipsoid',
               'octahedron', 'pyramid', 'rectangle', 'sphere', 'torus',
               'chain_mail', 'marble', 'maze', 'metal', 'metal_weave',
               'polka', 'rubber', 'rug', 'tiles', 'wood_plank']

COLOR = 'color'
COLOR_INDEX = 0
SHAPE = 'shape'
SHAPE_INDEX = 1
TEXTURE = 'texture'
TEXTURE_INDEX = 2
COMBINED = 'combined'
COMBINED_INDEX = 3
DIMENSION_NAMES = [COLOR, SHAPE, TEXTURE]

RESULT_SET_FIELDS = ['name', 'mean', 'std']
ANALYSIS_SET_FIELDS = ['examples', 'log_examples', 'accuracies', 'accuracy_drops',
                       'first_task_accuracies', 'new_task_accuracies', 'accuracy_counts',
                       'examples_by_task']
CONDITION_ANALYSES_FIELDS = DIMENSION_NAMES + [COMBINED]
TOTAL_CURVE_FIELDS = ['raw', 'mean', 'std', 'sem']

ResultSet = namedtuple('ResultSet', RESULT_SET_FIELDS)
AnalysisSet = namedtuple('AnalysisSet', ANALYSIS_SET_FIELDS)
ConditionAnalysesSet = namedtuple('ConditionAnalysesSet', CONDITION_ANALYSES_FIELDS)
TotalCurveResults = namedtuple('TotalCurveResults', TOTAL_CURVE_FIELDS)
TaskConditionalWeights = namedtuple('TaskConditionalWeights', ('additive', 'multiplicative'))

NAMED_TUPLE_CLASSES = (ResultSet, AnalysisSet, ConditionAnalysesSet, TotalCurveResults, TaskConditionalWeights)
for NamedTupleClass in NAMED_TUPLE_CLASSES:
    NamedTupleClass.__new__.__defaults__ = (None,) * len(NamedTupleClass._fields)

RESULT_SET_NAMES = ('Examples to criterion', 'Log examples to criterion',
                    'New task accuracy', 'New task accuracy delta',
                    'First task accuracy by epoch', 'New task accuracy by epoch')
ANALYSIS_FIELDS_TO_NAMES = {field: name for (name, field) in
                            zip(ANALYSIS_SET_FIELDS, RESULT_SET_NAMES)}
ANALYSIS_NAMES_TO_FIELDS = {name: field for (name, field) in
                            zip(ANALYSIS_SET_FIELDS, RESULT_SET_NAMES)}

CACHE_PATH = './analyses_caches/meta_learning_analyses_cache.pickle'
BACKUP_CACHE_PATH = './analyses_caches/meta_learning_analyses_cache_{date}.pickle'

def refresh_cache(new_values_dict=None, cache_path=CACHE_PATH):
    if new_values_dict is None:
        new_values_dict = {}

    if os.path.exists(cache_path):
        with open(cache_path, 'rb') as cache_file:
            cache = pickle.load(cache_file)

    else:
        cache = {}

    cache.update(new_values_dict)

    if os.path.exists(cache_path):
        os.rename(CACHE_PATH, BACKUP_CACHE_PATH.format(date=datetime.now().strftime('%Y-%m-%d_%H-%M-%S')))

    with open(cache_path, 'wb') as cache_file:
        pickle.dump(cache, cache_file)

    return cache

def examples_per_epoch(task, latest_task):
    if task == latest_task:
        return 22500

    return 22500 // (latest_task - 1)

def parse_run_results(current_run_id=None, current_run=None, samples=MAX_HISTORY_SAMPLES):
    if current_run_id is None and current_run is None:
        print('Must provide either a current run or its id')
        return

    if current_run is None:
        current_run = API.run(f'meta-learning-scaling/sequential-benchmark-baseline/{current_run_id}')

    current_df = current_run.history(pandas=True, samples=samples)

    examples_to_criterion = np.empty((10, 10))
    examples_to_criterion.fill(np.nan)
    absolute_accuracy = np.empty((9, 9))
    absolute_accuracy.fill(np.nan)
    accuracy_drop = np.empty((9, 9))
    accuracy_drop.fill(np.nan)

    first_task_accuracy_by_epoch = np.empty((10, samples))
    first_task_accuracy_by_epoch.fill(np.nan)
    new_task_accuracy_by_epoch = np.empty((10, samples))
    new_task_accuracy_by_epoch.fill(np.nan)

    first_row_blank = int(np.isnan(current_df['Test Accuracy'][0]))

    first_task_finished = current_df['Test Accuracy, Query #2'].first_valid_index() - first_row_blank
    examples_to_criterion[0, 0] = first_task_finished * examples_per_epoch(1, 1)
    absolute_accuracy[0, 0] = current_df['Test Accuracy, Query #1'][first_task_finished + 1]
    accuracy_drop[0, 0] = current_df['Test Accuracy, Query #1'][first_task_finished] - absolute_accuracy[0, 0]

    first_task_accuracy_by_epoch[0, 0:first_task_finished] = current_df['Test Accuracy, Query #1'][first_row_blank:first_task_finished + 1]
    new_task_accuracy_by_epoch[0, 0:first_task_finished] = current_df['Test Accuracy, Query #1'][first_row_blank:first_task_finished + 1]

    for current_task in range(2, 11):
        current_task_start = current_df[f'Test Accuracy, Query #{current_task}'].first_valid_index()

        if current_task == 10:
            current_task_end = current_df.shape[0]
        else:
            current_task_end = current_df[f'Test Accuracy, Query #{current_task + 1}'].first_valid_index()

        current_task_subset = current_df[TASK_ACC_COLS][current_task_start:current_task_end] > 0.95

        for task in range(1, current_task + 1):
            number_times_learned = current_task - task + 1
            number_total_tasks = current_task

            examples_to_criterion[number_times_learned - 1, number_total_tasks - 1] = examples_per_epoch(task, current_task) * \
                (current_task_subset[f'Test Accuracy, Query #{task}'].idxmax() - current_task_start + 1)

            if current_task < 10:
                absolute_accuracy[number_times_learned - 1, number_total_tasks - 1] = \
                    current_df[f'Test Accuracy, Query #{task}'][current_task_end]
                accuracy_drop[number_times_learned - 1, number_total_tasks - 1] = \
                    current_df[f'Test Accuracy, Query #{task}'][current_task_end - 1] - \
                    absolute_accuracy[number_times_learned - 1, number_total_tasks - 1]

        first_task_accuracy_by_epoch[current_task - 1, 0:current_task_end - current_task_start] = current_df['Test Accuracy, Query #1'][current_task_start:current_task_end]
        new_task_accuracy_by_epoch[current_task - 1, 0:current_task_end - current_task_start] = current_df[f'Test Accuracy, Query #{current_task}'][current_task_start:current_task_end]

    return examples_to_criterion, absolute_accuracy, accuracy_drop, first_task_accuracy_by_epoch, new_task_accuracy_by_epoch

def parse_run_results_with_new_task_accuracy_and_equal_size(current_run_id=None, current_run=None, samples=MAX_HISTORY_SAMPLES):
    if current_run_id is None and current_run is None:
        print('Must provide either a current run or its id')
        return

    if current_run is None:
        current_run = API.run(f'meta-learning-scaling/sequential-benchmark-baseline/{current_run_id}')

    current_df = current_run.history(pandas=True, samples=samples)

    examples_to_criterion = np.empty((10, 10))
    examples_to_criterion.fill(np.nan)
    absolute_accuracy = np.empty((10, 10))
    absolute_accuracy.fill(np.nan)
    absolute_accuracy_equal_size = np.empty((10, 10))
    absolute_accuracy_equal_size.fill(np.nan)

    first_task_accuracy_by_epoch = np.empty((10, samples))
    first_task_accuracy_by_epoch.fill(np.nan)
    new_task_accuracy_by_epoch = np.empty((10, samples))
    new_task_accuracy_by_epoch.fill(np.nan)

    first_row_blank = int(np.isnan(current_df['Test Accuracy'][0]))

    # import pdb; pdb.set_trace()

    first_task_finished = current_df['Test Accuracy, Query #2'].first_valid_index() - first_row_blank
    examples_to_criterion[0, 0] = first_task_finished * examples_per_epoch(1, 1)
    absolute_accuracy[0, 0] = current_df['Test Accuracy, Query #1'][first_row_blank] # index 0 is all NaNs, presumably the initalization
    absolute_accuracy_equal_size[0, 0] = current_df['Test Accuracy, Query #1'][first_row_blank]

    first_task_accuracy_by_epoch[0, 0:first_task_finished] = current_df['Test Accuracy, Query #1'][first_row_blank:first_task_finished]
    new_task_accuracy_by_epoch[0, 0:first_task_finished] = current_df['Test Accuracy, Query #1'][first_row_blank:first_task_finished]

    for current_task in range(2, 11):
        current_task_start = current_df[f'Test Accuracy, Query #{current_task}'].first_valid_index()

        if current_task == 10:
            current_task_end = current_df.shape[0]
        else:
            current_task_end = current_df[f'Test Accuracy, Query #{current_task + 1}'].first_valid_index()

        current_task_subset = current_df[TASK_ACC_COLS][current_task_start:current_task_end] > 0.95

        for task in range(1, current_task + 1):
            number_times_learned = current_task - task + 1
            number_total_tasks = current_task

            examples_to_criterion[number_times_learned - 1, number_total_tasks - 1] = examples_per_epoch(task, current_task) * \
                (current_task_subset[f'Test Accuracy, Query #{task}'].idxmax() - current_task_start + 1)

            absolute_accuracy[number_times_learned - 1, number_total_tasks - 1] = \
                current_df[f'Test Accuracy, Query #{task}'][current_task_start]

            if task == current_task:
                equal_size_epochs = current_task_start
            else:
                equal_size_epochs = min(current_task_start + current_task - 1, current_task_end - 1)

            absolute_accuracy_equal_size[number_times_learned - 1, number_total_tasks - 1] = \
                current_df[f'Test Accuracy, Query #{task}'][equal_size_epochs]

        first_task_accuracy_by_epoch[current_task - 1, 0:current_task_end - current_task_start] = current_df['Test Accuracy, Query #1'][current_task_start:current_task_end]
        new_task_accuracy_by_epoch[current_task - 1, 0:current_task_end - current_task_start] = current_df[f'Test Accuracy, Query #{current_task}'][current_task_start:current_task_end]

    return examples_to_criterion, absolute_accuracy, absolute_accuracy_equal_size, first_task_accuracy_by_epoch, new_task_accuracy_by_epoch

def parse_forgetting_results(current_run_id=None, current_run=None, samples=MAX_HISTORY_SAMPLES,
                             max_samples_per_curve=31):
    if current_run_id is None and current_run is None:
        print('Must provide either a current run or its id')
        return

    if current_run is None:
        current_run = API.run(f'meta-learning-scaling/sequential-benchmark-forgetting-experiment-revisited/{current_run_id}')

    print(f'Starting to parse run {current_run.name}')
    raw_df = current_run.history(pandas=True, samples=samples)

    forgetting_trjectories = np.empty((10, 10, max_samples_per_curve))
    forgetting_trjectories.fill(np.nan)

    if 'step_resumed_from' in current_run.config:
        step_resumed_from = current_run.config['step_resumed_from']
        post_resume_step = raw_df[900:1100]['_timestamp'].diff().idxmax()
        current_df = raw_df.drop(range(step_resumed_from + 1, post_resume_step + 1), axis=0, inplace=False)
        current_df.reset_index(inplace=True)

    else:
        current_df = raw_df

    for new_task in range(2, 11):
        forgetting_start = current_df[f'Test Accuracy, Query #{new_task}'].first_valid_index() - 1

        if new_task == 10:
            forgetting_end = len(current_df) - 1
        else:
            forgetting_end = current_df[f'Test Accuracy, Query #{new_task + 1}'].first_valid_index() - 1

        forgetting_len = min(forgetting_end - forgetting_start, max_samples_per_curve)

        for forgetting_task in range(1, new_task):
            number_times_learned = new_task - forgetting_task
            number_total_tasks = new_task
            trajectory = np.array(current_df[f'Test Accuracy, Query #{forgetting_task}'][forgetting_start:forgetting_start + forgetting_len], dtype=np.float64)

            if np.any(np.isnan(trajectory)):
                print(f'Found nans for new task {new_task} and forgetting task {forgetting_task}')
                print(forgetting_start, forgetting_end)
                print(trajectory)

            forgetting_trjectories[number_times_learned - 1, number_total_tasks - 1, :forgetting_len] = trajectory

    return forgetting_trjectories

def parse_total_task_training_curve(current_run_id=None, current_run=None, samples=MAX_HISTORY_SAMPLES):
    if current_run_id is None and current_run is None:
        print('Must provide either a current run or its id')
        return

    if current_run is None:
        current_run = API.run(f'meta-learning-scaling/sequential-benchmark-baseline/{current_run_id}')

    current_df = current_run.history(pandas=True, samples=samples)
    total_task_curve = defaultdict(dict)
    task_total_trials = defaultdict(lambda: 0)

    first_row_blank = int(np.isnan(current_df['Test Accuracy'][0]))
    first_task_finished = current_df['Test Accuracy, Query #2'].first_valid_index() - first_row_blank
    episode_trials = examples_per_epoch(1, 1)

    for i, acc in enumerate(current_df['Test Accuracy, Query #1'][first_row_blank:first_task_finished]):
        total_task_curve[1][(i + 1) * episode_trials] = acc

    task_total_trials[1] = (first_task_finished - first_row_blank) * episode_trials

    for current_task in range(2, 11):
        current_task_start = current_df[f'Test Accuracy, Query #{current_task}'].first_valid_index()

        if current_task == 10:
            current_task_end = current_df.shape[0]
        else:
            current_task_end = current_df[f'Test Accuracy, Query #{current_task + 1}'].first_valid_index()

        current_task_length = current_task_end - current_task_start

        for task in range(1, current_task + 1):
            episode_trials = examples_per_epoch(task, current_task)

            for i, acc in enumerate(current_df[f'Test Accuracy, Query #{task}'][current_task_start:current_task_end]):
                total_task_curve[task][task_total_trials[task] + (i + 1) * episode_trials] = acc

            task_total_trials[task] = task_total_trials[task] + (current_task_length * episode_trials)

    # print([(t, len(total_task_curve[t].keys())) for t in range(1, 11)])

    return total_task_curve

def parse_simultaneous_training(current_run_id=None, current_run=None, samples=MAX_HISTORY_SAMPLES):
    if current_run_id is None and current_run is None:
        print('Must provide either a current run or its id')
        return

    if current_run is None:
        current_run = API.run(f'meta-learning-scaling/sequential-benchmark-baseline/{current_run_id}')

    current_df = current_run.history(pandas=True, samples=samples)
    simultaneous_task_accuracies = current_df['Test Per-Query Accuracy (list)']
    first_row_blank = int(np.isnan(simultaneous_task_accuracies[0]))
    return np.array(list(simultaneous_task_accuracies[first_row_blank:]))

def process_multiple_runs_simultaneous_training(runs, debug=False, ignore_runs=None, samples=MAX_HISTORY_SAMPLES):
    simultaneous_training_results = []

    for i, run in enumerate(runs):
        if i > 0 and i % 10 == 0:
            print(run.name, i)
        else:
            print(run.name)

        if ignore_runs is not None and run.name in ignore_runs:
            continue

        simultaneous_training_results.append(parse_simultaneous_training(current_run=run, samples=samples))

    max_length = max([res.shape[0] for res in simultaneous_training_results])
    for i in range(len(simultaneous_training_results)):
        curr_length = simultaneous_training_results[i].shape[0]
        simultaneous_training_results[i] = np.pad(simultaneous_training_results[i], ((0, max_length - curr_length), (0, 0)),
                                                  'edge') #'constant', constant_values=np.nan)

    stacked_results = np.stack(simultaneous_training_results, axis=2)

    all_results_mean = np.nanmean(stacked_results, axis=(1, 2))
    all_results_std = np.nanstd(stacked_results, axis=(1, 2))
    all_results_sem = all_results_std / np.sqrt(np.prod(stacked_results.shape[1:]))

    per_task_results_mean = np.nanmean(stacked_results, axis=(2))
    per_task_results_std = np.nanstd(stacked_results, axis=(2))
    per_task_results_sem = per_task_results_std / np.sqrt(stacked_results.shape[2])

    return stacked_results, all_results_mean, all_results_std, all_results_sem, per_task_results_mean, per_task_results_std, per_task_results_sem

PRINT_HEADERS = ['###'] + [str(x) for x in range(1, 11)]

def pretty_print_results(results, **kwargs):
    result_rows = [[str(i + 1)] + list(results[i]) for i in range(len(results))]
    tab_args = dict(tablefmt='fancy_grid')
    tab_args.update(kwargs)
    print(tabulate.tabulate(result_rows, PRINT_HEADERS, **tab_args))

DEFAULT_PROJECT_PATH = 'meta-learning-scaling/sequential-benchmark-baseline'

def load_runs(max_rep_id, project_path=DEFAULT_PROJECT_PATH, split_runs_by_dimension=True,
              valid_run_ids=None, debug=False):
    runs = API.runs(project_path)

    results = ConditionAnalysesSet([], [], [], [])

    for run in runs:
        config = json.loads(run.json_config)
        run_id = config['dataset_random_seed']['value']

        if isinstance(run_id, dict):
            run_id = run_id['value']

        run_id = int(run_id)

        if debug: print(run.name, run_id)

        # run_name = run.description.split('\n')[0]
        # run_id = int(run_name[run_name.rfind('-') + 1:])

        if valid_run_ids is None or run_id in valid_run_ids:
            dimension = (run_id // 1000) - 1
            rep = run_id % 1000
            if rep < max_rep_id:
                # combined / all runs
                results[COMBINED_INDEX].append(run)
                # by-dimension
                if split_runs_by_dimension:
                        results[dimension].append(run)

    return results

def query_modulated_runs_by_dimension(max_rep_id):
    runs = API.runs('meta-learning-scaling/sequential-benchmark-task-modulated')

    results = {level: ConditionAnalysesSet([], [], [], []) for level in range(1, 5)}

    for run in runs:
        level, run_id = [int(x) for x in run.description.split('\n')[0][-6:].split('-')]
        dimension = (run_id // 1000) - 1
        rep = run_id % 1000
        if rep < max_rep_id:
            results[level][dimension].append(run)
            # combined / all runs
            results[level][3].append(run)

    return results

def epoch_to_trials(name, episode):
    if episode == 1 and 'power' not in name.lower():
        return 22500

    return 45000

def epochs_to_task_completions(runs, name=None, ignore_runs=None, samples=500, ipb_desc=None):
    if name is None:
        name = 'epochs_to_completion'

    results_per_run = []
    last_tasks_per_run = []

    for run in ipb(runs, desc=ipb_desc):
#         if i > 0 and i % 10 == 0:
#             print(run.name, i)
#         else:
#             print(run.name)

        if ignore_runs is not None and run.name in ignore_runs:
            continue

        df = run.history(pandas=True, samples=samples)
        first_row_blank = int(np.isnan(df['Test Accuracy'][0]))
        first_task_finished = df['Test Accuracy, Query #2'].first_valid_index() - first_row_blank

        task_finishes = [first_task_finished]
        last_tasks = [1]

        for current_task in range(2, 11):
            current_task_start = df[f'Test Accuracy, Query #{current_task}'].first_valid_index()

            if current_task == 10:
                current_task_end = df.shape[0]
            else:
                current_task_end = df[f'Test Accuracy, Query #{current_task + 1}'].first_valid_index()

            task_finishes.append(current_task_end)

            sub_df = df[[f'Test Accuracy, Query #{i + 1}' for i in range(current_task)]]
            min_column_name = sub_df.iloc[current_task_end - 2].idxmin()
            last_task_to_finish = int(min_column_name[min_column_name.find('#') + 1:])
            last_tasks.append(last_task_to_finish)

        results_per_run.append(task_finishes)
        last_tasks_per_run.append(last_tasks)

    results_per_run = np.array(results_per_run)
    trials_per_epoch = np.array([epoch_to_trials(name, e) for e in np.arange(1, 11)])
    trials_per_run = results_per_run * trials_per_epoch

    return ResultSet(name=name, mean=np.nanmean(results_per_run, axis=0), std=np.nanstd(results_per_run, axis=0)),\
            ResultSet(name=name, mean=np.nanmean(np.log(results_per_run), axis=0), std=np.nanstd(np.log(results_per_run), axis=0)),\
            ResultSet(name=name, mean=np.nanmean(trials_per_run, axis=0), std=np.nanstd(trials_per_run, axis=0)),\
            ResultSet(name=name, mean=np.nanmean(np.log(trials_per_run), axis=0), std=np.nanstd(np.log(trials_per_run), axis=0)),\
            last_tasks_per_run

def kth_diag_indices(a, k):
    rows, cols = np.diag_indices_from(a)
    if k < 0:
        return rows[-k:], cols[:k]
    elif k > 0:
        return rows[:-k], cols[k:]
    else:
        return rows, cols

def process_multiple_runs(runs, debug=False, ignore_runs=None, samples=MAX_HISTORY_SAMPLES, parse_func=parse_run_results):
    examples = []
    log_examples = []
    abs_accuracies = []
    accuracy_drops = []
    first_task_accuracies = []
    new_task_accuracies = []

    examples_by_task = np.zeros((30, 10, 10))
    log_examples_by_task = np.zeros((30, 10, 10))
    counts_by_task = np.zeros((30, 10, 10))

    for i, run in enumerate(runs):
        if i > 0 and i % 10 == 0:
            print(run.id, run.name, i)
        else:
            print(run.id, run.name)

        if ignore_runs is not None and run.name in ignore_runs:
            continue

        examples_to_criterion, absolute_accuracy, accuracy_drop, first_task_acc, new_task_acc = parse_func(current_run=run, samples=samples)
        # print(examples_to_criterion)
        # print(np.log(examples_to_criterion))
        examples.append(examples_to_criterion)
        log_examples.append(np.log(examples_to_criterion))
        abs_accuracies.append(absolute_accuracy)
        accuracy_drops.append(accuracy_drop)
        first_task_accuracies.append(first_task_acc)
        new_task_accuracies.append(new_task_acc)

        for index, task in enumerate(run.config['query_order']):
            task_examples = np.diag(examples_to_criterion, index)
            diag_r, diag_c = kth_diag_indices(examples_to_criterion, index)
            # examples_by_task[task,:10 - index] += task_examples
            # counts_by_task[task,:10 - index] += 1
            examples_by_task[task, diag_r, diag_c] += task_examples
            log_examples_by_task[task, diag_r, diag_c] += np.log(task_examples)
            counts_by_task[task, diag_r, diag_c] += 1

    # Removing all extraneous nans
    print('Removing extraneous nans')
    first_task_accuracies = np.array(first_task_accuracies)
    new_task_accuracies = np.array(new_task_accuracies)

    max_first_nan_idx = np.max(np.argmax(np.isnan(first_task_accuracies), axis=2))
    print('Max first nan index:', max_first_nan_idx)

    first_task_accuracies = first_task_accuracies[:, :, :max_first_nan_idx]
    new_task_accuracies = new_task_accuracies[:, :, :max_first_nan_idx]
    accuracy_counts = np.count_nonzero(~np.isnan(new_task_accuracies), axis=0)

    output = {}
    for result_set, name, field in zip((examples, log_examples, abs_accuracies, accuracy_drops,
                                        first_task_accuracies, new_task_accuracies),
                                       RESULT_SET_NAMES,
                                       ANALYSIS_SET_FIELDS):
        print(name, field)
        output[field] = ResultSet(name=name,
                                  mean=np.nanmean(result_set, axis=0),
                                  std=np.nanstd(result_set, axis=0))

    # eliminate zeros to allow dividing later
    accuracy_counts[accuracy_counts == 0] = 1
    output['accuracy_counts'] = accuracy_counts

    # to avoid division by zero
    counts_by_task[counts_by_task == 0] = 1
    average_examples_by_task = np.divide(examples_by_task, counts_by_task)
    average_log_examples_by_task = np.divide(log_examples_by_task, counts_by_task)
    output['examples_by_task'] = dict(raw=average_examples_by_task, log=average_log_examples_by_task)

    analysis = AnalysisSet(**output)

    if debug:
        return analysis, examples

    return analysis

def process_multiple_runs_total_task_training_curves(runs, debug=False, ignore_runs=None, samples=MAX_HISTORY_SAMPLES):
    aggregate_total_task_curve = defaultdict(lambda: defaultdict(list))

    for i, run in enumerate(runs):
        if i > 0 and i % 10 == 0:
            print(run.name, i)
        else:
            print(run.name)

        if ignore_runs is not None and run.name in ignore_runs:
            continue

        total_task_curve = parse_total_task_training_curve(current_run=run, samples=samples)
        for task in total_task_curve:
            for num_trials, acc in total_task_curve[task].items():
                aggregate_total_task_curve[task][num_trials].append(acc)

    aggregate_total_task_curve_raw = defaultdict(OrderedDict)
    aggregate_total_task_curve_mean = defaultdict(OrderedDict)
    aggregate_total_task_curve_std = defaultdict(OrderedDict)
    aggregate_total_task_curve_sem = defaultdict(OrderedDict)

    for task in range(1, 11):
        for num_trials in sorted(aggregate_total_task_curve[task]):
            accs = aggregate_total_task_curve[task][num_trials]

            aggregate_total_task_curve_raw[task][num_trials] = accs
            aggregate_total_task_curve_mean[task][num_trials] = np.nanmean(accs)

            std = np.nanstd(accs)
            if np.isnan(std):
                print(f'For accs {accs}, std is nan')

            aggregate_total_task_curve_std[task][num_trials] = std
            aggregate_total_task_curve_sem[task][num_trials] = std / (len(accs) ** 0.5)

    return aggregate_total_task_curve_raw, aggregate_total_task_curve_mean, aggregate_total_task_curve_std, aggregate_total_task_curve_sem

def sign_test(values):
    n = len(values)
    row_faster_results = np.empty((n, n))
    row_faster_results.fill(np.nan)

    col_faster_results = np.empty((n, n))
    col_faster_results.fill(np.nan)

    wilcoxon_statistics = np.empty((n, n))
    wilcoxon_statistics.fill(np.nan)

    wilcoxon_p_values = np.empty((n, n))
    wilcoxon_p_values.fill(np.nan)

    for row in range(n):
        for col in range(row + 1, n):
            row_means = values[row].mean
            row_means = row_means[~np.isnan(row_means)]
            col_means = values[col].mean
            col_means = col_means[~np.isnan(col_means)]

            # Run both explicit comparisons, to account for the possibility of ties
            row_faster = np.sum(row_means <col_means )
            row_faster_results[row, col] = row_faster

            col_faster = np.sum(row_means > col_means)
            col_faster_results[row, col] = col_faster

            s, p = wilcoxon(row_means, col_means)
            wilcoxon_statistics[row, col] = s
            wilcoxon_p_values[row, col] = p

    return row_faster_results, col_faster_results, wilcoxon_statistics, wilcoxon_p_values

MODULATION_LEVELS = ['\\thead[cl]{None}'] + [f'\\thead[cl]{{ {i} }}' for i in range(1, 5)]
SIGN_TEST_PRINT_HEADERS = ['\\thead[cl]{Modulation level}'] + [f'\\thead[cl]{{ {i} }}' for i in range(1, 5)]

def pretty_print_sign_test_results(row_faster_results, col_faster_results,
                                   wilcoxon_statistics=None, wilcoxon_p_values=None,
                                   higher_better=False, **kwargs):
    n = row_faster_results.shape[0]
    print_results = [[''] * n for _ in range(n)]
    if wilcoxon_statistics is not None:
        wilcoxon_results = [[''] * n for _ in range(n)]

    for row in range(n):
        for col in range(row + 1, n):
            row_val = int(row_faster_results[row, col])
            col_val = int(col_faster_results[row, col])
            if higher_better:
                row_val, col_val = col_val, row_val

            result = col_val # max(col_val, row_val)
            n_binom = col_val + row_val
            p = binom_test(max(col_val, row_val), n_binom)

            print_results[row][col] = f'\\makecell[cl]{{ {result} $(n={n_binom})$ \\\\ $p={p:.4f}{(p < 0.05 and p > 0.01) and "^{*}" or ""}{p < 0.01 and "^{**}" or ""}$}}'

            if wilcoxon_statistics is not None:
                wilcoxon_results[row][col] = f'\n{wilcoxon_statistics[row, col]:.4f}, {wilcoxon_p_values[row, col]:.4f}'

    tab_args = dict(tablefmt='fancy_grid')
    tab_args.update(kwargs)

    print_results = [[MODULATION_LEVELS[i]] + row[1:] for i, row in enumerate(print_results)]
    print(tabulate.tabulate(print_results[:-1], SIGN_TEST_PRINT_HEADERS, **tab_args))

    if wilcoxon_statistics is not None:
        wilcoxon_results = [[MODULATION_LEVELS[i]] + row[1:] for i, row in enumerate(wilcoxon_results)]
        print(tabulate.tabulate(wilcoxon_results[:-1], SIGN_TEST_PRINT_HEADERS, **tab_args))

def sign_test_with_sem(values, sample_sizes):
    n = len(values)
    row_faster_results = np.empty((n, n))
    row_faster_results.fill(np.nan)

    col_faster_results = np.empty((n, n))
    col_faster_results.fill(np.nan)

    wilcoxon_statistics = np.empty((n, n))
    wilcoxon_statistics.fill(np.nan)

    wilcoxon_p_values = np.empty((n, n))
    wilcoxon_p_values.fill(np.nan)

    for row in range(n):
        for col in range(row, n):
            row_means = values[row].mean
            row_means = row_means[~np.isnan(row_means)]
            row_sems = values[row].std / (sample_sizes[row] ** 0.5)
            row_sems = row_sems[~np.isnan(row_sems)]

            col_means = values[col].mean
            col_means = col_means[~np.isnan(col_means)]
            col_sems = values[col].std / (sample_sizes[col] ** 0.5)
            col_sems = col_sems[~np.isnan(col_sems)]

            # Run both explicit comparisons, to account for the possibility of ties
            row_faster_results[row, col] = np.sum(row_means + row_sems < col_means - col_sems)
            col_faster_results[row, col] = np.sum(row_means - row_sems > col_means + col_sems)

            signs = row_means < col_means
            signs[signs == 0] = -1

            shifted_row_means = row_means + np.multiply(row_sems, signs)
            shifted_col_means = col_means + np.multiply(col_sems, -1 * signs)
            s, p = wilcoxon(shifted_row_means, shifted_col_means)
            wilcoxon_statistics[row, col] = s
            wilcoxon_p_values[row, col] = p

    return row_faster_results, col_faster_results, wilcoxon_statistics, wilcoxon_p_values

DEFAULT_LEARNING_RATE = 5e-4
DEFAULT_WEIGHT_DECAY = 1e-4

def create_task_conditional_model(additive=True, multiplicative=True, checkpoint_path=None, name=None):
    mod_level = list(range(4))

    model = cnnmlp.TaskConditionalCNNMLP(
        mod_level=mod_level,
        multiplicative_mod=multiplicative,
        additive_mod=additive,
        query_length=30,
        conv_filter_sizes=(16, 32, 48, 64),
        conv_output_size=4480,
        mlp_layer_sizes=(512, 512, 512, 512),
        lr=DEFAULT_LEARNING_RATE,
        weight_decay=DEFAULT_WEIGHT_DECAY,
        use_lr_scheduler=False,
        conv_dropout=False,
        mlp_dropout=False,
        name=name)

    if checkpoint_path is not None:
        model.load_state(checkpoint_path)

    return model

def parse_task_conditional_weights(runs, additive=True, multiplicative=True,
                                   layers=range(4), download_root='/tmp', ignore_runs=None):
    weights_by_run_id = {}

    for i, run in enumerate(runs):
        if ignore_runs is not None and run.name in ignore_runs:
            continue

        j = i + 1
        if j % 10 == 0:
            print(run.name, j)
        else:
            print(run.name)

        dimension = run.config['benchmark_dimension']

        last_checkpoint_file = run.file(f'{run.name.replace("[0, 1, 2, 3]-", "")}-query-9.pth')
        last_checkpoint_file.download(replace=True, root=download_root)
        model = create_task_conditional_model(additive=additive, multiplicative=multiplicative,
                                             checkpoint_path=os.path.join(download_root, last_checkpoint_file.name))

        additive_weights = None
        multiplicative_weights = None

        if additive:
            additive_weights = [model.conv.additive_mod_layers[f'additive-{i}'].weight.detach().cpu().numpy()[:, dimension * 10:(dimension + 1) * 10]
                                for i in layers]

        if multiplicative:
            multiplicative_weights = [model.conv.multiplicative_mod_layers[f'multiplicative-{i}'].weight.detach().cpu().numpy()[:, dimension * 10:(dimension + 1) * 10]
                                      for i in layers]

        weights_by_run_id[run.config['latin_square_index']] = TaskConditionalWeights(additive=additive_weights, multiplicative=multiplicative_weights)

    return weights_by_run_id


from __future__ import print_function
import abc
import os
import math

import numpy as np
import logging
import torch
import torch.utils.data
from torch import nn
from torch.nn import init
from torch.nn import functional as F
from torch.autograd import Variable

class AbstractAutoEncoder(nn.Module):
    __metaclass__ = abc.ABCMeta

    @abc.abstractmethod
    def encode(self, x):
        return

    @abc.abstractmethod
    def decode(self, z):
        return

    @abc.abstractmethod
    def forward(self, x):
        """model return (reconstructed_x, *)"""
        return

    @abc.abstractmethod
    def sample(self, size):
        """sample new images from model"""
        return

    @abc.abstractmethod
    def loss_function(self, **kwargs):
        """accepts (original images, *) where * is the same as returned from forward()"""
        return

    @abc.abstractmethod
    def latest_losses(self):
        """returns the latest losses in a dictionary. Useful for logging."""
        return

class ResBlock(nn.Module):
    def __init__(self, in_channels, out_channels, mid_channels=None, bn=False):
        super(ResBlock, self).__init__()

        if mid_channels is None:
            mid_channels = out_channels

        layers = [
            nn.LeakyReLU(),
            nn.Conv2d(in_channels, mid_channels, kernel_size=3, stride=1, padding=1),
            nn.LeakyReLU(),
            nn.Conv2d(mid_channels, out_channels, kernel_size=1, stride=1, padding=0)]
        if bn:
            layers.insert(2, nn.BatchNorm2d(out_channels))
        self.convs = nn.Sequential(*layers)

    def forward(self, x):
        return x + self.convs(x)

class CVAE(AbstractAutoEncoder):
    def __init__(self, d, z, **kwargs):
        super(CVAE, self).__init__()

        self.encoder = nn.Sequential(
            nn.Conv2d(3, d // 2, kernel_size=4, stride=2, padding=1, bias=False),
            nn.BatchNorm2d(d // 2),
            nn.ReLU(inplace=True),
            nn.Conv2d(d // 2, d, kernel_size=4, stride=2, padding=1, bias=False),
            nn.BatchNorm2d(d),
            nn.ReLU(inplace=True),
            ResBlock(d, d, bn=True),
            nn.BatchNorm2d(d),
            ResBlock(d, d, bn=True),
        )

        self.decoder = nn.Sequential(
            ResBlock(d, d, bn=True),
            nn.BatchNorm2d(d),
            ResBlock(d, d, bn=True),
            nn.BatchNorm2d(d),

            nn.ConvTranspose2d(d, d // 2, kernel_size=4, stride=2, padding=1, bias=False),
            nn.BatchNorm2d(d // 2),
            nn.LeakyReLU(inplace=True),
            nn.ConvTranspose2d(d // 2, 3, kernel_size=4, stride=2, padding=1, bias=False),
        )
        self.xi_bn = nn.BatchNorm2d(3)

        self.f = 8
        self.d = d
        self.z = z
        self.fc11 = nn.Linear(d * self.f ** 2, self.z)
        self.fc12 = nn.Linear(d * self.f ** 2, self.z)
        self.fc21 = nn.Linear(self.z, d * self.f ** 2)

    def encode(self, x):
        h = self.encoder(x)
        h1 = h.view(-1, self.d * self.f ** 2)
        return h, self.fc11(h1), self.fc12(h1)

    def reparameterize(self, mu, logvar):
        if self.training:
            std = logvar.mul(0.5).exp_()
            eps = std.new(std.size()).normal_()
            return eps.mul(std).add_(mu)
        else:
           return mu

    def decode(self, z):
        z = z.view(-1, self.d, self.f, self.f)
        h3 = self.decoder(z)
        return torch.tanh(h3)

    def forward(self, x):

        _, mu, logvar = self.encode(x)
        hi = self.reparameterize(mu, logvar)
        hi_projected = self.fc21(hi)
        xi = self.decode(hi_projected)
        xi = self.xi_bn(xi)

        return xi

class CVAE_imagenet(nn.Module):
    def __init__(self, d, k=10, num_channels=3, **kwargs):
        super(CVAE_imagenet, self).__init__()

        self.encoder = nn.Sequential(
            nn.Conv2d(num_channels, d, kernel_size=4, stride=2, padding=1),
            nn.BatchNorm2d(d),
            nn.LeakyReLU(inplace=True),
            nn.Conv2d(d, d, kernel_size=4, stride=2, padding=1),
            nn.BatchNorm2d(d),
            nn.LeakyReLU(inplace=True),
            ResBlock(d, d),
            nn.BatchNorm2d(d),
            ResBlock(d, d),
            nn.BatchNorm2d(d),
        )
        self.decoder = nn.Sequential(
            ResBlock(d, d),
            nn.BatchNorm2d(d),
            ResBlock(d, d),
            nn.ConvTranspose2d(d, d, kernel_size=4, stride=2, padding=1),
            nn.BatchNorm2d(d),
            nn.LeakyReLU(inplace=True),
            nn.ConvTranspose2d(d, num_channels, kernel_size=4, stride=2, padding=1),
        )
        self.d = d
        self.emb = NearestEmbed(k, d)

        for l in self.modules():
            if isinstance(l, nn.Linear) or isinstance(l, nn.Conv2d):
                l.weight.detach().normal_(0, 0.02)
                torch.fmod(l.weight, 0.04)
                nn.init.constant_(l.bias, 0)

        self.encoder[-1].weight.detach().fill_(1 / 40)

        self.emb.weight.detach().normal_(0, 0.02)
        torch.fmod(self.emb.weight, 0.04)
        self.L_bn = nn.BatchNorm2d(num_channels)

    def encode(self, x):
        return self.encoder(x)

    def decode(self, x):
        return torch.tanh(self.decoder(x))

    def forward(self, x):

        z_e = self.encode(x)

        z_q, _ = self.emb(z_e, weight_sg=True)
        emb, _ = self.emb(z_e.detach())

        l = self.decode(z_q)
        xi = self.L_bn(l)

        return  xi