from qrem.types.datastructure_base import DataStructureBase
from qrem.common import math
from typing import Dict, Tuple, List
import operator
from functools import reduce
import numpy as np
import itertools


from qrem.common import math, probability




#TODO: TUZJAN, prepare full local test
#TODO: Find a way of exporting to jason dictionary, in which key is a tuple. This is needed to export CNModelData to json, e.g. noise matrices are stored in a dictionary with tupels as keys.  
class CNModelData(DataStructureBase):
  

    """
    The class contains parameters of Clusters and Neighbors CN noise model (currently it does not support neighbors). 
    
    It is used to:

    1) store results of CN noise model reconstruction algorithm (noise characterization routine)

    2) Mitigate redout errors on marginal probability distributions (error mitigation routine)  


    Parameters
    ----------
    clusters_tuple : a tuple of tuples 
        Stores information about clusters of qubits. Eeach tuple contains indices of qubits belonging to a cluster  e.g. ((0,2),(1,3))   

    noise_matrices: dictionary (key - tuple, value - array)
        Stores noise matrices associated to clusters. Key in the dictionary is a tuple with indices of qubits belonging to a cluster. 
        Corresponding value is an array storing noise matrix associated with the cluster. Eg. (0:2) : [[0.8,0,0,0], [0.2,1,0.1,0],[0,0,0.9,0],[0,0,0,1]].  

    inverse_noise_matrices : dictionary (key - tuple, value - array)
        Stores noise matrices associated to clusters. Key in the dictionary is a tuple with indices of qubits belonging to a cluster. 
        Corresponding value is an array storing inverse noise matrix associated with the cluster. Eg. (0:2) : [[0.8,0,0,0], [0.2,1,0.1,0],[0,0,0.9,0],[0,0,0,1]].
        This property is filled automatically when noise_matrices property is set via set_noise_matrices_dictionary method.

    qubit_in_cluster_membership: dictionary (key: tuple, value: tuple)
        Stores information about qubits membership in clusters. Key in the dictionary is a tuple constructed form a qubit index. Corresponding value is a cluster to which the qubit belongs.
        E.g. (2,) : (0,2,5). This property is filled automatically when noise_matrices property is set via set_noise_matrices_dictionary method

    
    Methods
    ----------
    set_noise_model:  
        Stores information about clusters of qubits. Eeach tuple contains indices of qubits belonging to a cluster  e.g. ((0,2),(1,3))   

    noise_matrices: dictionary (key - tuple, value - array)
        Stores noise matrices associated to clusters. Key in the dictionary is a tuple with indices of qubits belonging to a cluster. 
        Corresponding value is an array storing noise matrix associated with the cluster. Eg. (0:2) : [[0.8,0,0,0], [0.2,1,0.1,0],[0,0,0.9,0],[0,0,0,1]].  

    inverse_noise_matrices : dictionary (key - tuple, value - array)
        Stores noise matrices associated to clusters. Key in the dictionary is a tuple with indices of qubits belonging to a cluster. 
        Corresponding value is an array storing inverse noise matrix associated with the cluster. Eg. (0:2) : [[0.8,0,0,0], [0.2,1,0.1,0],[0,0,0.9,0],[0,0,0,1]].
        This property is filled automatically when noise_matrices property is set via set_noise_matrices_dictionary method.

    qubit_in_cluster_membership: dictionary (key: tuple, value: tuple)
        Stores information about qubits membership in clusters. Key in the dictionary is a tuple constructed form a qubit index. Corresponding value is a cluster to which the qubit belongs.
        E.g. (2,) : (0,2,5). This property is filled automatically when noise_matrices property is set via set_noise_matrices_dictionary method
    

'
"""
    def __init__(self, number_of_qubits: int):
        super().__init__()
        self.number_of_qubits = number_of_qubits
        self.locality = None  #integer
        self.clusters_tuple=()
        self.noise_matrices = {}
        self.inverse_noise_matrices ={}
        self.qubit_in_cluster_membership ={}
        self.composite_inverse_noise_matrices = {}
        self.clusters_neighborhoods = {}
    

    def _establish_qubit_in_cluster_dictionary(self):
        """
        Method used to establish membership of qubits in clusters. 
        It constructs a dictionary with key - tuple consisting of a qubit index e.g. (0,) and value - tuple storing corresponding cluster e.g. (0,3)
        It is automatically called when set_clusters_tuple method is used.     

        """
        for cluster in self.clusters_tuple:
            for qubit in cluster:
                self.qubit_in_cluster_membership[(qubit,)]= cluster
    
    def _compute_inverse_noise_matrices(self):
        """
        Method that creates inverse noise matrix dictionary for an instance of CN noise model. 
        In the dictionary key corresponds to a tuple storing a cluster, and value to an numpy array storing inverse noise matrix e.g. (0,5): 4 by 4 numpy array.
        It is automatically called when set_noise_matrices_dictionary is used.     

        """

        for cluster, noise_matrix in self.noise_matrices.items():
            
            if bool(self.clusters_neighborhoods):

                try:
                    self.inverse_noise_matrices[(cluster)] = np.linalg.inv(noise_matrix['averaged'])
                except:
                    print(f"MATRIX FOR SUBSET: {(cluster)} is not invertible, computing pseudinverse",'','red')
                    self.inverse_noise_matrices[(cluster)] = np.linalg.pinv(noise_matrix)['averaged']


            else:
                try:
                    self.inverse_noise_matrices[(cluster)] = np.linalg.inv(noise_matrix)
                except:
                    print(f"MATRIX FOR SUBSET: {(cluster)} is not invertible, computing pseudinverse",'','red')
                    self.inverse_noise_matrices[(cluster)] = np.linalg.pinv(noise_matrix)
    

    def set_clusters_tuple(self,clusters_tuple: tuple):
        """
        Method fills clusters_tuple property. It is called when properties of a CN noise model are established via set_noise_model method.   
        In addition the method calls _establish_qubit_in_cluster_dictionary, which fills qubit_in_cluster_membership property 

        """
        if isinstance(clusters_tuple,tuple) and isinstance(clusters_tuple[0],tuple):
            self.clusters_tuple = clusters_tuple
            self._establish_qubit_in_cluster_dictionary()
        else:
             raise Exception("Wrong data format. Cluster list should be a tuple of tuples")
    
    
        
    def set_noise_matrices_dictionary(self, noise_matrices_dictionary: Dict):

        """
        Method fills noise_matrices property. It is called when properties of a CN noise model are established via set_noise_model method.   
        In addition the method calls _compute_inverse_noise_matrices, which fills inverse_noise_matrices property.

        Parameters
        ----------
        noise_matrices_dictionary : dictionary (key - tuple of integers, value - array)
            Dictionary storing noise matrices for a CN noise model. Key in the dictionary is a tuple with indices of qubits belonging to a cluster. 
            Corresponding value is an array storing noise matrix associated with the cluster. Eg. (0:2) : [[0.8,0,0,0], [0.2,1,0.1,0],[0,0,0.9,0],[0,0,0,1]].
        
        

        """


        if isinstance(noise_matrices_dictionary,dict):
            self.noise_matrices =  noise_matrices_dictionary
            self._compute_inverse_noise_matrices()
        else:
             raise Exception("Wrong data format. Cluster list should be a list of tuples")
        
    def set_noise_model(self,noise_matrices_dictionary):

        """
        Main method used to fill properties of CN noise model. The following properties are filled when the method is called:    
        1) clusters_tuple
        2) qubit_in_cluster_membership
        3) noise_matrices 
        4) inverse_noise_matrices 
          

        Parameters
        ----------
        noise_matrices_dictionary : dictionary (key - tuple of integers, value - array)
            Dictionary storing noise matrices for a CN noise model. Key in the dictionary is a tuple with indices of qubits belonging to a cluster. 
            Corresponding value is an array storing noise matrix associated with the cluster. Eg. (0:2) : [[0.8,0,0,0], [0.2,1,0.1,0],[0,0,0.9,0],[0,0,0,1]].

        """

        self.set_clusters_tuple(clusters_tuple=tuple(noise_matrices_dictionary.keys()))
        self.set_noise_matrices_dictionary(noise_matrices_dictionary=noise_matrices_dictionary)

    def compute_extended_inverse_noise_matrices(self, clusters_in_marginal_list:List[Tuple])-> np.array:

        """
        Function determines inverse noise matrix for a given marginal. It is used as a step in mitigation_marginal function.
        The matrix is constructed as a tensor product on inverse noise matrices involved in a marginal, when it is necessary
        the matrix is permuted to ensure proper ordering of qubits indices. E.g. for a cluster_in_marginal_list 
        [(0,2),(1,3)], the returned matrix corresponds to qubits ordered as (0,1,2,3).   


        Parameters
        ----------
        noise_model : object of CNModelData class
            An object of CNModelData class 

        clusters_in_marginal_list:
            A list of tuples with clusters involved in the marginal   
            

        Returns
        -------

            An inverse noise matrix for qubits specified in clusters_in_marginal_list, qubits are sorted in ascending order 
            (e.g. for a clusters_in_marginal_list =[(0,4),(1,8)], indices of the inverse noise matrix indices correspond to qubits in the order (0,1,4,8) )
        
        """ 

        # a list of qubits is established
        unordered_qubits_in_marginal_list = list(reduce(operator.concat, clusters_in_marginal_list))

        #the inverse noise matrix is initialized as a unit numpy array 
        marginal_inverse_noise_matrix = np.array([1])
        
        # a loop over clusters in clusters_in_marginal_list
        for cluster in clusters_in_marginal_list:
        
            # total noise matrix is updated by taking tensor product with a inverse noise matrix of the currect cluster
            marginal_inverse_noise_matrix = np.kron(marginal_inverse_noise_matrix, self.inverse_noise_matrices[cluster])
    
        #final noise matrix is returned, permuted to preserve ascending order of qubit indices if necessary 
        marginal_inverse_noise_matrix = math.permute_composite_matrix(qubits_list=unordered_qubits_in_marginal_list,noise_matrix=marginal_inverse_noise_matrix)

        unordered_qubits_in_marginal_list.sort()

        key = tuple(unordered_qubits_in_marginal_list)

        self.inverse_noise_matrices[key] = marginal_inverse_noise_matrix 

        return marginal_inverse_noise_matrix 




