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 cholesky_logdet_maha(const float* cov, const float* diff, int D, float* L) {
    for(int i=0; i<D*D; i++) L[i] = cov[i];
    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];
            L[i*D + j] = (L[i*D + j] - sum) / L[j*D + j];
        }
        float sum = 0.0f;
        for(int k=0; k<i; k++) sum += L[i*D + k] * L[i*D + k];
        L[i*D + i] = sqrtf(L[i*D + i] - sum);
    }
    float log_det = 0.0f;
    for(int i=0; i<D; i++) log_det += logf(L[i*D + i]);
    float y[16];
    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 i=0; i<D; i++) maha += y[i] * y[i];
    return log_det + maha;
}

__global__ void em_step_kernel(
    const float* __restrict__ x, const float* __restrict__ mus, const float* __restrict__ covs, const float* __restrict__ log_pi_in,
    float* __restrict__ new_mus, float* __restrict__ new_covs, float* __restrict__ new_log_pi, float* __restrict__ log_lik_out,
    int N, int D, int K
) {
    extern __shared__ float log_resp[];
    float log_pi[16];
    float sum_exp = 0.0f;
    for(int k=0; k<K; k++) sum_exp += expf(log_pi_in[k]);
    float log_sum = logf(sum_exp);
    for(int k=0; k<K; k++) log_pi[k] = log_pi_in[k] - log_sum;

    float total_log_lik = 0.0f;
    float L[256];
    float diff[16];
    for(int n=threadIdx.x; n<N; n += blockDim.x) {
        float max_lr = -1e30f;
        for(int k=0; k<K; k++) {
            for(int d=0; d<D; d++) diff[d] = x[n*D + d] - mus[k*D + d];
            float lm = cholesky_logdet_maha(covs + k*D*D, diff, D, L);
            float lr = -0.5f * (D * logf(2*3.14159265f) + lm) + log_pi[k];
            log_resp[n*K + k] = lr;
            if (lr > max_lr) max_lr = lr;
        }
        float sum = 0.0f;
        for(int k=0; k<K; k++) sum += expf(log_resp[n*K + k] - max_lr);
        float norm = max_lr + logf(sum);
        total_log_lik += norm;
        for(int k=0; k<K; k++) {
            log_resp[n*K + k] = expf(log_resp[n*K + k] - norm);
        }
    }
    // reduce log_lik, simple atomic for now
    atomicAdd(log_lik_out, total_log_lik);

    __syncthreads();
    // only thread 0 does M step
    if(threadIdx.x == 0) {
    float N_k[16];
    for(int k=0; k<K; k++) {
        N_k[k] = 0.0f;
        for(int n=0; n<N; n++) N_k[k] += log_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 += log_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 di = x[n*D + i] - new_mus[k*D + i];
                    float dj = x[n*D + j] - new_mus[k*D + j];
                    sum += log_resp[n*K + k] * di * dj;
                }
                new_covs[k*D*D + i*D + j] = sum / N_k[k];
            }
        }
        for(int d=0; d<D; d++) new_covs[k*D*D + d*D + d] += 1e-4f;
    }
    float Nlog = logf((float)N);
    for(int k=0; k<K; k++) new_log_pi[k] = logf(N_k[k]) - Nlog;
    }
}

std::vector<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 options = torch::TensorOptions().device(x.device()).dtype(torch::kFloat32);
    auto new_mus = torch::zeros({K, D}, options);
    auto new_covs = torch::zeros({K, D, D}, options);
    auto new_log_pi = torch::zeros({K}, options);
    auto log_lik_t = torch::zeros({}, options);
    int shmem = N * K * sizeof(float);
    int threads = 256;
    em_step_kernel<<<1, threads, shmem>>>(
        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_t.data_ptr<float>(),
        N, D, K);
    return {new_mus, new_covs, new_log_pi, log_lik_t};
}
"""

cpp_source = """
#include <vector>
std::vector<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=True,
    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:
        outs = module.em_forward(x.contiguous(), self.mus.contiguous(), self.covs.contiguous(), self.log_pi.contiguous())
        return tuple(outs)