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

# Compile CUDA kernel at module load time
cuda_source = """
#include <torch/extension.h>
#include <cuda.h>
#include <cuda_runtime.h>
#include <cmath>

__global__ void compute_log_resp_kernel(
    const float* __restrict__ x,
    const float* __restrict__ mus,
    const float* __restrict__ L,
    const float* __restrict__ log_det,
    const float* __restrict__ log_pi,
    float* __restrict__ log_resp,
    int N, int K, int D
) {
    int n = blockIdx.x * blockDim.x + threadIdx.x;
    int k = blockIdx.y * blockDim.y + threadIdx.y;
    
    if (n >= N || k >= K) return;
    
    const float log_2pi = 1.8378770664093453f;  // log(2*pi)
    
    int idx = n * K + k;
    int x_offset = n * D;
    int mu_offset = k * D;
    int L_offset = k * D * D;
    
    // Allocate diff and y in registers/local memory
    extern __shared__ float shared_mem[];
    float* diff = shared_mem;  // Not used in this simple version
    
    // Compute diff and triangular solve on-the-fly
    // Solve L @ y = diff, where L is lower triangular
    // Then compute y^T @ y = ||y||^2
    
    float y_sq_sum = 0.0f;
    
    // Forward substitution: solve row by row
    for (int i = 0; i < D; i++) {
        // diff[i] = x[n, i] - mus[k, i]
        float diff_i = x[x_offset + i] - mus[mu_offset + i];
        
        // Subtract L[k, i, j] * y[j] for j < i
        float sum = diff_i;
        for (int j = 0; j < i; j++) {
            // L[k, i, j] is at L_offset + i * D + j
            sum -= L[L_offset + i * D + j] * y_sq_sum;  // This is wrong, need separate y array
        }
        
        // Actually, let me redo this correctly with a local y array
        // The above is incorrect
    }
}

__global__ void log_gaussian_kernel(
    const float* __restrict__ x,
    const float* __restrict__ mus,
    const float* __restrict__ L,
    const float* __restrict__ log_det,
    const float* __restrict__ log_pi,
    float* __restrict__ log_resp,
    int N, int K, int D
) {
    // Each thread handles one (n, k) pair
    int n = blockIdx.x * blockDim.x + threadIdx.x;
    int k = blockIdx.y * blockDim.y + threadIdx.y;
    
    if (n >= N || k >= K) return;
    
    const float log_2pi = 1.8378770664093453f;
    
    int idx = n * K + k;
    
    // Compute diff = x[n] - mus[k] and solve L[k] @ y = diff
    // Using forward substitution (L is lower triangular)
    
    // We'll compute maha = ||y||^2 directly
    
    float maha = 0.0f;
    
    for (int i = 0; i < D; i++) {
        // diff[i] = x[n, i] - mus[k, i]
        float diff_i = x[n * D + i] - mus[k * D + i];
        
        // y[i] = (diff[i] - sum_{j<i} L[k,i,j] * y[j]) / L[k,i,i]
        // We need to track y values...
        
        // Let's use a different approach: compute y incrementally
    }
    
    // Actually, let me write this more carefully with a local y buffer
    // For small D (like 8), this fits in registers
}

__global__ void compute_mahalanobis_kernel(
    const float* __restrict__ x,
    const float* __restrict__ mus,
    const float* __restrict__ L,
    const float* __restrict__ log_det,
    const float* __restrict__ log_pi,
    float* __restrict__ log_resp,
    int N, int K, int D
) {
    // Each thread handles one (n, k) pair
    // D is expected to be small (e.g., 8)
    int n = blockIdx.x * blockDim.x + threadIdx.x;
    int k = blockIdx.y * blockDim.y + threadIdx.y;
    
    if (n >= N || k >= K) return;
    
    const float log_2pi = 1.8378770664093453f;  // log(2*pi)
    
    // Pointers to the relevant data
    const float* x_ptr = x + n * D;
    const float* mu_ptr = mus + k * D;
    const float* L_ptr = L + k * D * D;
    
    // Compute diff = x - mu and solve L @ y = diff using forward substitution
    // L is lower triangular, stored row-major
    // We compute ||y||^2 directly
    
    float y[32];  // Assume D <= 32
    float maha = 0.0f;
    
    for (int i = 0; i < D; i++) {
        // diff[i] = x[i] - mu[i]
        float diff_i = x_ptr[i] - mu_ptr[i];
        
        // y[i] = (diff_i - sum_{j < i} L[i,j] * y[j]) / L[i,i]
        float sum = diff_i;
        for (int j = 0; j < i; j++) {
            sum -= L_ptr[i * D + j] * y[j];
        }
        y[i] = sum / L_ptr[i * D + i];
        maha += y[i] * y[i];
    }
    
    // log_resp[n,k] = -0.5 * (D * log(2*pi) + log_det[k] + maha) + log_pi[k]
    float log_gauss = -0.5f * (D * log_2pi + log_det[k] + maha);
    log_resp[n * K + k] = log_gauss + log_pi[k];
}

// Kernel for M-step covariance computation
// Each thread computes one element of new_covs[k]
__global__ void compute_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 K, int D
) {
    int k = blockIdx.x * blockDim.x + threadIdx.x;
    int i = blockIdx.y * blockDim.y + threadIdx.y;
    int j = blockIdx.z * blockDim.z + threadIdx.z;
    
    if (k >= K || i >= D || j >= D) return;
    
    // new_covs[k, i, j] = sum_n resp[n,k] * (x[n,i] - new_mus[k,i]) * (x[n,j] - new_mus[k,j]) / N_k[k]
    // Only compute lower triangular part (cov is symmetric)
    
    float sum = 0.0f;
    float mu_i = new_mus[k * D + i];
    float mu_j = new_mus[k * D + j];
    
    for (int n = 0; n < N; n++) {
        float diff_i = x[n * D + i] - mu_i;
        float diff_j = x[n * D + j] - mu_j;
        sum += resp[n * K + k] * diff_i * diff_j;
    }
    
    new_covs[k * D * D + i * D + j] = sum / N_k[k];
}

// Optimized covariance kernel using shared memory for partial sums
__global__ void compute_covs_fast_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 K, int D
) {
    // One block per (k, i, j) output element
    // Threads in block cooperate to sum over N
    extern __shared__ float sdata[];
    
    int k = blockIdx.x / (D * D);
    int ij = blockIdx.x % (D * D);
    int i = ij / D;
    int j = ij % D;
    
    int tid = threadIdx.x;
    int blockSize = blockDim.x;
    
    // Initialize shared memory for reduction
    sdata[tid] = 0.0f;
    
    float mu_i = new_mus[k * D + i];
    float mu_j = new_mus[k * D + j];
    
    // Each thread sums over a subset of N
    for (int n = tid; n < N; n += blockSize) {
        float diff_i = x[n * D + i] - mu_i;
        float diff_j = x[n * D + j] - mu_j;
        sdata[tid] += resp[n * K + k] * diff_i * diff_j;
    }
    __syncthreads();
    
    // Block reduction
    for (int s = blockSize / 2; s > 0; s >>= 1) {
        if (tid < s) {
            sdata[tid] += sdata[tid + s];
        }
        __syncthreads();
    }
    
    // Write result
    if (tid == 0) {
        new_covs[k * D * D + i * D + j] = sdata[0] / N_k[k];
    }
}

// Log-sum-exp kernel: compute logsumexp along K dimension
__global__ void logsumexp_kernel(
    const float* __restrict__ log_resp,
    float* __restrict__ result,
    float* __restrict__ normalized,  // optional: if not null, store normalized values
    int N, int K
) {
    int n = blockIdx.x * blockDim.x + threadIdx.x;
    if (n >= N) return;
    
    // Find max for numerical stability
    float max_val = log_resp[n * K];
    for (int k = 1; k < K; k++) {
        float val = log_resp[n * K + k];
        if (val > max_val) max_val = val;
    }
    
    // Compute sum of exp
    float sum = 0.0f;
    for (int k = 0; k < K; k++) {
        sum += expf(log_resp[n * K + k] - max_val);
    }
    
    // logsumexp = max + log(sum)
    float lse = max_val + logf(sum);
    result[n] = lse;
    
    // Optionally normalize
    if (normalized != nullptr) {
        for (int k = 0; k < K; k++) {
            normalized[n * K + k] = log_resp[n * K + k] - lse;
        }
    }
}
"""

