import numpy as np
import utils
import optimize

class Experiment:
     def __init__(self,df,num_roles,num_segments,algo_ver,dest,tuning_params,num_levels=None,refValue = 0):
         self.df=df
         self.num_roles=num_roles
         self.num_segments=num_segments
         self.algo_ver = algo_ver
         self.dest = dest
         self.num_levels = num_levels
         self.refValue = refValue
         self.tuning_params =  tuning_params
         
     def execute(self):
        # Extract nodes
        nodes_arr = np.union1d(self.df['target'],self.df['source']).astype(int) 
        # list of nodes         
        nodes = nodes_arr.tolist()
        num_vertices = len(nodes)
        
        # Extract graph
        G = utils.getGraph(nodes_arr,self.df.values)
        
        # User parameters
        num_roles=self.num_roles
        num_segments=self.num_segments
        # Time stamps
        t_df = self.df["timestamp"]
        t_df = sorted(t_df)    
        print('TIME')
        print(len(t_df))               
        # initilaize group assignment randomly
        group_assignment= np.random.randint(num_roles, size=(num_vertices))                   
        # node-group dictionary
        group_dic = {}
        keys = nodes
        values = list(group_assignment)
        group_dic = dict(zip(keys,values))
        
        # create a new dictionary - key: node-pair , value:  list of timestamps
        dic=self.df.groupby(['source','target'])['timestamp'].apply(list).to_dict()
        
        print(len(dic.values()))
        
        # initialize change points
        change_points_arr = np.zeros((num_roles, num_roles, num_segments+1) , dtype=int)
               
        #Initialize  lamda
        lambda_estimates = np.zeros((num_roles, num_roles,num_segments) , dtype=float)
        
        global itr
        
        ###      K-segmentation         ###        
        if self.algo_ver == 1:
            opt = optimize.Optimize( group_dic,lambda_estimates,change_points_arr,nodes,num_roles,num_segments,dic,None,self.tuning_params)    
            [group_dic,lambda_estimates,change_points_arr,likelihood] = opt.k_seg()

        ### (K,H)-segmentation variant-1 ###                        
        elif self.algo_ver == 2:
            opt = optimize.Optimize( group_dic,lambda_estimates,change_points_arr,nodes,num_roles,num_segments,dic, self.num_levels,self.tuning_params)    
            [group_dic,lambda_estimates,change_points_arr,likelihood] = opt.k_h_seg_var_1()  

        ### (K,H)-segmentation variant-2 ###
        elif self.algo_ver == 3:  
            
            opt = optimize.Optimize( group_dic,lambda_estimates,change_points_arr,nodes,num_roles,num_segments,dic,self.num_levels,self.tuning_params)    
            [group_dic,lambda_estimates,change_points_arr,likelihood,itr] = opt.k_h_seg_var_2()   
            print (itr)
        
        ### Level-dependent (K,H)-segmentation variant-2  ###
        elif self.algo_ver == 4:              
            opt = optimize.Optimize( group_dic,lambda_estimates,change_points_arr,nodes,num_roles,num_segments,dic,self.num_levels,self.tuning_params)    
            [group_dic,lambda_estimates,change_points_arr,likelihood,g_mapping,itr] = opt.mm_k_h_seg_var_2()  
            print (itr)
            # print('g_mapping_discoverd {}'.format(g_mapping))
            # for e_h in range(0,self.num_levels):
            #     g_a = group_dic[e_h]
            #     list_of_groups=  [[] for _ in range(num_roles)]
            #     for idx, val in g_a.items():
            #         list_of_groups[val].append(idx)
            #     print('group assignments {}: {}'.format(e_h,list_of_groups)) 
        # Plotting
        # dest_folder= self.dest  + str(self.algo_ver)+'/'
        # utils.generate_plots(G,group_dic,lambda_estimates,num_roles,num_segments,dest_folder,nodes,change_points_arr,t_df,self.refValue)
        
        # list_of_groups=  [[] for _ in range(num_roles)]
        # for idx, val in group_dic.items():
        #     list_of_groups[val].append(idx)
        # print('group assignments: {}'.format(list_of_groups))            
        # print('lambdas: {}'.format(lambda_estimates))
        # return [likelihood,group_dic,lambda_estimates,change_points_arr]
        # return [likelihood,group_dic]
        return [itr,likelihood,group_dic,lambda_estimates,change_points_arr]
        # return [_itr,likelihood]