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

# Define the CUDA kernel with proper C++ wrappers
cuda_source = """
#include <torch/extension.h>
#include <cuda.h>
#include <cuda_runtime.h>
#include <cmath>
#include <c10/cuda/CUDAException.h>

constexpr int BLOCK_SIZE = 256;
constexpr double PI = 3.14159265358979323846;

// Forward substitution: solve L * y = b where L is lower triangular
__device__ __forceinline__ void forward_substitution_inline(const double* L, const double* b, double* y, int D) {
    #pragma unroll
    for (int i = 0; i < 8; i++) {
        if (i >= D) break;
        double sum = b[i];
        #pragma unroll
        for (int j = 0; j < 8; j++) {
            if (j >= i) break;
            sum -= L[i * 8 + j] * y[j];
        }
        y[i] = sum / L[i * 8 + i];
    }
}

// Kernel for E-step with precomputed Cholesky
// Uses shared memory for mus and L matrices
__global__ void e_step_optimized_kernel(
    const double* __restrict__ x,        // (N, D)
    const double* __restrict__ mus,      // (K, D)
    const double* __restrict__ L_chol,   // (K, D, D) - precomputed Cholesky L
    const double* __restrict__ log_det,  // (K,) - precomputed log determinants
    const double* __restrict__ log_pi,   // (K,)
    double* __restrict__ log_resp,       // (N, K)
    int N, int D, int K
) {
    extern __shared__ double shared_data[];
    
    // Shared memory layout: 
    // mus[K][D] + L_chol[K][D*D] + log_det[K] + log_pi[K]
    double* s_mus = shared_data;
    double* s_L = s_mus + K * D;
    double* s_log_det = s_L + K * D * D;
    double* s_log_pi = s_log_det + K;
    
    // Collaborative load of shared memory
    int tid = threadIdx.x;
    int total_shared = K * D + K * D * D + K + K;
    
    for (int i = tid; i < total_shared; i += blockDim.x) {
        if (i < K * D) {
            s_mus[i] = mus[i];
        } else if (i < K * D + K * D * D) {
            s_L[i - K * D] = L_chol[i - K * D];
        } else if (i < K * D + K * D * D + K) {
            s_log_det[i - K * D - K * D * D] = log_det[i - K * D - K * D * D];
        } else {
            s_log_pi[i - K * D - K * D * D - K] = log_pi[i - K * D - K * D * D - K];
        }
    }
    __syncthreads();
    
    // Each thread handles one (n, k) pair
    int n = blockIdx.x * blockDim.x + threadIdx.x;
    int k = blockIdx.y;
    
    if (n >= N || k >= K) return;
    
    // Get pointers to shared data for component k
    const double* mu_k = s_mus + k * D;
    const double* L_k = s_L + k * D * D;
    double log_det_k = s_log_det[k];
    double log_pi_k = s_log_pi[k];
    
    // Load x[n] into registers - coalesced access
    double x_local[8];
    for (int d = 0; d < D; d++) {
        x_local[d] = x[n * D + d];
    }
    
    // Compute diff = x[n] - mu[k]
    double diff[8];
    for (int d = 0; d < D; d++) {
        diff[d] = x_local[d] - mu_k[d];
    }
    
    // Solve L * y = diff using forward substitution
    double y[8];
    forward_substitution_inline(L_k, diff, y, D);
    
    // Compute Mahalanobis distance
    double maha = 0.0;
    for (int d = 0; d < D; d++) {
        maha += y[d] * y[d];
    }
    
    // Log Gaussian
    double log_gauss = -0.5 * (D * log(2.0 * PI) + log_det_k + maha);
    log_resp[n * K + k] = log_gauss + log_pi_k;
}

// Kernel to compute logsumexp and normalize log_resp
__global__ void logsumexp_normalize_kernel(
    double* __restrict__ log_resp,    // (N, K)
    double* __restrict__ log_lik,     // scalar output
    double* __restrict__ resp,        // (N, K) output
    int N, int K
) {
    int n = blockIdx.x * blockDim.x + threadIdx.x;
    if (n >= N) return;
    
    // Find max for numerical stability
    double max_val = log_resp[n * K];
    for (int k = 1; k < K; k++) {
        max_val = fmax(max_val, log_resp[n * K + k]);
    }
    
    // Compute sum of exp(log_resp - max_val)
    double sum = 0.0;
    for (int k = 0; k < K; k++) {
        sum += exp(log_resp[n * K + k] - max_val);
    }
    
    double log_sum = max_val + log(sum);
    
    // Normalize and convert to resp
    for (int k = 0; k < K; k++) {
        resp[n * K + k] = exp(log_resp[n * K + k] - log_sum);
        log_resp[n * K + k] -= log_sum;
    }
    
    // Add to log likelihood using atomic add
    atomicAdd(log_lik, log_sum);
}

// M-step kernel: compute N_k
__global__ void compute_Nk_kernel(
    const double* __restrict__ resp,  // (N, K)
    double* __restrict__ N_k,         // (K,)
    int N, int K
) {
    int k = blockIdx.x * blockDim.x + threadIdx.x;
    if (k >= K) return;
    
    double sum = 0.0;
    for (int n = 0; n < N; n++) {
        sum += resp[n * K + k];
    }
    N_k[k] = fmax(sum, 1e-8);
}

// M-step kernel: compute new_mus
__global__ void compute_new_mus_kernel(
    const double* __restrict__ x,     // (N, D)
    const double* __restrict__ resp,  // (N, K)
    const double* __restrict__ N_k,   // (K,)
    double* __restrict__ new_mus,     // (K, D)
    int N, int D, int K
) {
    int k = blockIdx.x;
    int d = threadIdx.x;
    
    if (k >= K || d >= D) return;
    
    double sum = 0.0;
    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];
}

// M-step kernel: compute new_covs
__global__ void compute_new_covs_kernel(
    const double* __restrict__ x,        // (N, D)
    const double* __restrict__ resp,     // (N, K)
    const double* __restrict__ new_mus,  // (K, D)
    const double* __restrict__ N_k,      // (K,)
    double* __restrict__ new_covs,       // (K, D, D)
    int N, int D, int K
) {
    int k = blockIdx.x;
    int idx = threadIdx.x;
    int i = idx / D;
    int j = idx % D;
    
    if (k >= K || i >= D || j >= D) return;
    
    double sum = 0.0;
    for (int n = 0; n < N; n++) {
        double diff_i = x[n * D + i] - new_mus[k * D + i];
        double 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];
    
    // Add regularization to diagonal
    if (i == j) {
        new_covs[k * D * D + i * D + j] += 1e-4;
    }
}

// M-step kernel: compute new_log_pi
__global__ void compute_new_log_pi_kernel(
    const double* __restrict__ N_k,      // (K,)
    double* __restrict__ new_log_pi,     // (K,)
    int N, int K
) {
    int k = blockIdx.x * blockDim.x + threadIdx.x;
    if (k >= K) return;
    
    new_log_pi[k] = log(N_k[k]) - log((double)N);
}

// C++ wrapper functions
torch::Tensor e_step_optimized_cuda(torch::Tensor x, torch::Tensor mus, torch::Tensor L_chol, 
                                     torch::Tensor log_det, torch::Tensor log_pi, int N, int D, int K) {
    auto options = torch::TensorOptions().dtype(torch::kFloat64).device(x.device());
    auto log_resp = torch::empty({N, K}, options);
    
    dim3 grid((N + BLOCK_SIZE - 1) / BLOCK_SIZE, K);
    dim3 block(BLOCK_SIZE);
    
    // Shared memory size: mus(K*D) + L(K*D*D) + log_det(K) + log_pi(K)
    size_t shared_mem_size = (K * D + K * D * D + K + K) * sizeof(double);
    
    e_step_optimized_kernel<<<grid, block, shared_mem_size>>>(
        x.data_ptr<double>(),
        mus.data_ptr<double>(),
        L_chol.data_ptr<double>(),
        log_det.data_ptr<double>(),
        log_pi.data_ptr<double>(),
        log_resp.data_ptr<double>(),
        N, D, K
    );
    
    C10_CUDA_CHECK(cudaGetLastError());
    return log_resp;
}

std::tuple<torch::Tensor, torch::Tensor> logsumexp_cuda(torch::Tensor log_resp, int N, int K) {
    auto options = torch::TensorOptions().dtype(torch::kFloat64).device(log_resp.device());
    auto resp = torch::empty({N, K}, options);
    auto log_lik = torch::zeros({}, options);
    
    dim3 grid((N + BLOCK_SIZE - 1) / BLOCK_SIZE);
    dim3 block(BLOCK_SIZE);
    
    logsumexp_normalize_kernel<<<grid, block>>>(
        log_resp.data_ptr<double>(),
        log_lik.data_ptr<double>(),
        resp.data_ptr<double>(),
        N, K
    );
    
    C10_CUDA_CHECK(cudaGetLastError());
    return std::make_tuple(log_lik, resp);
}

torch::Tensor compute_Nk_cuda(torch::Tensor resp, int N, int K) {
    auto options = torch::TensorOptions().dtype(torch::kFloat64).device(resp.device());
    auto N_k = torch::empty({K}, options);
    
    dim3 grid((K + BLOCK_SIZE - 1) / BLOCK_SIZE);
    dim3 block(BLOCK_SIZE);
    
    compute_Nk_kernel<<<grid, block>>>(
        resp.data_ptr<double>(),
        N_k.data_ptr<double>(),
        N, K
    );
    
    C10_CUDA_CHECK(cudaGetLastError());
    return N_k;
}

torch::Tensor compute_new_mus_cuda(torch::Tensor x, torch::Tensor resp, torch::Tensor N_k, int N, int D, int K) {
    auto options = torch::TensorOptions().dtype(torch::kFloat64).device(x.device());
    auto new_mus = torch::empty({K, D}, options);
    
    dim3 grid(K);
    dim3 block(D);
    
    compute_new_mus_kernel<<<grid, block>>>(
        x.data_ptr<double>(),
        resp.data_ptr<double>(),
        N_k.data_ptr<double>(),
        new_mus.data_ptr<double>(),
        N, D, K
    );
    
    C10_CUDA_CHECK(cudaGetLastError());
    return new_mus;
}

torch::Tensor compute_new_covs_cuda(torch::Tensor x, torch::Tensor resp, torch::Tensor new_mus, torch::Tensor N_k, int N, int D, int K) {
    auto options = torch::TensorOptions().dtype(torch::kFloat64).device(x.device());
    auto new_covs = torch::empty({K, D, D}, options);
    
    dim3 grid(K);
    dim3 block(D * D);
    
    compute_new_covs_kernel<<<grid, block>>>(
        x.data_ptr<double>(),
        resp.data_ptr<double>(),
        new_mus.data_ptr<double>(),
        N_k.data_ptr<double>(),
        new_covs.data_ptr<double>(),
        N, D, K
    );
    
    C10_CUDA_CHECK(cudaGetLastError());
    return new_covs;
}

torch::Tensor compute_new_log_pi_cuda(torch::Tensor N_k, int N, int K) {
    auto options = torch::TensorOptions().dtype(torch::kFloat64).device(N_k.device());
    auto new_log_pi = torch::empty({K}, options);
    
    dim3 grid((K + BLOCK_SIZE - 1) / BLOCK_SIZE);
    dim3 block(BLOCK_SIZE);
    
    compute_new_log_pi_kernel<<<grid, block>>>(
        N_k.data_ptr<double>(),
        new_log_pi.data_ptr<double>(),
        N, K
    );
    
    C10_CUDA_CHECK(cudaGetLastError());
    return new_log_pi;
}
"""

