import torch
import numpy as np
from math import pi
import warnings

class GaussianMixture(torch.nn.Module):
    """
    Fits a mixture of k=1,..,K Gaussians to the input data (K is supplied via n_components).
    Input tensors are expected to be flat with dimensions (n: number of samples, d: number of features).
    The model then extends them to (n, 1, d).
    The model parametrization (mu, sigma) is stored as (1, k, d),
    probabilities are shaped (n, k, 1) if they relate to an individual sample,
    or (1, k, 1) if they assign membership probabilities to one of the mixture components.
    """
    def __init__(self, n_components, n_features, device, covariance_type="full", eps=1e-6, init_params="kmeans++", mu_init=None, var_init=None):
        """
        Initializes the model and brings all tensors into their required shape.
        The class expects data to be fed as a flat tensor in (n, d).
        The class owns:
            x:               torch.Tensor (n, 1, d)
            mu:              torch.Tensor (1, k, d)
            var:             torch.Tensor (1, k, d) or (1, k, d, d)
            pi:              torch.Tensor (1, k, 1)
            covariance_type: str
            eps:             float
            init_params:     str
            log_likelihood:  float
            n_components:    int
            n_features:      int
        args:
            n_components:    int
            n_features:      int
        options:
            mu_init:         torch.Tensor (1, k, d)
            var_init:        torch.Tensor (1, k, d) or (1, k, d, d)
            covariance_type: str
            eps:             float
            init_params:     str
        """
        super(GaussianMixture, self).__init__()

        self.n_components = n_components
        self.n_features = n_features

        self.mu_init = mu_init
        self.var_init = var_init
        self.eps = eps

        self.log_likelihood = -np.inf

        self.covariance_type = covariance_type
        self.init_params = init_params
        
        self.device = device
        self.reg_covar = 1e-06

        assert self.covariance_type in ["full", "diag"]
        assert self.init_params in ["kmeans", "random", "kmeans++"]


    def _init_params(self, x):
        """Initialize parameters for the Gaussian Mixture Model.
        If mu_init or var_init are provided, they will be used instead of random initialization.
        """
        device = self.device
        
        # Initialize means (mu)
        if self.mu_init is not None:
            expected_shape = (1, self.n_components, self.n_features)
            if self.mu_init.size() != expected_shape:
                raise ValueError(
                    f"Input mu_init does not have required tensor dimensions {expected_shape}, "
                    f"got {tuple(self.mu_init.size())}"
                )
            self.mu = torch.nn.Parameter(self.mu_init.to(device), requires_grad=False)
        else:
            if self.init_params == "kmeans" and self.mu_init is None:
                self.mu = torch.nn.Parameter(self.get_kmeans_mu(x, n_centers=self.n_components), requires_grad=False)
            elif self.init_params == "kmeans++" and self.mu_init is None:
                self.mu = torch.nn.Parameter(self.get_kmeans_mu_plus_plus(x, n_centers=self.n_components), requires_grad=False)

        # Initialize covariances (var)
        if self.covariance_type == "diag":
            if self.var_init is not None:
                expected_shape = (1, self.n_components, self.n_features)
                if self.var_init.size() != expected_shape:
                    raise ValueError(
                        f"Input var_init does not have required tensor dimensions {expected_shape}, "
                        f"got {tuple(self.var_init.size())}"
                    )
                self.var = torch.nn.Parameter(self.var_init.to(device), requires_grad=False)
            else:
                # Initialize with ones but add small noise for symmetry breaking
                self.var = torch.nn.Parameter(
                    torch.ones(1, self.n_components, self.n_features, device=device) + \
                    torch.randn(1, self.n_components, self.n_features, device=device) * 0.01,
                    requires_grad=False
                )
        
        elif self.covariance_type == "full":
            if self.var_init is not None:
                expected_shape = (1, self.n_components, self.n_features, self.n_features)
                if self.var_init.size() != expected_shape:
                    raise ValueError(
                        f"Input var_init does not have required tensor dimensions {expected_shape}, "
                        f"got {tuple(self.var_init.size())}"
                    )
                self.var = torch.nn.Parameter(self.var_init.to(device), requires_grad=False)
            else:
                # Initialize as identity matrix with small random noise
                identity = torch.eye(self.n_features, device=device)
                noise = torch.randn(self.n_features, self.n_features, device=device) * 0.01
                self.var = torch.nn.Parameter(
                    (identity + noise).reshape(1, 1, self.n_features, self.n_features)
                    .repeat(1, self.n_components, 1, 1),
                    requires_grad=False
                )

        # Initialize mixing coefficients (pi)
        self.pi = torch.nn.Parameter(
            torch.full((1, self.n_components, 1), 1. / self.n_components, device=device),
            requires_grad=False
        )

        # Optional: add small epsilon to prevent numerical instability
        if hasattr(self, 'reg_covar') and self.reg_covar > 0:
            if self.covariance_type == "diag":
                self.var.data += self.reg_covar
            else:  # full
                self.var.data += torch.eye(self.n_features, device=device) * self.reg_covar

        self.params_fitted = False


    def check_size(self, x):
        if len(x.size()) == 2:
            # (n, d) --> (n, 1, d)
            x = x.unsqueeze(1)

        return x.to(self.device)


    def bic(self, x):
        """
        Bayesian information criterion for a batch of samples.
        args:
            x:      torch.Tensor (n, d) or (n, 1, d)
        returns:
            bic:    float
        """
        x = self.check_size(x)
        n = x.shape[0]

        # Free parameters for covariance, means and mixture components
        free_params = self.n_features * self.n_components + self.n_features + self.n_components - 1

        bic = -2. * self.__score(x, as_average=False).mean() * n + free_params * np.log(n)

        return bic


    def fit(self, x, delta=1e-3, n_iter=100, n_init=10, warm_start=False):
        """
        Fits model to the data multiple times and keeps the best parameters.
        
        args:
            x:          torch.Tensor (n, d) or (n, k, d)
        options:
            delta:      float, tolerance for convergence
            n_iter:     int, maximum number of iterations per initialization
            n_init:     int, number of times the algorithm will be run with different initializations
            warm_start: bool, if True, use previous parameters as initialization
        """
        x = self.check_size(x)
        device = self.device
        
        # if warm_start is enabled, we will have a unique initialization
        do_init = not (warm_start and self.params_fitted)
        n_init = n_init if do_init else 1
        
        max_log_likelihood = -np.inf
        best_params = None
        
        for init in range(n_init):
            if not do_init and self.params_fitted:
                # If warm_start is True and we have fitted parameters, skip initialization
                current_log_likelihood = self.log_likelihood
            else:
                # Initialize parameters
                self._init_params(x)
                
                current_log_likelihood = -np.inf
                
            # Store initial parameters
            best_mu = self.mu.data.clone()
            best_var = self.var.data.clone()
            best_pi = self.pi.data.clone() if hasattr(self, 'pi') else None
            
            i = 0
            j = np.inf
            converged = False
            
            while (i <= n_iter) and (j >= delta):
                log_likelihood_old = current_log_likelihood
                mu_old = self.mu.data.clone()
                var_old = self.var.data.clone()
                
                # EM step
                self.__em(x)
                current_log_likelihood = self.__score(x)
                
                # Check for numerical stability
                if torch.isinf(current_log_likelihood.abs()) or torch.isnan(current_log_likelihood):
                    # When the log-likelihood assumes unbound values, reinitialize parameters
                    self._init_params(x)
                    for p in self.parameters():
                        p.data = p.data.to(device)
                    if self.init_params == "kmeans":
                        self.mu.data = self.get_kmeans_mu(x, n_centers=self.n_components)
                    elif self.init_params == "kmeans++":
                        self.mu.data = self.get_kmeans_mu_plus_plus(x, n_centers=self.n_components)
                    current_log_likelihood = -np.inf
                
                i += 1
                j = current_log_likelihood - log_likelihood_old
                
                # If likelihood improved, update best parameters
                if current_log_likelihood > max_log_likelihood:
                    max_log_likelihood = current_log_likelihood
                    best_mu = self.mu.data.clone()
                    best_var = self.var.data.clone()
                    if hasattr(self, 'pi'):
                        best_pi = self.pi.data.clone()
                    
                if j <= delta:
                    # When score decreases or convergence reached, revert to best parameters
                    converged = True
                    break
            
            # If this initialization gave us the best log likelihood so far, store the parameters
            if current_log_likelihood > max_log_likelihood:
                max_log_likelihood = current_log_likelihood
                best_params = {
                    'mu': best_mu,
                    'var': best_var,
                    'pi': best_pi if hasattr(self, 'pi') else None
                }
        
        # Set the best parameters found
        if best_params is not None:
            self.mu.data = best_params['mu']
            self.var.data = best_params['var']
            if hasattr(self, 'pi') and best_params['pi'] is not None:
                self.pi.data = best_params['pi']
            self.log_likelihood = max_log_likelihood
        
        self.params_fitted = True
        
        # Warn if none of the initializations converged
        if not converged and n_iter > 0:
            warnings.warn(
                "Best performing initialization did not converge. "
                "Try different init parameters, increase n_iter, "
                "or decrease delta.",
                RuntimeWarning
            )
        
        return self


    def predict(self, x, probs=False):
        """
        Assigns input data to one of the mixture components by evaluating the likelihood under each.
        If probs=True returns normalized probabilities of class membership.
        args:
            x:          torch.Tensor (n, d) or (n, 1, d)
            probs:      bool
        returns:
            p_k:        torch.Tensor (n, k)
            (or)
            y:          torch.LongTensor (n)
        """
        x = self.check_size(x)

        weighted_log_prob = self._estimate_log_prob(x) + torch.log(self.pi)

        if probs:
            p_k = torch.exp(weighted_log_prob)
            return torch.squeeze(p_k / (p_k.sum(1, keepdim=True)))
        else:
            return torch.squeeze(torch.max(weighted_log_prob, 1)[1].type(torch.LongTensor))


    def predict_proba(self, x):
        """
        Returns normalized probabilities of class membership.
        args:
            x:          torch.Tensor (n, d) or (n, 1, d)
        returns:
            y:          torch.LongTensor (n)
        """
        return self.predict(x, probs=True)


    def sample(self, n):
        """
        Samples from the model.
        args:
            n:          int
        returns:
            x:          torch.Tensor (n, d)
            y:          torch.Tensor (n)
        """
        counts = torch.distributions.multinomial.Multinomial(total_count=n, probs=self.pi.squeeze()).sample()
        x = torch.empty(0, device=counts.device)
        y = torch.cat([torch.full([int(sample)], j, device=counts.device) for j, sample in enumerate(counts)])

        # Only iterate over components with non-zero counts
        for k in np.arange(self.n_components)[counts > 0]: 
            if self.covariance_type == "diag":
                x_k = self.mu[0, k] + torch.randn(int(counts[k]), self.n_features, device=self.device) * torch.sqrt(self.var[0, k])
            elif self.covariance_type == "full":
                d_k = torch.distributions.multivariate_normal.MultivariateNormal(self.mu[0, k], self.var[0, k])
                x_k = torch.stack([d_k.sample() for _ in range(int(counts[k]))])

            x = torch.cat((x, x_k), dim=0)

        return x, y


    def score_samples(self, x):
        """
        Computes log-likelihood of samples under the current model.
        args:
            x:          torch.Tensor (n, d) or (n, 1, d)
        returns:
            score:      torch.LongTensor (n)
        """
        x = self.check_size(x)

        score = self.__score(x, as_average=False)
        return score


    def _estimate_log_prob(self, x):
        """
        Returns a tensor with dimensions (n, k, 1), which indicates the log-likelihood that samples belong to the k-th Gaussian.
        args:
            x:            torch.Tensor (n, d) or (n, 1, d)
        returns:
            log_prob:     torch.Tensor (n, k, 1)
        """
        x = self.check_size(x)

        if self.covariance_type == "full":
            mu = self.mu
            var = self.var

            precision = torch.inverse(var)
            d = x.shape[-1]

            log_2pi = d * np.log(2. * pi)

            log_det = self._calculate_log_det(precision)

            x_mu_T = (x - mu).unsqueeze(-2)
            x_mu = (x - mu).unsqueeze(-1)

            x_mu_T_precision = calculate_matmul_n_times(self.n_components, x_mu_T, precision)
            x_mu_T_precision_x_mu = calculate_matmul(x_mu_T_precision, x_mu)

            return -.5 * (log_2pi - log_det + x_mu_T_precision_x_mu)

        elif self.covariance_type == "diag":
            mu = self.mu
            prec = torch.rsqrt(self.var)

            log_p = torch.sum((mu * mu + x * x - 2 * x * mu) * prec, dim=2, keepdim=True)
            log_det = torch.sum(torch.log(prec), dim=2, keepdim=True)

            return -.5 * (self.n_features * np.log(2. * pi) + log_p - log_det)


    def _calculate_log_det(self, var):
        """
        Calculate log determinant in log space, to prevent overflow errors.
        args:
            var:            torch.Tensor (1, k, d, d)
        """
        log_det = torch.empty(size=(self.n_components,)).to(var.device)
        
        for k in range(self.n_components):
            log_det[k] = 2 * torch.log(torch.diagonal(torch.linalg.cholesky(var[0,k]))).sum()

        return log_det.unsqueeze(-1)


    def _e_step(self, x):
        """
        Computes log-responses that indicate the (logarithmic) posterior belief (sometimes called responsibilities) that a data point was generated by one of the k mixture components.
        Also returns the mean of the mean of the logarithms of the probabilities (as is done in sklearn).
        This is the so-called expectation step of the EM-algorithm.
        args:
            x:              torch.Tensor (n, d) or (n, 1, d)
        returns:
            log_prob_norm:  torch.Tensor (1)
            log_resp:       torch.Tensor (n, k, 1)
        """
        x = self.check_size(x)

        weighted_log_prob = self._estimate_log_prob(x) + torch.log(self.pi)

        log_prob_norm = torch.logsumexp(weighted_log_prob, dim=1, keepdim=True)
        log_resp = weighted_log_prob - log_prob_norm

        return torch.mean(log_prob_norm), log_resp


    def _m_step(self, x, log_resp):
        """
        From the log-probabilities, computes new parameters pi, mu, var (that maximize the log-likelihood). This is the maximization step of the EM-algorithm.
        args:
            x:          torch.Tensor (n, d) or (n, 1, d)
            log_resp:   torch.Tensor (n, k, 1)
        returns:
            pi:         torch.Tensor (1, k, 1)
            mu:         torch.Tensor (1, k, d)
            var:        torch.Tensor (1, k, d)
        """
        x = self.check_size(x)

        resp = torch.exp(log_resp)

        pi = torch.sum(resp, dim=0, keepdim=True) + self.eps
        mu = torch.sum(resp * x, dim=0, keepdim=True) / pi

        if self.covariance_type == "full":
            eps = (torch.eye(self.n_features) * self.eps).to(x.device)
            var = torch.sum((x - mu).unsqueeze(-1).matmul((x - mu).unsqueeze(-2)) * resp.unsqueeze(-1), dim=0,
                            keepdim=True) / torch.sum(resp, dim=0, keepdim=True).unsqueeze(-1) + eps

        elif self.covariance_type == "diag":
            x2 = (resp * x * x).sum(0, keepdim=True) / pi
            mu2 = mu * mu
            xmu = (resp * mu * x).sum(0, keepdim=True) / pi
            var = x2 - 2 * xmu + mu2 + self.eps

        pi = pi / x.shape[0]

        return pi, mu, var


    def __em(self, x):
        """
        Performs one iteration of the expectation-maximization algorithm by calling the respective subroutines.
        args:
            x:          torch.Tensor (n, 1, d)
        """
        _, log_resp = self._e_step(x)
        pi, mu, var = self._m_step(x, log_resp)

        self.__update_pi(pi)
        self.__update_mu(mu)
        self.__update_var(var)


    def __score(self, x, as_average=True):
        """
        Computes the log-likelihood of the data under the model.
        args:
            x:                  torch.Tensor (n, 1, d)
            sum_data:           bool
        returns:
            score:              torch.Tensor (1)
            (or)
            per_sample_score:   torch.Tensor (n)

        """
        weighted_log_prob = self._estimate_log_prob(x) + torch.log(self.pi)
        per_sample_score = torch.logsumexp(weighted_log_prob, dim=1)

        if as_average:
            return per_sample_score.mean()
        else:
            return torch.squeeze(per_sample_score)


    def __update_mu(self, mu):
        """
        Updates mean to the provided value.
        args:
            mu:         torch.FloatTensor
        """
        assert mu.size() in [(self.n_components, self.n_features), (1, self.n_components, self.n_features)], "Input mu does not have required tensor dimensions (%i, %i) or (1, %i, %i)" % (self.n_components, self.n_features, self.n_components, self.n_features)

        if mu.size() == (self.n_components, self.n_features):
            self.mu = mu.unsqueeze(0)
        elif mu.size() == (1, self.n_components, self.n_features):
            self.mu.data = mu


    def __update_var(self, var):
        """
        Updates variance to the provided value.
        args:
            var:        torch.FloatTensor
        """
        if self.covariance_type == "full":
            assert var.size() in [(self.n_components, self.n_features, self.n_features), (1, self.n_components, self.n_features, self.n_features)], "Input var does not have required tensor dimensions (%i, %i, %i) or (1, %i, %i, %i)" % (self.n_components, self.n_features, self.n_features, self.n_components, self.n_features, self.n_features)

            if var.size() == (self.n_components, self.n_features, self.n_features):
                self.var = var.unsqueeze(0)
            elif var.size() == (1, self.n_components, self.n_features, self.n_features):
                self.var.data = var

        elif self.covariance_type == "diag":
            assert var.size() in [(self.n_components, self.n_features), (1, self.n_components, self.n_features)], "Input var does not have required tensor dimensions (%i, %i) or (1, %i, %i)" % (self.n_components, self.n_features, self.n_components, self.n_features)

            if var.size() == (self.n_components, self.n_features):
                self.var = var.unsqueeze(0)
            elif var.size() == (1, self.n_components, self.n_features):
                self.var.data = var


    def __update_pi(self, pi):
        """
        Updates pi to the provided value.
        args:
            pi:         torch.FloatTensor
        """
        assert pi.size() in [(1, self.n_components, 1)], "Input pi does not have required tensor dimensions (%i, %i, %i)" % (1, self.n_components, 1)

        self.pi.data = pi


    def get_kmeans_mu(self, x, n_centers, init_times=50, min_delta=1e-3):
        """
        Find an initial value for the mean. Requires a threshold min_delta for the k-means algorithm to stop iterating.
        The algorithm is repeated init_times often, after which the best centerpoint is returned.
        args:
            x:            torch.FloatTensor (n, d) or (n, 1, d)
            init_times:   init
            min_delta:    int
        """
        x = x.to(self.device)
        if len(x.size()) == 3:
            x = x.squeeze(1)
        x_min, x_max = x.min(), x.max()
        x = (x - x_min) / (x_max - x_min)
        
        min_cost = np.inf

        for i in range(init_times):
            tmp_center = x[np.random.choice(np.arange(x.shape[0]), size=n_centers, replace=False), ...]
            l2_dis = torch.norm((x.unsqueeze(1).repeat(1, n_centers, 1) - tmp_center), p=2, dim=2)
            l2_cls = torch.argmin(l2_dis, dim=1)

            cost = 0
            for c in range(n_centers):
                cost += torch.norm(x[l2_cls == c] - tmp_center[c], p=2, dim=1).mean()

            if cost < min_cost:
                min_cost = cost
                center = tmp_center

        delta = np.inf

        while delta > min_delta:
            l2_dis = torch.norm((x.unsqueeze(1).repeat(1, n_centers, 1) - center), p=2, dim=2)
            l2_cls = torch.argmin(l2_dis, dim=1)
            center_old = center.clone()

            for c in range(n_centers):
                center[c] = x[l2_cls == c].mean(dim=0)

            delta = torch.norm((center_old - center), dim=1).max()

        return (center.unsqueeze(0)*(x_max - x_min) + x_min)
    
    def get_kmeans_mu_plus_plus(self, x, n_centers, min_delta=1e-3):
        """
        Implementación de K-means++ para encontrar valores iniciales de centroides.
        La inicialización k-means++ selecciona centroides iniciales basándose en 
        la distancia a los centroides existentes.
        
        args:
            x: torch.FloatTensor (n, d) o (n, 1, d)
            n_centers: int, número de centroides
            min_delta: float, umbral para detener las iteraciones
        returns:
            centers: torch.FloatTensor (1, n_centers, d)
        """
        x = x.to(self.device)
        # Manejo inicial de dimensiones como en el método original
        if len(x.size()) == 3:
            x = x.squeeze(1)
        
        # Normalización de datos
        x_min, x_max = x.min(), x.max()
        x = (x - x_min) / (x_max - x_min)
        
        # Obtener dimensiones
        n_samples, n_features = x.size()
        
        # Inicializar lista de centroides
        centers = []
        
        # Seleccionar el primer centroide aleatoriamente
        first_center_idx = torch.randint(0, n_samples, (1,))
        centers.append(x[first_center_idx])
        
        # Seleccionar los centroides restantes usando k-means++
        for _ in range(1, n_centers):
            # Calcular distancias al cuadrado a los centroides más cercanos
            distances = torch.zeros(n_samples)
            
            # Convertir centros actuales a tensor
            current_centers = torch.cat(centers)
            
            # Para cada punto, encontrar la distancia mínima a cualquier centroide existente
            for i in range(n_samples):
                # Calcular distancias a todos los centroides existentes
                point_distances = torch.norm(x[i] - current_centers, p=2, dim=1)
                # Guardar la distancia mínima al cuadrado
                distances[i] = torch.min(point_distances) ** 2
                
            # Normalizar las distancias para crear una distribución de probabilidad
            distances = distances / torch.sum(distances)
            
            # Seleccionar el siguiente centro basado en las probabilidades
            next_center_idx = torch.multinomial(distances, 1)
            centers.append(x[next_center_idx])
        
        # Convertir lista de centroides a tensor
        center = torch.cat(centers)
        
        # Proceso de refinamiento iterativo (similar al método original)
        delta = np.inf
        
        while delta > min_delta:
            # Calcular distancias a todos los centroides
            l2_dis = torch.norm((x.unsqueeze(1).repeat(1, n_centers, 1) - center), p=2, dim=2)
            # Asignar puntos al centroide más cercano
            l2_cls = torch.argmin(l2_dis, dim=1)
            center_old = center.clone()
            
            # Actualizar centroides
            for c in range(n_centers):
                mask = l2_cls == c
                if torch.sum(mask) > 0:  # Evitar división por cero
                    center[c] = x[mask].mean(dim=0)
            
            # Calcular el cambio máximo en los centroides
            delta = torch.norm((center_old - center), dim=1).max()
        
        # Desnormalizar los centroides
        return (center.unsqueeze(0) * (x_max - x_min) + x_min)

def calculate_matmul_n_times(n_components, mat_a, mat_b):
    """
    Calculate matrix product of two matrics with mat_a[0] >= mat_b[0].
    Bypasses torch.matmul to reduce memory footprint.
    args:
        mat_a:      torch.Tensor (n, k, 1, d)
        mat_b:      torch.Tensor (1, k, d, d)
    """
    res = torch.zeros(mat_a.shape).to(mat_a.device)
    
    for i in range(n_components):
        mat_a_i = mat_a[:, i, :, :].squeeze(-2)
        mat_b_i = mat_b[0, i, :, :].squeeze()
        res[:, i, :, :] = mat_a_i.mm(mat_b_i).unsqueeze(1)
    
    return res

def calculate_matmul(mat_a, mat_b):
    """
    Calculate matrix product of two matrics with mat_a[0] >= mat_b[0].
    Bypasses torch.matmul to reduce memory footprint.
    args:
        mat_a:      torch.Tensor (n, k, 1, d)
        mat_b:      torch.Tensor (n, k, d, 1)
    """
    assert mat_a.shape[-2] == 1 and mat_b.shape[-1] == 1
    return torch.sum(mat_a.squeeze(-2) * mat_b.squeeze(-1), dim=2, keepdim=True)