import torch
import torch.nn as nn
import math


class Model(nn.Module):
    """
    One iteration of Expectation-Maximization for a Gaussian Mixture Model.
    E-step: compute soft assignments (responsibilities).
    M-step: update means, covariances, and mixing weights.
    The dominant cost is the batched evaluation of K Gaussian densities
    over N data points.
    """

    def __init__(self, dim, num_components):
        super().__init__()
        self.dim = dim
        self.K = num_components
        self.register_buffer("mus", torch.randn(num_components, dim))
        self.register_buffer("covs", torch.eye(dim).unsqueeze(0).expand(num_components, -1, -1).clone())
        self.register_buffer("log_pi", torch.zeros(num_components))

    def _log_gaussian(self, x, mu, cov):
        """Log N(x | mu, cov) for batched x."""
        D = mu.shape[-1]
        diff = x - mu.unsqueeze(0)  # (N, D)
        L = torch.linalg.cholesky(cov)
        solve = torch.linalg.solve_triangular(L, diff.unsqueeze(-1), upper=False).squeeze(-1)
        maha = (solve ** 2).sum(-1)
        log_det = 2.0 * L.diagonal().log().sum()
        return -0.5 * (D * math.log(2 * math.pi) + log_det + maha)

    def forward(self, x: torch.Tensor) -> tuple:
        """
        Args:
            x: (N, dim) – data points
        Returns:
            (new_mus, new_covs, new_log_pi, log_lik):
                new_mus:    (K, dim)
                new_covs:   (K, dim, dim)
                new_log_pi: (K,)
                log_lik:    scalar – total log-likelihood
        """
        N, D = x.shape
        K = self.K
        log_pi = self.log_pi - torch.logsumexp(self.log_pi, dim=0)

        # E-step
        log_resp = torch.stack([
            self._log_gaussian(x, self.mus[k], self.covs[k]) + log_pi[k]
            for k in range(K)
        ], dim=-1)  # (N, K)
        log_lik = torch.logsumexp(log_resp, dim=-1).sum()
        log_resp = log_resp - torch.logsumexp(log_resp, dim=-1, keepdim=True)
        resp = torch.exp(log_resp)  # (N, K)

        # M-step
        N_k = resp.sum(dim=0).clamp(min=1e-8)  # (K,)
        new_mus = (resp.t() @ x) / N_k.unsqueeze(-1)
        new_covs = torch.zeros_like(self.covs)
        for k in range(K):
            diff = x - new_mus[k].unsqueeze(0)
            weighted = diff * resp[:, k:k+1]
            new_covs[k] = (weighted.t() @ diff) / N_k[k] + 1e-4 * torch.eye(D, device=x.device)

        new_log_pi = N_k.log() - math.log(N)

        return new_mus, new_covs, new_log_pi, log_lik


dim = 8
num_components = 5
num_data = 2048


def get_inputs():
    p = popcorn_pri
    return [torch.randn(num_data, dim)]


def get_init_inputs():
    return [dim, num_components]

import torch
import torch.nn as nn
from torch.utils.cpp_extension import load_inline
import math