cpp_source = """
std::tuple<torch::Tensor, torch::Tensor, torch::Tensor> gmm_e_step_cuda(
    torch::Tensor x,
    torch::Tensor mus,
    torch::Tensor L,  // Cholesky factors
    torch::Tensor log_det,
    torch::Tensor log_pi
);

torch::Tensor gmm_compute_covs_cuda(
    torch::Tensor x,
    torch::Tensor new_mus,
    torch::Tensor resp,
    torch::Tensor N_k
);
"""

# Actually let me rewrite this more cleanly
cuda_source = """
#include <torch/extension.h>
#include <cuda.h>
#include <cuda_runtime.h>
#include <cmath>

// Compute log_resp for all (n, k) pairs
__global__ void log_gaussian_batched_kernel(
    const float* __restrict__ x,
    const float* __restrict__ mus,
    const float* __restrict__ L,
    const float* __restrict__ log_det,
    const float* __restrict__ log_pi,
    float* __restrict__ log_resp,
    int N, int K, int D
) {
    int n = blockIdx.x * blockDim.x + threadIdx.x;
    int k = blockIdx.y * blockDim.y + threadIdx.y;
    
    if (n >= N || k >= K) return;
    
    const float log_2pi = 1.8378770664093453f;
    
    const float* x_ptr = x + n * D;
    const float* mu_ptr = mus + k * D;
    const float* L_ptr = L + k * D * D;
    
    // Forward substitution to solve L @ y = (x - mu)
    // and compute ||y||^2
    float y_sq_sum = 0.0f;
    
    // For small D, use register array
    float y[16];  // Assume D <= 16
    
    for (int i = 0; i < D; i++) {
        float diff_i = x_ptr[i] - mu_ptr[i];
        float sum = diff_i;
        for (int j = 0; j < i; j++) {
            sum -= L_ptr[i * D + j] * y[j];
        }
        y[i] = sum / L_ptr[i * D + i];
        y_sq_sum += y[i] * y[i];
    }
    
    float log_gauss = -0.5f * (D * log_2pi + log_det[k] + y_sq_sum);
    log_resp[n * K + k] = log_gauss + log_pi[k];
}

// Logsumexp and normalization kernel
__global__ void logsumexp_norm_kernel(
    const float* __restrict__ log_resp,
    float* __restrict__ lse_out,
    float* __restrict__ log_resp_norm,
    int N, int K
) {
    int n = blockIdx.x * blockDim.x + threadIdx.x;
    if (n >= N) return;
    
    const float* row = log_resp + n * K;
    float* out_row = log_resp_norm + n * K;
    
    // Find max
    float max_val = row[0];
    for (int k = 1; k < K; k++) {
        max_val = fmaxf(max_val, row[k]);
    }
    
    // Sum of exp(x - max)
    float sum = 0.0f;
    for (int k = 0; k < K; k++) {
        sum += expf(row[k] - max_val);
    }
    
    float lse = max_val + logf(sum);
    lse_out[n] = lse;
    
    // Normalize
    for (int k = 0; k < K; k++) {
        out_row[k] = row[k] - lse;
    }
}

// Compute N_k and new_mus
__global__ void compute_stats_kernel(
    const float* __restrict__ resp,
    const float* __restrict__ x,
    float* __restrict__ N_k,
    float* __restrict__ new_mus,
    int N, int K, int D
) {
    int k = blockIdx.x * blockDim.x + threadIdx.x;
    if (k >= K) return;
    
    float n_k = 0.0f;
    float mu[16];  // Assume D <= 16
    for (int d = 0; d < D; d++) mu[d] = 0.0f;
    
    for (int n = 0; n < N; n++) {
        float r = resp[n * K + k];
        n_k += r;
        for (int d = 0; d < D; d++) {
            mu[d] += r * x[n * D + d];
        }
    }
    
    N_k[k] = fmaxf(n_k, 1e-8f);
    for (int d = 0; d < D; d++) {
        new_mus[k * D + d] = mu[d] / N_k[k];
    }
}

// Compute weighted covariance contributions
// Each thread handles one (k, i, j) element
__global__ void compute_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 K, int D
) {
    int k = blockIdx.x;
    int i = threadIdx.x;
    int j = threadIdx.y;
    
    if (k >= K || i >= D || j >= D) return;
    
    float sum = 0.0f;
    float mu_i = new_mus[k * D + i];
    float mu_j = new_mus[k * D + j];
    
    for (int n = 0; n < N; n++) {
        float r = resp[n * K + k];
        float diff_i = x[n * D + i] - mu_i;
        float diff_j = x[n * D + j] - mu_j;
        sum += r * diff_i * diff_j;
    }
    
    new_covs[k * D * D + i * D + j] = sum / N_k[k];
}

// Host functions
torch::Tensor compute_log_resp_cuda(
    torch::Tensor x,
    torch::Tensor mus,
    torch::Tensor L,
    torch::Tensor log_det,
    torch::Tensor log_pi
) {
    int N = x.size(0);
    int K = mus.size(0);
    int D = x.size(1);
    
    auto log_resp = torch::empty({N, K}, x.options());
    
    dim3 block(32, 4);  // 128 threads per block
    dim3 grid((N + block.x - 1) / block.x, (K + block.y - 1) / block.y);
    
    log_gaussian_batched_kernel<<<grid, block>>>(
        x.data_ptr<float>(),
        mus.data_ptr<float>(),
        L.data_ptr<float>(),
        log_det.data_ptr<float>(),
        log_pi.data_ptr<float>(),
        log_resp.data_ptr<float>(),
        N, K, D
    );
    
    return log_resp;
}

std::tuple<torch::Tensor, torch::Tensor, torch::Tensor> e_step_cuda(
    torch::Tensor x,
    torch::Tensor mus,
    torch::Tensor L,
    torch::Tensor log_det,
    torch::Tensor log_pi
) {
    int N = x.size(0);
    int K = mus.size(0);
    int D = x.size(1);
    
    auto log_resp = torch::empty({N, K}, x.options());
    auto lse = torch::empty({N}, x.options());
    auto log_resp_norm = torch::empty({N, K}, x.options());
    
    dim3 block1(32, 4);
    dim3 grid1((N + block1.x - 1) / block1.x, (K + block1.y - 1) / block1.y);
    
    log_gaussian_batched_kernel<<<grid1, block1>>>(
        x.data_ptr<float>(),
        mus.data_ptr<float>(),
        L.data_ptr<float>(),
        log_det.data_ptr<float>(),
        log_pi.data_ptr<float>(),
        log_resp.data_ptr<float>(),
        N, K, D
    );
    
    // Logsumexp and normalization
    int block2 = 256;
    int grid2 = (N + block2 - 1) / block2;
    
    logsumexp_norm_kernel<<<grid2, block2>>>(
        log_resp.data_ptr<float>(),
        lse.data_ptr<float>(),
        log_resp_norm.data_ptr<float>(),
        N, K
    );
    
    return std::make_tuple(log_resp_norm, lse, log_resp);
}

std::tuple<torch::Tensor, torch::Tensor> m_step_stats_cuda(
    torch::Tensor resp,
    torch::Tensor x
) {
    int N = x.size(0);
    int K = resp.size(1);
    int D = x.size(1);
    
    auto N_k = torch::empty({K}, x.options());
    auto new_mus = torch::empty({K, D}, x.options());
    
    int block = 256;
    int grid = (K + block - 1) / block;
    
    compute_stats_kernel<<<grid, block>>>(
        resp.data_ptr<float>(),
        x.data_ptr<float>(),
        N_k.data_ptr<float>(),
        new_mus.data_ptr<float>(),
        N, K, D
    );
    
    return std::make_tuple(N_k, new_mus);
}

torch::Tensor compute_covs_cuda(
    torch::Tensor x,
    torch::Tensor new_mus,
    torch::Tensor resp,
    torch::Tensor N_k
) {
    int N = x.size(0);
    int K = new_mus.size(0);
    int D = x.size(1);
    
    auto new_covs = torch::zeros({K, D, D}, x.options());
    
    // One block per k
    // 2D thread block for (i, j)
    dim3 block(D, D);
    dim3 grid(K);
    
    compute_covs_kernel<<<grid, block>>>(
        x.data_ptr<float>(),
        new_mus.data_ptr<float>(),
        resp.data_ptr<float>(),
        N_k.data_ptr<float>(),
        new_covs.data_ptr<float>(),
        N, K, D
    );
    
    return new_covs;
}
"""