#MOVE TO CNModelData class   

    def get_clusters_in_marginal_list(self, marginal : Tuple[int]) -> List[Tuple]:
        
        """
        Function creates a list of clusters that are involved in a marginal. Used in mitigation routines.
        For a given marginal inspects provided noise model and checks clusters membership of qubits form
        marginal.
        
        Parameters
        ----------
        marginal : tuple
            A tuple specifying marginal 

        noise_model : object of CNModelData class
            An object of CNModelData class  
            

        Returns
        -------

        clusters_in_marginal_list
            A list of tuples involved in the input marginal 
            
        
        """  

        #list storing clusters involved in the marginal, empty for start
        clusters_in_marginal_list=[]
        
        #a loop over qubits in the marginal
        for qubit in marginal:
            
            #a cluster to which the qubit belongs is determined 
            cluster = self.qubit_in_cluster_membership[(qubit,)]
            
            #this cluster is appended to clusters_in_marginal_list if it is not there 
            if cluster not in clusters_in_marginal_list:
                clusters_in_marginal_list.append(self.qubit_in_cluster_membership[(qubit,)])
        
        #results are returned
        return clusters_in_marginal_list
    

    def get_clusters_and_neighborhoods_in_marginal_dictionary(self, marginal : Tuple[int]) -> Dict[Tuple[int], Tuple[int]]:
        
        """
        Function creates a list of clusters that are involved in a marginal. Used in mitigation routines.
        For a given marginal inspects provided noise model and checks clusters membership of qubits form
        marginal.
        
        Parameters
        ----------
        marginal : tuple
            A tuple specifying marginal 

        noise_model : object of CNModelData class
            An object of CNModelData class  
            

        Returns
        -------

        clusters_in_marginal_list
            A list of tuples involved in the input marginal 
            
        
        """  

        #list storing clusters involved in the marginal, empty for start
        clusters_and_neighbors_in_marginal_dictionary={}
        
        if bool(self.clusters_neighborhoods):
        #a loop over qubits in the marginal
            for qubit in marginal:
                
                #a cluster to which the qubit belongs is determined 
                cluster = self.qubit_in_cluster_membership[(qubit,)]

                
                
                #this cluster is appended to clusters_in_marginal_list if it is not there 
                if cluster not in clusters_and_neighbors_in_marginal_dictionary.keys():
                    clusters_and_neighbors_in_marginal_dictionary[cluster] = self.clusters_neighborhoods[cluster]
            
            #results are returned
        return clusters_and_neighbors_in_marginal_dictionary
    


    def get_neighbors_in_marginal(self, marginal : Tuple[int]) -> Dict[Tuple[int], Tuple[int]]:
        
        """
        Function creates a dictionary with clusters and neighbors that belong to a marginal (it can happen that clusters' neighbors do not belong to other clusters that form a marginal .
        For a given marginal, and clusters and neighbors, membership of neighbors in the marginal is checked  
        
        Parameters
        ----------
        marginal : tuple
            A tuple specifying marginal 

        noise_model : object of CNModelData class
            An object of CNModelData class  
            

        Returns
        -------

        clusters_in_marginal_list
            A list of tuples involved in the input marginal 
            
        
        """  
        
        #here a dictionary is created to store neighbors that take part in a marginal 

        #dictionary with keys corresponding to clusters and values corresponding to neighbors that are in a marginal, and their position relative to a neighborhood 
        neighbors_coinciding_with_clusters_dictionary = {}

        clusters_and_neighbors_in_marginal_dictionary = self.get_clusters_and_neighborhoods_in_marginal_dictionary(marginal=marginal)

        qubits_in_marginal_list = list(reduce(operator.concat, list(clusters_and_neighbors_in_marginal_dictionary.keys())))


        #a loop over clusters and neighbors 
        for cluster, neighbors in clusters_and_neighbors_in_marginal_dictionary.items():

            # a temporary list storing neighbors
            neighbor_in_marginal = []

            #if there are neighbors for a cluster the check is performed 
            if neighbors!= None:

                #qubits in neighborhood are enumerated to store their relative position
                for index, qubit in enumerate(neighbors):

                    #if a qubit form this neighborhood is present in the marginal, its index and relative position is stored in a list 
                    if qubit in qubits_in_marginal_list:

                        neighbor_in_marginal.append((index,qubit))
                    
                #if after a previous loop neighbors were identified a dictionary entry is cerated
                if len(neighbor_in_marginal) > 0:

                    neighbors_coinciding_with_clusters_dictionary[cluster] = tuple(neighbor_in_marginal)
            
        return neighbors_coinciding_with_clusters_dictionary 
    

    def get_noise_matrices_indexes(self, marginal : Tuple[int] ,state):
        
        """
        Function creates a dictionary with indices of state dependent noise matrices that should be used to compute noise matrix inverse, 

        It is used during error  mitigation of CN noise model with neighbors.  
        
        Parameters
        ----------
        marginal : tuple
            A tuple specifying marginal 

        noise_model : object of CNModelData class
            An object of CNModelData class  
            

        Returns
        -------

        neighbors_indices_to_compute_average 
            A dictionary with a key - cluster and value - tuple of indicies to be averaged over or a string 'averaged' 
            
              
            
        
        """  
        
        
        
        #a dictionary to store information about noise matrices needed to compute an inverse for a given cluster is created 
        neighbors_indices_to_compute_average ={} 

        # this loop goes over all clusters that take part in a marginal, and their neighbors 

        neighbors_coinciding_with_clusters_dictionary = self.get_neighbors_in_marginal(marginal=marginal)

        clusters_and_neighbors_in_marginal_dictionary = self.get_clusters_and_neighborhoods_in_marginal_dictionary(marginal=marginal) 

        
        for cluster, neighbors in clusters_and_neighbors_in_marginal_dictionary.items():
            
            #this conditions verify whether given cluster has a neighbor that actual coincides with the marginal    
            if cluster in neighbors_coinciding_with_clusters_dictionary.keys():

                #here all possible states of the clusters neighbors are created 
                possible_states_list = list(itertools.product([0,1], repeat = len(neighbors)))
                
                #here we check, which states of neighbors are needed to compute inverse noise matrix 
                for neighbor_tuple in neighbors_coinciding_with_clusters_dictionary[cluster]:

                #possible_state is ordered according to position of a qubit in its neighborhood (relative index), while the position of a qubit in state is just its index
                    possible_states_list = [ possible_state for possible_state in possible_states_list if possible_state[neighbor_tuple[0]] == int(state[neighbor_tuple[1]]) ]
                
                neighbors_indices_to_compute_average[cluster] = tuple(possible_states_list)
            
            #if there are no neighbors in the marginal averaged noise matrix is used 
            else:

                neighbors_indices_to_compute_average[cluster] = 'averaged'
        
        return neighbors_indices_to_compute_average 


    def compute_extended_inverse_noise_matrices_state_dependent(self, clusters_in_marginal_list:List[Tuple],marginal,state)-> np.array:

        """
        Function determines inverse noise matrix for a given marginal. It is used as a step in mitigation_marginal function.
        The matrix is constructed as a tensor product on inverse noise matrices involved in a marginal, when it is necessary
        the matrix is permuted to ensure proper ordering of qubits indices. E.g. for a cluster_in_marginal_list 
        [(0,2),(1,3)], the returned matrix corresponds to qubits ordered as (0,1,2,3).   


        Parameters
        ----------
        noise_model : object of CNModelData class
            An object of CNModelData class 

        clusters_in_marginal_list:
            A list of tuples with clusters involved in the marginal   
            

        Returns
        -------

            An inverse noise matrix for qubits specified in clusters_in_marginal_list, qubits are sorted in ascending order 
            (e.g. for a clusters_in_marginal_list =[(0,4),(1,8)], indices of the inverse noise matrix indices correspond to qubits in the order (0,1,4,8) )
        
        """ 

        # a list of qubits is established
        unordered_qubits_in_marginal_list = list(reduce(operator.concat, clusters_in_marginal_list))

        #indicies of needed noise matrix are established
        noise_matrices_indicies = self.get_noise_matrices_indexes(marginal=marginal,state=state)

        #the inverse noise matrix is initialized as a unit numpy array 
        marginal_inverse_noise_matrix = np.array([1])
        
        # a loop over clusters in clusters_in_marginal_list
        for cluster in clusters_in_marginal_list:
            
            if noise_matrices_indicies[cluster]  == 'averaged':      
                
                # total noise matrix is updated by taking tensor product with a inverse noise matrix of the currect cluster
                marginal_inverse_noise_matrix = np.kron(marginal_inverse_noise_matrix, self.inverse_noise_matrices[cluster])

            else:

                current_noise_matrix = self.compute_specific_state_dependent_noise_matrix(cluster=cluster, indicies=noise_matrices_indicies[cluster])

                current_inverse_noise_matrix = np.linalg.inv(current_noise_matrix)

                marginal_inverse_noise_matrix = np.kron(marginal_inverse_noise_matrix, current_inverse_noise_matrix)




    
        #final noise matrix is returned, permuted to preserve ascending order of qubit indices if necessary 
        marginal_inverse_noise_matrix = math.permute_composite_matrix(qubits_list=unordered_qubits_in_marginal_list,noise_matrix=marginal_inverse_noise_matrix)

        unordered_qubits_in_marginal_list.sort()

        key = tuple(unordered_qubits_in_marginal_list)

        #self.inverse_noise_matrices[key] = marginal_inverse_noise_matrix 

        return marginal_inverse_noise_matrix 
    

    def compute_specific_state_dependent_noise_matrix(self, cluster: Tuple[int], indicies: Tuple[Tuple[int]]) -> np.array:

        noise_matrix = self.noise_matrices[cluster]
        
        normalization_factor = len(indicies)
        
        dim = len(noise_matrix['averaged'] )

        averaged_noise_matrix = np.zeros((dim,dim))

        for index in indicies:

            averaged_noise_matrix += noise_matrix[index]

        averaged_noise_matrix = 1/normalization_factor * averaged_noise_matrix

        return averaged_noise_matrix

        



    
   