cuda_source = """
#include <torch/extension.h>
#include <cuda.h>
#include <cuda_runtime.h>
#include <math.h>

__device__ float log_gaussian(const float* x, const float* mu, const float* cov, int D, float log2pi) {
    float diff[8];
    for(int d=0; d<D; d++) diff[d] = x[d] - mu[d];
    float L[8*8] = {0};
    for(int i=0; i<D; i++) {
        for(int j=0; j<=i; j++) {
            float sum = 0.0f;
            for(int k=0; k<j; k++) sum += L[i*D + k] * L[j*D + k];
            if(i==j) {
                L[i*D + j] = sqrtf(cov[i*D + j] - sum);
            } else {
                L[i*D + j] = (cov[i*D + j] - sum) / L[j*D + j];
            }
        }
    }
    float log_det = 0.0f;
    for(int d=0; d<D; d++) log_det += logf(L[d*D + d]);
    log_det *= 2.0f;
    float y[8];
    for(int i=0; i<D; i++) {
        float sum = 0.0f;
        for(int j=0; j<i; j++) sum += L[i*D + j] * y[j];
        y[i] = (diff[i] - sum) / L[i*D + i];
    }
    float maha = 0.0f;
    for(int d=0; d<D; d++) maha += y[d]*y[d];
    return -0.5f * (D * log2pi + log_det + maha);
}

__global__ void em_step_kernel(
    const float* x, const float* mus, const float* covs, const float* log_pi_in,
    float* new_mus, float* new_covs, float* new_log_pi, float* log_lik_out,
    int N, int D, int K
) {
    if(threadIdx.x == 0) {
        float log2pi = logf(2.0f * 3.1415926535f);
        float max_log = -1e30f;
        for(int k=0; k<K; k++) if(log_pi_in[k] > max_log) max_log = log_pi_in[k];
        float sum_exp = 0.0f;
        for(int k=0; k<K; k++) sum_exp += expf(log_pi_in[k] - max_log);
        float log_sum = max_log + logf(sum_exp);
        float log_pi[5];
        for(int k=0; k<K; k++) log_pi[k] = log_pi_in[k] - log_sum;
        float log_resp[2048*5];
        float total_log_lik = 0.0f;
        for(int n=0; n<N; n++) {
            float maxr = -1e30f;
            for(int k=0; k<K; k++) {
                float lg = log_gaussian(&x[n*D], &mus[k*D], &covs[k*D*D], D, log2pi);
                log_resp[n*K + k] = lg + log_pi[k];
                if(log_resp[n*K + k] > maxr) maxr = log_resp[n*K + k];
            }
            float sumr = 0.0f;
            for(int k=0; k<K; k++) sumr += expf(log_resp[n*K + k] - maxr);
            float lse = maxr + logf(sumr);
            total_log_lik += lse;
            for(int k=0; k<K; k++) log_resp[n*K + k] -= lse;
        }
        *log_lik_out = total_log_lik;
        float resp[2048*5];
        for(int i=0; i<N*K; i++) resp[i] = expf(log_resp[i]);
        float N_k[5] = {0};
        for(int k=0; k<K; k++) {
            for(int n=0; n<N; n++) N_k[k] += resp[n*K + k];
            if(N_k[k] < 1e-8f) N_k[k] = 1e-8f;
        }
        for(int k=0; k<K; k++) {
            for(int d=0; d<D; d++) {
                float sum = 0.0f;
                for(int n=0; n<N; n++) sum += resp[n*K + k] * x[n*D + d];
                new_mus[k*D + d] = sum / N_k[k];
            }
        }
        for(int k=0; k<K; k++) {
            for(int i=0; i<D; i++) for(int j=0; j<D; j++) {
                float sum = 0.0f;
                for(int n=0; n<N; n++) {
                    float diff_i = x[n*D + i] - new_mus[k*D + i];
                    float diff_j = x[n*D + j] - new_mus[k*D + j];
                    sum += resp[n*K + k] * diff_i * diff_j;
                }
                new_covs[k*D*D + i*D + j] = sum / N_k[k] + (i==j ? 1e-4f : 0.0f);
            }
        }
        float log_N = logf((float)N);
        for(int k=0; k<K; k++) new_log_pi[k] = logf(N_k[k]) - log_N;
    }
}

std::tuple<torch::Tensor, torch::Tensor, torch::Tensor, torch::Tensor> em_forward(
    torch::Tensor x, torch::Tensor mus, torch::Tensor covs, torch::Tensor log_pi
) {
    int N = x.size(0);
    int D = x.size(1);
    int K = mus.size(0);
    auto new_mus = torch::zeros({K, D}, x.options());
    auto new_covs = torch::zeros({K, D, D}, x.options());
    auto new_log_pi = torch::zeros({K}, x.options());
    auto log_lik = torch::zeros({1}, x.options());
    em_step_kernel<<<1, 1, 0>>>(x.data_ptr<float>(), mus.data_ptr<float>(), covs.data_ptr<float>(), log_pi.data_ptr<float>(),
        new_mus.data_ptr<float>(), new_covs.data_ptr<float>(), new_log_pi.data_ptr<float>(), log_lik.data_ptr<float>(),
        N, D, K);
    return std::make_tuple(new_mus, new_covs, new_log_pi, log_lik.squeeze());
}
"""

cpp_source = """
std::tuple<torch::Tensor, torch::Tensor, torch::Tensor, torch::Tensor> em_forward(
    torch::Tensor x, torch::Tensor mus, torch::Tensor covs, torch::Tensor log_pi);
"""

module = load_inline(
    name="gmm_em",
    cpp_sources=cpp_source,
    cuda_sources=cuda_source,
    functions=["em_forward"],
    verbose=False,
    extra_cuda_cflags=["-O3"]
)

class ModelNew(nn.Module):
    def __init__(self, dim, num_components):
        super().__init__()
        self.dim = dim
        self.K = num_components
        self.register_buffer("mus", torch.randn(num_components, dim))
        self.register_buffer("covs", torch.eye(dim).unsqueeze(0).expand(num_components, -1, -1).clone())
        self.register_buffer("log_pi", torch.zeros(num_components))

    def forward(self, x: torch.Tensor) -> tuple:
        return module.em_forward(x, self.mus, self.covs, self.log_pi)