import sbm_core
import math
import numpy as np

#########################################################################################################################
### Algorithms for Maximum Likelihood Estimation of the intensity functions and communities in stochastic blockmodels ###
#########################################################################################################################


class Optimize:
    
    def __init__(self, group_assignment,lambda_estimates,change_points_arr,nodes,num_roles,num_segments,df,num_levels,tuning_params):
       self.group_assignment = group_assignment
       self.lambda_estimates = lambda_estimates
       self.change_points_arr = change_points_arr
       self.nodes = nodes
       self.num_roles= num_roles
       self.num_segments =num_segments
       self.df=df
       self.num_levels=num_levels
       self.tuning_params = tuning_params
       self.g_mapping =  None
       
    def k_seg(self):        
        """K-segmentation"""

        liklihood_sum = 0
        #  initialize lambdas randomly
        # self.num_levels = self.num_segments
        # self.com_h_lvl_lambda()
        self.lambda_estimates=sbm_core.estimate_lamda(self.group_assignment,self.lambda_estimates,self.change_points_arr,self.num_roles,self.num_segments,self.df,self.tuning_params)
        liklihood_sum = sbm_core.compute_cost(self.group_assignment,self.lambda_estimates,self.change_points_arr,self.num_roles,self.num_segments,self.df)
        
        _prev_val = math.inf 
        _curr_val = liklihood_sum
        _itr = 0
        
        while round(_prev_val,2) != round(_curr_val,2):
            print("iteration no........... %d " %(_itr+1)) 
           
            self.group_assignment=sbm_core.assign_groups(self.group_assignment,self.lambda_estimates,self.change_points_arr,self.nodes,self.num_roles,self.num_segments,self.df)
            # self.group_assignment=sbm_core.group_assignment_ver2_2(self.nodes,self.num_roles,self.num_segments,self.lambda_estimates,self.group_assignment,self.change_points_arr,self.df)
            print('after grouping')
            liklihood_sum = sbm_core.compute_cost(self.group_assignment,self.lambda_estimates,self.change_points_arr,self.num_roles,self.num_segments,self.df)
           
            self.lambda_estimates=sbm_core.estimate_lamda(self.group_assignment,self.lambda_estimates,self.change_points_arr,self.num_roles,self.num_segments,self.df,self.tuning_params)
            print('after lambda estimate')
            liklihood_sum = sbm_core.compute_cost(self.group_assignment,self.lambda_estimates,self.change_points_arr,self.num_roles,self.num_segments,self.df)
            
            print('after seg')
            self.change_points_arr = sbm_core.dyn_prog_seg(self.group_assignment,self.lambda_estimates,self.change_points_arr,self.num_roles,self.num_segments,self.df)    
            liklihood_sum = sbm_core.compute_cost(self.group_assignment,self.lambda_estimates,self.change_points_arr,self.num_roles,self.num_segments,self.df)                    
            self.lambda_estimates=sbm_core.estimate_lamda(self.group_assignment,self.lambda_estimates,self.change_points_arr,self.num_roles,self.num_segments,self.df,self.tuning_params)
            print('after lambda estimate')
            liklihood_sum = sbm_core.compute_cost(self.group_assignment,self.lambda_estimates,self.change_points_arr,self.num_roles,self.num_segments,self.df)
   
            _prev_val = _curr_val
            _curr_val = liklihood_sum
            _itr+=1
        return [self.group_assignment,self.lambda_estimates,self.change_points_arr,liklihood_sum]
    
    def com_h_lvl_lambda(self):
        H = self.num_levels
        print('k-h levels %d'%(self.num_levels))
                
        # h-level lambda estimates    
        lambda_estimates_h = np.random.rand(self.num_roles, self.num_roles, H)
        # print(lambda_estimates_h)
                
        l1 =list(range(0, H))
        l2 = []
        if self.num_segments > self.num_levels:
            l2 = [np.random.randint(0,H) for i in range(self.num_segments-H)]  
        
        # Mapping from segment to a level
        g_mapping= np.array(l1 + l2)
        # print('g mapping {}'.format(g_mapping))
               
        for d in range(0, self.num_segments):
            self.lambda_estimates[:,:, d]= lambda_estimates_h[:,:,g_mapping[d]]

    def k_h_seg_var_1(self):
        """(K,H)-segmentation variant-1"""
        
        
        # initilaization algorithm: initialize lambdas randomly and segments through linear seg. ver 2.            
        self.com_h_lvl_lambda()
                    
        [self.lambda_estimatess,self.change_points_arr]=sbm_core.linear_seg_ver_2(self.num_roles,self.num_segments,self.group_assignment,self.lambda_estimates,self.change_points_arr,self.df)
        self.lambda_estimates=sbm_core.estimate_lamda_kh(self.num_roles,self.num_segments,self.lambda_estimates,self.group_assignment,self.change_points_arr,self.df,self.tuning_params)        
        liklihood_sum = sbm_core.compute_cost(self.group_assignment,self.lambda_estimates,self.change_points_arr,self.num_roles,self.num_segments,self.df)
        _prev_val = math.inf 
        _curr_val = liklihood_sum
        _itr = 0        

        while round(_prev_val,2) != round(_curr_val,2):
            print("iteration no........... %d " %(_itr+1)) 
            # self.group_assignment=sbm_core.assign_groups(self.group_assignment,self.lambda_estimates,self.change_points_arr,self.nodes,self.num_roles,self.num_segments,self.df)
            self.group_assignment=sbm_core.group_assignment_ver2_2(self.nodes,self.num_roles,self.num_segments,self.lambda_estimates,self.group_assignment,self.change_points_arr,self.df)
            # print('after grouping')
            # liklihood_sum = sbm_core.compute_cost(self.group_assignment,self.lambda_estimates,self.change_points_arr,self.num_roles,self.num_segments,self.df)
            
            self.lambda_estimates=sbm_core.estimate_lamda_kh(self.num_roles,self.num_segments,self.lambda_estimates,self.group_assignment,self.change_points_arr,self.df,self.tuning_params)
            # print('after lambda estimate')
            # liklihood_sum = sbm_core.compute_cost(self.group_assignment,self.lambda_estimates,self.change_points_arr,self.num_roles,self.num_segments,self.df)

            # print('after seg')
            self.change_points_arr = sbm_core.dyn_prog_seg(self.group_assignment,self.lambda_estimates,self.change_points_arr,self.num_roles,self.num_segments,self.df)
            # [self.lambda_estimates,self.change_points_arr]=sbm_core.linear_seg_ver_2(self.num_roles,self.num_segments,self.group_assignment,self.lambda_estimates,self.change_points_arr,self.df)
            # liklihood_sum = sbm_core.compute_cost(self.group_assignment,self.lambda_estimates,self.change_points_arr,self.num_roles,self.num_segments,self.df)            
            
            self.lambda_estimates=sbm_core.estimate_lamda_kh(self.num_roles,self.num_segments,self.lambda_estimates,self.group_assignment,self.change_points_arr,self.df,self.tuning_params)
            # print('after lambda estimate')
            liklihood_sum = sbm_core.compute_cost(self.group_assignment,self.lambda_estimates,self.change_points_arr,self.num_roles,self.num_segments,self.df)            
            # liklihood_sum = sbm_core.com_cost(self.num_roles,self.num_segments,self.lambda_estimates,self.change_points_arr,self.group_assignment,self.df)
            print(' %d %f'%(_itr+1,liklihood_sum))
            _prev_val = _curr_val
            _curr_val = liklihood_sum
            _itr+=1            
        return [self.group_assignment,self.lambda_estimates,self.change_points_arr,liklihood_sum]
    
    def k_h_seg_var_2(self):
        """(K,H)-segmentation variant-2"""
        
        # initilaization algorithm: initialize lambdas randomly and segments through linear seg. ver 2.            
        self.com_h_lvl_lambda()
                    
        [self.lambda_estimatess,self.change_points_arr]=sbm_core.linear_seg_ver_2(self.num_roles,self.num_segments,self.group_assignment,self.lambda_estimates,self.change_points_arr,self.df)
        self.lambda_estimates=sbm_core.estimate_lamda_kh(self.num_roles,self.num_segments,self.lambda_estimates,self.group_assignment,self.change_points_arr,self.df,self.tuning_params)        
        liklihood_sum = sbm_core.compute_cost(self.group_assignment,self.lambda_estimates,self.change_points_arr,self.num_roles,self.num_segments,self.df)
        _prev_val = math.inf 
        _curr_val = liklihood_sum
        _itr = 0        

        while round(_prev_val,2) != round(_curr_val,2):
            print("iteration no........... %d " %(_itr+1)) 
            
            self.group_assignment=sbm_core.group_assignment_ver2_2(self.nodes,self.num_roles,self.num_segments,self.lambda_estimates,self.group_assignment,self.change_points_arr,self.df)
            # print('after grouping')
            liklihood_sum = sbm_core.compute_cost(self.group_assignment,self.lambda_estimates,self.change_points_arr,self.num_roles,self.num_segments,self.df)
            
            self.lambda_estimates=sbm_core.estimate_lamda_kh(self.num_roles,self.num_segments,self.lambda_estimates,self.group_assignment,self.change_points_arr,self.df,self.tuning_params)
            # print('after lambda estimate')
            liklihood_sum = sbm_core.compute_cost(self.group_assignment,self.lambda_estimates,self.change_points_arr,self.num_roles,self.num_segments,self.df)

            # print('after seg')
            # self.change_points_arr = sbm_core.dyn_prog_seg(self.group_assignment,self.lambda_estimates,self.change_points_arr,self.num_roles,self.num_segments,self.df)
            [self.lambda_estimates,self.change_points_arr]=sbm_core.linear_seg_ver_2(self.num_roles,self.num_segments,self.group_assignment,self.lambda_estimates,self.change_points_arr,self.df)
            liklihood_sum = sbm_core.compute_cost(self.group_assignment,self.lambda_estimates,self.change_points_arr,self.num_roles,self.num_segments,self.df)            
            
            self.lambda_estimates=sbm_core.estimate_lamda_kh(self.num_roles,self.num_segments,self.lambda_estimates,self.group_assignment,self.change_points_arr,self.df,self.tuning_params)
            # print('after lambda estimate')
            # liklihood_sum = sbm_core.compute_cost(self.group_assignment,self.lambda_estimates,self.change_points_arr,self.num_roles,self.num_segments,self.df)            
            print(' %d %f'%(_itr+1,liklihood_sum))
            _prev_val = _curr_val
            _curr_val = liklihood_sum
            _itr+=1            
        return [self.group_assignment,self.lambda_estimates,self.change_points_arr,liklihood_sum,_itr]

    
    def com_h_lvl_lambda_group(self):
        H = self.num_levels
        print('k-h levels %d'%(self.num_levels))
                
        # h-level lambda estimates    
        # lambda_estimates_h = 1e-1*np.random.rand(self.num_roles, self.num_roles, H)
        lambda_estimates_h = np.random.rand(self.num_roles, self.num_roles, H)
        # print(lambda_estimates_h)
                
        l1 =list(range(0, H))
        l2 = []
        if self.num_segments > self.num_levels:
            l2 = [np.random.randint(0,H) for i in range(self.num_segments-H)]  
        
        # Mapping from segment to a level
        g_mapping= np.array(l1 + l2)
        print('g mapping {}'.format(g_mapping))
        
        R = int(self.num_roles)
        for k in range(0, R):
            for g in range(0, R):
                for d in range(0, self.num_levels):
                    self.lambda_estimates[k,g, d]= lambda_estimates_h[k,g,d]
                    self.lambda_estimates[g,k, d]= lambda_estimates_h[k,g,d]
        
        num_vertices = len(self.nodes)
        group_assignment_arr= np.random.randint(self.num_roles, size=(self.num_levels,num_vertices))  
            # node-group dictionary
        group_dic = {}
                            
        for i in range(0,self.num_levels):
            level = i
            
            group_dic_level = {}
            keys = self.nodes
            values = list(group_assignment_arr[level])
            group_dic_level = dict(zip(keys,values))
            group_dic[i] = group_dic_level
              
                
        self.group_assignment = group_dic   
        self.g_mapping = g_mapping
    
    def mm_k_h_seg_var_2(self):
        """(K,H)-segmentation variant-2"""
        
        # initilaization algorithm: initialize lambdas randomly and segments through linear seg. ver 2.            
        self.com_h_lvl_lambda_group()
                    
        [self.change_points_arr,self.g_mapping,self.group_assignment]=sbm_core.mm_linear_seg_ver_2(self.num_roles,self.num_segments,self.group_assignment,self.lambda_estimates,self.change_points_arr,self.df,self.g_mapping)
        self.lambda_estimates=sbm_core.mm_estimate_lamda_kh(self.num_roles,self.num_segments,self.lambda_estimates,self.group_assignment,self.change_points_arr,self.df,self.g_mapping,self.tuning_params)        
        liklihood_sum = sbm_core.mm_compute_cost(self.group_assignment,self.lambda_estimates,self.change_points_arr,self.num_roles,self.num_segments,self.df,self.g_mapping)
        _prev_val = math.inf 
        _curr_val = liklihood_sum
        _itr = 0        

        while round(_prev_val,2) != round(_curr_val,2):
            print("iteration no........... %d " %(_itr+1)) 
            
            self.group_assignment=sbm_core.mm_group_assignment_ver2_2(self.nodes,self.num_roles,self.num_segments,self.lambda_estimates,self.group_assignment,self.change_points_arr,self.df,self.g_mapping)
            # print('after grouping')
            # liklihood_sum = sbm_core.mm_compute_cost(self.group_assignment,self.lambda_estimates,self.change_points_arr,self.num_roles,self.num_segments,self.df,self.g_mapping)
            
            self.lambda_estimates=sbm_core.mm_estimate_lamda_kh(self.num_roles,self.num_segments,self.lambda_estimates,self.group_assignment,self.change_points_arr,self.df,self.g_mapping,self.tuning_params)
            # print('after lambda estimate')
            # liklihood_sum = sbm_core.mm_compute_cost(self.group_assignment,self.lambda_estimates,self.change_points_arr,self.num_roles,self.num_segments,self.df,self.g_mapping)

            print('after seg')
            # self.change_points_arr = sbm_core.dyn_prog_seg(self.group_assignment,self.lambda_estimates,self.change_points_arr,self.num_roles,self.num_segments,self.df)
            [self.change_points_arr,self.g_mapping,self.group_assignment]=sbm_core.mm_linear_seg_ver_2(self.num_roles,self.num_segments,self.group_assignment,self.lambda_estimates,self.change_points_arr,self.df,self.g_mapping)
            # liklihood_sum = sbm_core.mm_compute_cost(self.group_assignment,self.lambda_estimates,self.change_points_arr,self.num_roles,self.num_segments,self.df,self.g_mapping)            
            
            self.lambda_estimates=sbm_core.mm_estimate_lamda_kh(self.num_roles,self.num_segments,self.lambda_estimates,self.group_assignment,self.change_points_arr,self.df,self.g_mapping,self.tuning_params)
            # print('after lambda estimate')
            liklihood_sum = sbm_core.mm_compute_cost(self.group_assignment,self.lambda_estimates,self.change_points_arr,self.num_roles,self.num_segments,self.df,self.g_mapping)            
            print(' %d %f'%(_itr+1,liklihood_sum))
            _prev_val = _curr_val
            _curr_val = liklihood_sum
            _itr+=1            
        return [self.group_assignment,self.lambda_estimates,self.change_points_arr,liklihood_sum,self.g_mapping,_itr]