# from qrem.functions_qrem import ancillary_functions as anf

# def t1(test_model:type[CNModelData]):
#     # Test of json export

    
#     json_dict_test = test_model.to_json()

#     new_model =CNModelData(number_of_qubits=5)

#     new_model.import_json(json_dict_test)

# def t2(test_model:type[CNModelData]):
#     # Test of pickle export

 
#     pickle_data = test_model.export_pickle("test_data")

#     new_model =CNModelData(number_of_qubits=5)

#     new_model.import_pickle("test_data")

#     print(new_model.clusters_tuple)

    


if __name__=='__main__':
   #set clusters tuple
    qubits_in_cluster_1=(0,3)
    qubits_in_cluster_2=(1,2,4)

    #generation noise matrice for clusters with built in function probability.random_stochastic_matrix
    noise_matrix_1= probability.random_stochastic_matrix(2**2)
    noise_matrix_2= probability.random_stochastic_matrix(2**3)
    
    noise_matrices_dictionary_test = {qubits_in_cluster_1: noise_matrix_1, qubits_in_cluster_2: noise_matrix_2}

    #create noise model 
    test_model= CNModelData(number_of_qubits=5)
    test_model.set_noise_model(noise_matrices_dictionary=noise_matrices_dictionary_test) 

    t2(test_model)
 

        


    #}



   # test_model.load_from_dict(dictionary_to_load)
   # print('after loading:\n', test_model.get_dict_format())

    #NOTE: Is it possible to have a nested dictionary with keys that are not string? 
    #json_dict_test = test_model.to_json()
 
    #pickle_dict_test = test_model.get_pickle()
    #test_model.export_pickle('exported_pickle', overwrite=True)
    #print('after change:\n', test_model.get_dict_format())
    #test_model.import_pickle('exported_pickle')