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
import math
from torch.utils.cpp_extension import load_inline

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

constexpr int BLOCK_SIZE = 256;
constexpr int TILE_SIZE = 16;

// Cholesky decomposition for small matrices (dim x dim)
__global__ void cholesky_kernel(const float* __restrict__ covs, float* __restrict__ L, int K, int D) {
    int k = blockIdx.x;
    if (k >= K) return;
    
    const float* cov = covs + k * D * D;
    float* l = L + k * D * D;
    
    // Load into shared memory
    __shared__ float A[64];  // Max dim=8, so 64 floats
    __shared__ float L_local[64];
    
    for (int i = threadIdx.x; i < D * D; i += blockDim.x) {
        A[i] = cov[i];
        L_local[i] = 0.0f;
    }
    __syncthreads();
    
    // Cholesky: L[i][j] = (A[i][j] - sum(L[i][m]*L[j][m])) / L[j][j]
    // Only one thread does the computation for small matrices
    if (threadIdx.x == 0) {
        for (int i = 0; i < D; i++) {
            for (int j = 0; j <= i; j++) {
                float sum = 0.0f;
                for (int m = 0; m < j; m++) {
                    sum += L_local[i * D + m] * L_local[j * D + m];
                }
                if (i == j) {
                    float val = A[i * D + i] - sum;
                    L_local[i * D + i] = sqrtf(fmaxf(val, 1e-8f));
                } else {
                    L_local[i * D + j] = (A[i * D + j] - sum) / L_local[j * D + j];
                }
            }
        }
    }
    __syncthreads();
    
    // Copy back
    for (int i = threadIdx.x; i < D * D; i += blockDim.x) {
        l[i] = L_local[i];
    }
}

// Compute log-gaussian for all components: returns log_resp (N, K)
// Also returns log_det (K,) for each component
__global__ void log_gaussian_kernel(
    const float* __restrict__ x,      // (N, D)
    const float* __restrict__ mus,   // (K, D)
    const float* __restrict__ L,     // (K, D, D) - Cholesky factors
    const float* __restrict__ log_pi, // (K,)
    float* __restrict__ log_resp,     // (N, K)
    int N, int D, int K
) {
    int n = blockIdx.x * blockDim.x + threadIdx.x;
    if (n >= N) return;
    
    // Load x[n] into registers
    float x_local[8];  // Max D=8
    for (int d = 0; d < D; d++) {
        x_local[d] = x[n * D + d];
    }
    
    // For each component k
    for (int k = 0; k < K; k++) {
        const float* mu = mus + k * D;
        const float* l = L + k * D * D;
        
        // Compute diff = x - mu
        float diff[8];
        for (int d = 0; d < D; d++) {
            diff[d] = x_local[d] - mu[d];
        }
        
        // Solve L @ solve = diff (lower triangular)
        float solve[8];
        for (int i = 0; i < D; i++) {
            float sum = diff[i];
            for (int j = 0; j < i; j++) {
                sum -= l[i * D + j] * solve[j];
            }
            solve[i] = sum / l[i * D + i];
        }
        
        // Mahalanobis distance = sum(solve^2)
        float maha = 0.0f;
        for (int d = 0; d < D; d++) {
            maha += solve[d] * solve[d];
        }
        
        // Log-det from L diagonal (precomputed in another kernel or here)
        // We need to compute this efficiently
        
        // Log probability
        // log_p = -0.5 * (D * log(2*pi) + log_det + maha)
        // We'll compute log_det separately
        
        log_resp[n * K + k] = -0.5f * maha + log_pi[k];
    }
}

// Compute log-determinant from L diagonal
__global__ void compute_log_det_kernel(
    const float* __restrict__ L,
    float* __restrict__ log_det,
    int K, int D
) {
    int k = blockIdx.x;
    if (k >= K) return;
    
    const float* l = L + k * D * D;
    float sum = 0.0f;
    for (int d = 0; d < D; d++) {
        sum += logf(l[d * D + d]);
    }
    log_det[k] = 2.0f * sum;  // log(det(A)) = 2 * sum(log(diag(L)))
}

// Add log_det and constant to log_resp
__global__ void add_log_det_kernel(
    float* __restrict__ log_resp,
    const float* __restrict__ log_det,
    int N, int K, int D
) {
    int n = blockIdx.x * blockDim.x + threadIdx.x;
    if (n >= N) return;
    
    float log_const = -0.5f * D * logf(2.0f * 3.14159265358979323846f);
    
    for (int k = 0; k < K; k++) {
        log_resp[n * K + k] += log_const - 0.5f * log_det[k];
    }
}