cpp_source = """
torch::Tensor e_step_optimized_cuda(torch::Tensor x, torch::Tensor mus, torch::Tensor L_chol, 
                                     torch::Tensor log_det, torch::Tensor log_pi, int N, int D, int K);
std::tuple<torch::Tensor, torch::Tensor> logsumexp_cuda(torch::Tensor log_resp, int N, int K);
torch::Tensor compute_Nk_cuda(torch::Tensor resp, int N, int K);
torch::Tensor compute_new_mus_cuda(torch::Tensor x, torch::Tensor resp, torch::Tensor N_k, int N, int D, int K);
torch::Tensor compute_new_covs_cuda(torch::Tensor x, torch::Tensor resp, torch::Tensor new_mus, torch::Tensor N_k, int N, int D, int K);
torch::Tensor compute_new_log_pi_cuda(torch::Tensor N_k, int N, int K);
"""

# Load the extension
gmm_extension = load_inline(
    name='gmm_extension_v4',
    cpp_sources=cpp_source,
    cuda_sources=cuda_source,
    functions=['e_step_optimized_cuda', 'logsumexp_cuda', 'compute_Nk_cuda', 'compute_new_mus_cuda', 
               'compute_new_covs_cuda', 'compute_new_log_pi_cuda'],
    extra_cuda_cflags=['-O3'],
    verbose=True
)


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))
        
        # Cached Cholesky decompositions
        self.register_buffer("L_chol", torch.zeros(num_components, dim, dim))
        self.register_buffer("log_det_chol", torch.zeros(num_components))
        self._update_chol_cache()

    def _update_chol_cache(self):
        """Precompute Cholesky decompositions and log determinants."""
        K, D = self.K, self.dim
        covs_d = self.covs.double()
        for k in range(K):
            L = torch.linalg.cholesky(covs_d[k])
            self.L_chol[k] = L
            self.log_det_chol[k] = 2.0 * L.diagonal().log().sum()

    def forward(self, x: torch.Tensor) -> tuple:
        N, D = x.shape
        K = self.K
        
        # Normalize log_pi using PyTorch (this is a small K-sized operation)
        log_pi = self.log_pi - torch.logsumexp(self.log_pi, dim=0)
        
        # Convert to double for precision
        x_d = x.double().contiguous()
        mus_d = self.mus.double().contiguous()
        L_chol_d = self.L_chol.double().contiguous()
        log_det_d = self.log_det_chol.double().contiguous()
        log_pi_d = log_pi.double().contiguous()
        
        # E-step using custom CUDA kernels
        log_resp = gmm_extension.e_step_optimized_cuda(x_d, mus_d, L_chol_d, log_det_d, log_pi_d, N, D, K)
        log_lik, resp = gmm_extension.logsumexp_cuda(log_resp, N, K)
        
        # M-step using custom CUDA kernels
        N_k = gmm_extension.compute_Nk_cuda(resp, N, K)
        new_mus = gmm_extension.compute_new_mus_cuda(x_d, resp, N_k, N, D, K)
        new_covs = gmm_extension.compute_new_covs_cuda(x_d, resp, new_mus, N_k, N, D, K)
        new_log_pi = gmm_extension.compute_new_log_pi_cuda(N_k, N, K)
        
        # Convert back to original dtype
        return new_mus.float(), new_covs.float(), new_log_pi.float(), log_lik.float()