import numpy as np

class LOTBarycenter:
    @staticmethod
    def generate_barycenters_within_class(pclouds, labels, weights=None, output_labels=None, uniform=True, n=1):
        """
        Generate barycenters of point clouds within the same class.

        Parameters
        ----------
        pclouds : np.array
            A 2D array where each row represents a point cloud.
            Shape: (number_of_point_clouds, number_of_features)
        
        labels : np.array
            A 1D array of labels corresponding to the point clouds in `pclouds`.
            Shape: (number_of_point_clouds,)
        
        weights : list of np.array, optional
            A list of weight vectors for generating barycenters. If `None`, random weights will be generated.
            Default is None.

        output_labels : list or np.array, optional
            A list of labels corresponding to the output barycenters. Must be given if weights is given, and must be same length as weights. Specifies which class each weight vector corresponds to.
        
        uniform : bool, optional
            Whether the returned barycenter uses the uniform weight vector. Default is true, overridden if weights is given
        
        n : int, optional
            The number of random weight vectors to generate for each unique label if `weights` is not provided.
            Default is 1.

        Returns
        -------
        barycenters : np.array
            A 2D array where each row is a barycenter generated by combining point clouds in `pclouds` according
            to the corresponding weight vector. Shape: (number_of_barycenters, number_of_features)
        
        output_labels : np.array
            A 1D array of labels corresponding to the generated barycenters. Shape: (number_of_barycenters,)
        
        weights : np.array
            A 2D array of the weight vectors used to generate the barycenters.
        """
        # Identify unique labels in the dataset
        unique_labels = np.unique(labels)
        
        # If uniform is True, only generate 1 barycenter per class
        if uniform: 
            n = 1
        
        # If no weights are provided, generate random or uniform weights
        if weights is None:
            weights = []
            output_labels = []
            
            # Iterate through each unique label
            for label in unique_labels:
                # Select the point clouds corresponding to the current label
                pclouds_class = pclouds[labels == label]
                
                # create output_labels list
                output_labels += [label] * n

                # Generate 'n' weight vectors
                for _ in range(n):
                    # Generate random weights
                    lambd = np.random.exponential(scale=1, size=pclouds_class.shape[0])  
                    if uniform: lambd = np.ones(pclouds_class.shape[0])  # If uniform, set all weights to 1
                    lambd /= lambd.sum()  # Normalize the weights to sum to 1
                    weights.append(lambd)
                
        # else:
        #     n = len(weights)
        #     weights = np.vstack([weights]*len(unique_labels))
            

        # List to store the generated barycenters
        barycenters = []
        
        for i in range(len(weights)):
            label = output_labels[i]
            pclouds_class = pclouds[labels == label]
            barycenters.append(np.dot(weights[i],pclouds_class))
        
        # # Iterate through each unique label to generate barycenters
        # for i, label in enumerate(unique_labels):
        #     # Select the point clouds corresponding to the current label
        #     pclouds_class = pclouds[labels == label]
            
        #     # Generate barycenters using the weights
        #     for weight in weights[i * n:(i + 1) * n]:
        #         barycenters.append(np.dot(weight, pclouds_class))

        # Return the generated barycenters, corresponding labels, and weights used
        return np.array(barycenters), np.array(output_labels), weights

    @staticmethod
    def generate_barycenters_between_classes(pclouds, labels, class_pairs, weights=None):
        """
        Generate barycenters by combining point clouds from different classes using randomly selected representatives.

        Parameters
        ----------
        pclouds : np.array
            A 2D array where each row represents a point cloud.
            Shape: (number_of_point_clouds, number_of_features)
        
        labels : np.array
            A 1D array of labels corresponding to the point clouds in `pclouds`.
            Shape: (number_of_point_clouds,)
        
        class_pairs : list of tuple
            A list of tuples where each tuple contains two class labels to combine for generating barycenters.
        
        weights : np.array, optional
            A list of weight vectors for generating barycenters by combining the point clouds in the specified classes.
            If `None`, random weights will be generated. Each weight vector should sum to 1.

        Returns
        -------
        barycenters : np.array
            A 2D array where each row is a barycenter generated by combining point clouds from the two classes.
            Shape: (number_of_barycenters, number_of_features)
        
        representatives : list of tuples
            A list of tuples where each tuple contains the two randomly selected representatives from the classes.
        
        weights : list of np.array
            A list of the weight vectors used to generate the barycenters.
        """
        # Initialize lists to store the generated barycenters and representatives
        barycenters = []
        representatives = []
        
        # Initialize the weights list if it's not provided
        if weights is None: weights_list = []
        else: weights_list = weights

        # Iterate through each pair of classes specified in class_pairs
        for (label1, label2) in class_pairs:
            # Select point clouds corresponding to the first class in the pair
            pclouds_class1 = pclouds[labels == label1]
            # Select point clouds corresponding to the second class in the pair
            pclouds_class2 = pclouds[labels == label2]

            # Randomly select one representative from each class
            rep1 = pclouds_class1[np.random.choice(len(pclouds_class1))]
            rep2 = pclouds_class2[np.random.choice(len(pclouds_class2))]

            # Store the selected representatives
            representatives.append((rep1, rep2))

            # If weights are not provided, generate random weights
            if weights is None:
                lambd = np.random.rand(2)
                lambd /= lambd.sum()  # Normalize the weights to sum to 1
                weights_list.append(lambd)
            else:
                lambd = weights_list.pop(0)
            
            # Generate the barycenter using the selected representatives and weights
            barycenter = np.dot(lambd, np.array([rep1, rep2]))
            barycenters.append(barycenter)
        
        # Return the generated barycenters, the representatives used, and the weights used
        return np.array(barycenters), representatives, weights_list


    @staticmethod
    def generate_barycenters_general(pclouds, weights):
        """
        Generate barycenters from arbitrary point clouds using specified weights.

        Parameters
        ----------
        pclouds : np.array
            A 2D array where each row represents a point cloud.
            Shape: (number_of_point_clouds, number_of_features)
        
        weights : np.array
            A 2D array where each row is a weight vector for generating a barycenter by combining the point clouds.

        Returns
        -------
        barycenters : np.array
            A 2D array where each row is a barycenter generated by combining point clouds in `pclouds` according
            to the corresponding weight vector. Shape: (number_of_barycenters, number_of_features)
        """
        # Generate barycenters by computing the weighted sum of point clouds based on the provided weights
        barycenters = np.dot(weights, pclouds)
        
        # Return the generated barycenters
        return barycenters