// Softmax: log_resp = log_resp - logsumexp(log_resp, dim=-1)
__global__ void log_softmax_kernel(
    float* __restrict__ log_resp,
    int N, int K
) {
    int n = blockIdx.x * blockDim.x + threadIdx.x;
    if (n >= N) return;
    
    float* row = log_resp + n * K;
    
    // Find max for numerical stability
    float max_val = row[0];
    for (int k = 1; k < K; k++) {
        max_val = fmaxf(max_val, row[k]);
    }
    
    // Compute logsumexp
    float sum = 0.0f;
    for (int k = 0; k < K; k++) {
        sum += expf(row[k] - max_val);
    }
    float log_sum = max_val + logf(sum);
    
    // Subtract logsumexp
    for (int k = 0; k < K; k++) {
        row[k] -= log_sum;
    }
}

// Compute total log-likelihood (reduction)
__global__ void logsumexp_reduce_kernel(
    const float* __restrict__ log_resp,
    float* __restrict__ partial_sums,
    int N, int K
) {
    __shared__ float sdata[BLOCK_SIZE];
    
    int tid = threadIdx.x;
    int n = blockIdx.x * blockDim.x + threadIdx.x;
    
    float val = 0.0f;
    if (n < N) {
        // Compute logsumexp for this row
        const float* row = log_resp + n * K;
        float max_val = row[0];
        for (int k = 1; k < K; k++) {
            max_val = fmaxf(max_val, row[k]);
        }
        float sum = 0.0f;
        for (int k = 0; k < K; k++) {
            sum += expf(row[k] - max_val);
        }
        val = max_val + logf(sum);
    }
    
    sdata[tid] = val;
    __syncthreads();
    
    // Block reduction
    for (int s = blockDim.x >> 1; s > 0; s >>= 1) {
        if (tid < s) {
            sdata[tid] += sdata[tid + s];
        }
        __syncthreads();
    }
    
    if (tid == 0) {
        partial_sums[blockIdx.x] = sdata[0];
    }
}

// M-step: compute new_mus = (resp^T @ x) / N_k
__global__ void m_step_mus_kernel(
    const float* __restrict__ resp,
    const float* __restrict__ x,
    float* __restrict__ new_mus,
    float* __restrict__ N_k,
    int N, int D, int K
) {
    int k = blockIdx.x;
    int d = threadIdx.x;
    
    if (k >= K || d >= D) return;
    
    float sum = 0.0f;
    float n_sum = 0.0f;
    
    for (int n = 0; n < N; n++) {
        float r = resp[n * K + k];
        sum += r * x[n * D + d];
        if (d == 0) {
            n_sum += r;
        }
    }
    
    if (d == 0) {
        N_k[k] = fmaxf(n_sum, 1e-8f);
    }
    __syncthreads();
    
    new_mus[k * D + d] = sum / N_k[k];
}

// M-step: compute new_covs[k] = (diff^T @ weighted_diff) / N_k[k] + reg
__global__ void m_step_covs_kernel(
    const float* __restrict__ x,
    const float* __restrict__ new_mus,
    const float* __restrict__ resp,
    const float* __restrict__ N_k,
    float* __restrict__ new_covs,
    int N, int D, int K
) {
    int k = blockIdx.x;
    if (k >= K) return;
    
    __shared__ float mu_local[8];
    __shared__ float cov_local[64];
    
    // Load mu to shared memory
    if (threadIdx.x < D) {
        mu_local[threadIdx.x] = new_mus[k * D + threadIdx.x];
    }
    
    // Initialize cov to 0
    for (int i = threadIdx.x; i < D * D; i += blockDim.x) {
        cov_local[i] = 0.0f;
    }
    __syncthreads();
    
    // Accumulate outer products
    for (int n = threadIdx.x; n < N; n += blockDim.x) {
        float r = resp[n * K + k];
        float diff[8];
        for (int d = 0; d < D; d++) {
            diff[d] = x[n * D + d] - mu_local[d];
        }
        
        // Outer product weighted by r
        for (int i = 0; i < D; i++) {
            for (int j = 0; j < D; j++) {
                atomicAdd(&cov_local[i * D + j], r * diff[i] * diff[j]);
            }
        }
    }
    __syncthreads();
    
    // Normalize and add regularization
    float n_k = N_k[k];
    for (int i = threadIdx.x; i < D * D; i += blockDim.x) {
        int row = i / D;
        int col = i % D;
        float val = cov_local[i] / n_k;
        if (row == col) {
            val += 1e-4f;
        }
        new_covs[k * D * D + i] = val;
    }
}

// Compute new_log_pi = log(N_k) - log(N)
__global__ void compute_log_pi_kernel(
    const float* __restrict__ N_k,
    float* __restrict__ new_log_pi,
    int K, int N
) {
    int k = blockIdx.x * blockDim.x + threadIdx.x;
    if (k >= K) return;
    
    new_log_pi[k] = logf(N_k[k]) - logf((float)N);
}