gmm_module = load_inline(
    name='gmm_cuda',
    cpp_sources='',
    cuda_sources=cuda_source,
    functions=['e_step_cuda', 'm_step_stats_cuda', 'compute_covs_cuda'],
    verbose=False,
    extra_cuda_cflags=['-O3']
)


class ModelNew(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.
    Using custom CUDA kernels for performance.
    """

    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:
        """
        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
        
        # Normalize log_pi
        log_pi = self.log_pi - torch.logsumexp(self.log_pi, dim=0)
        
        # Compute Cholesky decomposition of covariances
        L = torch.linalg.cholesky(self.covs)  # (K, D, D)
        log_det = 2.0 * L.diagonal(dim1=-2, dim2=-1).log().sum(-1)  # (K,)
        
        # E-step using CUDA kernel
        log_resp_norm, lse, _ = gmm_module.e_step_cuda(
            x.contiguous(), 
            self.mus.contiguous(),
            L.contiguous(),
            log_det.contiguous(),
            log_pi.contiguous()
        )
        
        log_lik = lse.sum()
        resp = torch.exp(log_resp_norm)  # (N, K)
        
        # M-step using CUDA kernels
        N_k, new_mus = gmm_module.m_step_stats_cuda(
            resp.contiguous(),
            x.contiguous()
        )
        
        new_covs = gmm_module.compute_covs_cuda(
            x.contiguous(),
            new_mus.contiguous(),
            resp.contiguous(),
            N_k.contiguous()
        )
        
        # Add regularization to diagonal
        new_covs = new_covs + 1e-4 * torch.eye(D, device=x.device, dtype=x.dtype).unsqueeze(0)
        
        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():
    return [torch.randn(num_data, dim)]


def get_init_inputs():
    return [dim, num_components]