// Host function to launch kernels
std::tuple<torch::Tensor, torch::Tensor, torch::Tensor, torch::Tensor> 
gmm_em_forward_cuda(
    torch::Tensor x,
    torch::Tensor mus,
    torch::Tensor covs,
    torch::Tensor log_pi,
    int K
) {
    int N = x.size(0);
    int D = x.size(1);
    
    auto options = x.options();
    
    // Compute Cholesky decomposition
    auto L = torch::empty({K, D, D}, options);
    int cholesky_blocks = K;
    cholesky_kernel<<<cholesky_blocks, BLOCK_SIZE>>>(
        covs.data_ptr<float>(), L.data_ptr<float>(), K, D
    );
    C10_CUDA_CHECK(cudaGetLastError());
    
    // Compute log_det from L
    auto log_det = torch::empty({K}, options);
    compute_log_det_kernel<<<K, BLOCK_SIZE>>>(
        L.data_ptr<float>(), log_det.data_ptr<float>(), K, D
    );
    C10_CUDA_CHECK(cudaGetLastError());
    
    // Normalize log_pi
    auto log_pi_data = log_pi.data_ptr<float>();
    // This should be done on CPU or in a kernel
    
    // Compute log_resp
    auto log_resp = torch::empty({N, K}, options);
    int log_gaussian_blocks = (N + BLOCK_SIZE - 1) / BLOCK_SIZE;
    log_gaussian_kernel<<<log_gaussian_blocks, BLOCK_SIZE>>>(
        x.data_ptr<float>(),
        mus.data_ptr<float>(),
        L.data_ptr<float>(),
        log_pi.data_ptr<float>(),
        log_resp.data_ptr<float>(),
        N, D, K
    );
    C10_CUDA_CHECK(cudaGetLastError());
    
    // Add log_det and constant
    add_log_det_kernel<<<log_gaussian_blocks, BLOCK_SIZE>>>(
        log_resp.data_ptr<float>(),
        log_det.data_ptr<float>(),
        N, K, D
    );
    C10_CUDA_CHECK(cudaGetLastError());
    
    // Compute total log-likelihood
    int num_blocks = (N + BLOCK_SIZE - 1) / BLOCK_SIZE;
    auto partial_sums = torch::empty({num_blocks}, options);
    logsumexp_reduce_kernel<<<num_blocks, BLOCK_SIZE>>>(
        log_resp.data_ptr<float>(),
        partial_sums.data_ptr<float>(),
        N, K
    );
    C10_CUDA_CHECK(cudaGetLastError());
    
    float log_lik_val = partial_sums.sum().item<float>();
    auto log_lik = torch::tensor(log_lik_val, options);
    
    // Normalize log_resp
    log_softmax_kernel<<<log_gaussian_blocks, BLOCK_SIZE>>>(
        log_resp.data_ptr<float>(),
        N, K
    );
    C10_CUDA_CHECK(cudaGetLastError());
    
    // Compute responsibilities
    auto resp = torch::exp(log_resp);
    
    // M-step: compute N_k and new_mus
    auto N_k = torch::empty({K}, options);
    auto new_mus = torch::empty({K, D}, options);
    m_step_mus_kernel<<<K, D>>>(
        resp.data_ptr<float>(),
        x.data_ptr<float>(),
        new_mus.data_ptr<float>(),
        N_k.data_ptr<float>(),
        N, D, K
    );
    C10_CUDA_CHECK(cudaGetLastError());
    
    // M-step: compute new_covs
    auto new_covs = torch::empty({K, D, D}, options);
    m_step_covs_kernel<<<K, BLOCK_SIZE>>>(
        x.data_ptr<float>(),
        new_mus.data_ptr<float>(),
        resp.data_ptr<float>(),
        N_k.data_ptr<float>(),
        new_covs.data_ptr<float>(),
        N, D, K
    );
    C10_CUDA_CHECK(cudaGetLastError());
    
    // Compute new_log_pi
    auto new_log_pi = torch::empty({K}, options);
    compute_log_pi_kernel<<<1, K>>>(
        N_k.data_ptr<float>(),
        new_log_pi.data_ptr<float>(),
        K, N
    );
    C10_CUDA_CHECK(cudaGetLastError());
    
    return std::make_tuple(new_mus, new_covs, new_log_pi, log_lik);
}
"""

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

gmm_module = load_inline(
    name='gmm_em_cuda',
    cpp_sources=cpp_source,
    cuda_sources=cuda_source,
    functions=['gmm_em_forward_cuda'],
    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:
        # Normalize log_pi
        log_pi = self.log_pi - torch.logsumexp(self.log_pi, dim=0)
        
        # Call CUDA kernel
        new_mus, new_covs, new_log_pi, log_lik = gmm_module.gmm_em_forward_cuda(
            x, self.mus, self.covs, log_pi, self.K
        )
        
        return new_mus, new_covs, new_log_pi, log_lik