import torch
import torch.nn as nn
import torch.nn.functional as F
import math

# From https://github.com/karpathy/minGPT/blob/master/mingpt/model.py

class NewGELU(nn.Module):
    """
    Implementation of the GELU activation function currently in Google BERT repo (identical to OpenAI GPT).
    Reference: Gaussian Error Linear Units (GELU) paper: https://arxiv.org/abs/1606.08415
    """
    def __init__(self):
        super(NewGELU, self).__init__()
    
    def forward(self, x):
        return 0.5 * x * (1.0 + torch.tanh(math.sqrt(2.0 / math.pi) * (x + 0.044715 * torch.pow(x, 3.0))))

class CausalSelfAttention(nn.Module):
    """
    A vanilla multi-head masked self-attention layer with a projection at the end.
    It is possible to use torch.nn.MultiheadAttention here but I am including an
    explicit implementation here to show that there is nothing too scary here.
    """

    def __init__(self, n_embd, n_head, attn_pdrop, resid_pdrop, max_seqlen):
        super().__init__()
        assert n_embd % n_head == 0
        # key, query, value projections for all heads, but in a batch
        self.c_attn = nn.Linear(n_embd, 3 * n_embd)
        # output projection
        self.c_proj = nn.Linear(n_embd, n_embd)
        # regularization
        self.attn_dropout = nn.Dropout(attn_pdrop)
        self.resid_dropout = nn.Dropout(resid_pdrop)
        # causal mask to ensure that attention is only applied to the left in the input sequence
        self.register_buffer("bias", torch.tril(torch.ones(max_seqlen, max_seqlen))
                                     .view(1, 1, max_seqlen, max_seqlen))
        self.n_head = n_head
        self.n_embd = n_embd

    def forward(self, x):
        B, T, C = x.size() # batch size, sequence length, embedding dimensionality (n_embd)

        # calculate query, key, values for all heads in batch and move head forward to be the batch dim
        q, k ,v  = self.c_attn(x).split(self.n_embd, dim=2)
        k = k.view(B, T, self.n_head, C // self.n_head).transpose(1, 2) # (B, nh, T, hs)
        q = q.view(B, T, self.n_head, C // self.n_head).transpose(1, 2) # (B, nh, T, hs)
        v = v.view(B, T, self.n_head, C // self.n_head).transpose(1, 2) # (B, nh, T, hs)

        # causal self-attention; Self-attend: (B, nh, T, hs) x (B, nh, hs, T) -> (B, nh, T, T)
        att = (q @ k.transpose(-2, -1)) * (1.0 / math.sqrt(k.size(-1)))
        att = att.masked_fill(self.bias[:,:,:T,:T] == 0, float('-inf'))
        att = F.softmax(att, dim=-1)
        att = self.attn_dropout(att)
        y = att @ v # (B, nh, T, T) x (B, nh, T, hs) -> (B, nh, T, hs)
        y = y.transpose(1, 2).contiguous().view(B, T, C) # re-assemble all head outputs side by side

        # output projection
        y = self.resid_dropout(self.c_proj(y))
        return y
    
class Model(nn.Module):
    """ an unassuming Transformer block """

    def __init__(self, n_embd, n_head, attn_pdrop, resid_pdrop, max_seqlen):
        super().__init__()
        self.ln_1 = nn.LayerNorm(n_embd)
        self.attn = CausalSelfAttention(n_embd, n_head, attn_pdrop, resid_pdrop, max_seqlen)
        self.ln_2 = nn.LayerNorm(n_embd)
        self.mlp = nn.ModuleDict(dict(
            c_fc    = nn.Linear(n_embd, 4 * n_embd),
            c_proj  = nn.Linear(4 * n_embd, n_embd),
            act     = NewGELU(),
            dropout = nn.Dropout(resid_pdrop),
        ))
        m = self.mlp
        self.mlpf = lambda x: m.dropout(m.c_proj(m.act(m.c_fc(x)))) # MLP forward

    def forward(self, x):
        x = x + self.attn(self.ln_1(x))
        x = x + self.mlpf(self.ln_2(x))
        return x

batch_size = 128
max_seqlen = 1024
seq_len = 512
n_embd = 768
n_head = 8
attn_pdrop = 0.0
resid_pdrop = 0.0

def get_inputs():
    return [torch.rand(batch_size, seq_len, n_embd)]

def get_init_inputs():
    return [n_embd, n_head, attn_pdrop, resid_pdrop, max_seqlen]

import torch
import torch.nn as nn
import torch.nn.functional as F
from torch.utils.cpp_extension import load_inline

_cpp_src = r"""
#include <torch/extension.h>
#include <vector>
torch::Tensor gelu_tanh_cuda(torch::Tensor x);
torch::Tensor layernorm_cuda(torch::Tensor x, torch::Tensor gamma, torch::Tensor beta, double eps);
std::vector<torch::Tensor> add_layernorm_cuda(torch::Tensor x, torch::Tensor y, torch::Tensor gamma, torch::Tensor beta, double eps);
"""

_cuda_src = r"""
#include <torch/extension.h>
#include <ATen/cuda/CUDAContext.h>
#include <cuda.h>
#include <cuda_runtime.h>
#include <vector>

__global__ void gelu_tanh_kernel(const float* __restrict__ x, float* __restrict__ y, long long n) {
    long long idx = (long long)blockIdx.x * blockDim.x + threadIdx.x;
    long long stride = (long long)blockDim.x * gridDim.x;
    const float kAlpha = 0.7978845608028654f;
    const float kBeta = 0.044715f;
    for (long long i = idx; i < n; i += stride) {
        float v = x[i];
        float u = kAlpha * (v + kBeta * v * v * v);
        y[i] = 0.5f * v * (1.0f + __tanhf(u));
    }
}

torch::Tensor gelu_tanh_cuda(torch::Tensor x) {
    auto y = torch::empty_like(x);
    long long n = x.numel();
    if (n == 0) return y;
    int threads = 256;
    int blocks = (int)((n + threads - 1) / threads);
    if (blocks > 65535) blocks = 65535;
    gelu_tanh_kernel<<<blocks, threads, 0, at::cuda::getCurrentCUDAStream()>>>(x.data_ptr<float>(), y.data_ptr<float>(), n);
    return y;
}

__device__ __forceinline__ void warp_reduce_pair(float &a, float &b) {
    for (int offset = 16; offset > 0; offset >>= 1) {
        a += __shfl_down_sync(0xffffffff, a, offset);
        b += __shfl_down_sync(0xffffffff, b, offset);
    }
}

__global__ void layernorm_kernel(const float* __restrict__ x, const float* __restrict__ gamma, const float* __restrict__ beta,
                                 float* __restrict__ out, int rows, int C, float eps) {
    int row = blockIdx.x;
    int tid = threadIdx.x;
    float sum = 0.0f, ss = 0.0f;
    const float* xr = x + (long long)row * C;
    for (int c = tid; c < C; c += blockDim.x) { float v = xr[c]; sum += v; ss += v * v; }
    __shared__ float sh_sum[32];
    __shared__ float sh_ss[32];
    __shared__ float sh_mean;
    __shared__ float sh_inv;
    warp_reduce_pair(sum, ss);
    if ((tid & 31) == 0) { sh_sum[tid >> 5] = sum; sh_ss[tid >> 5] = ss; }
    __syncthreads();
    sum = (tid < (blockDim.x >> 5)) ? sh_sum[tid] : 0.0f;
    ss = (tid < (blockDim.x >> 5)) ? sh_ss[tid] : 0.0f;
    if (tid < 32) warp_reduce_pair(sum, ss);
    if (tid == 0) {
        float mean = sum / (float)C;
        float var = ss / (float)C - mean * mean;
        sh_mean = mean;
        sh_inv = rsqrtf(var + eps);
    }
    __syncthreads();
    float mean = sh_mean, inv = sh_inv;
    float* orow = out + (long long)row * C;
    for (int c = tid; c < C; c += blockDim.x) { orow[c] = (xr[c] - mean) * inv * gamma[c] + beta[c]; }
}

__global__ void add_layernorm_kernel(const float* __restrict__ x, const float* __restrict__ y,
                                     const float* __restrict__ gamma, const float* __restrict__ beta,
                                     float* __restrict__ res, float* __restrict__ out, int rows, int C, float eps) {
    int row = blockIdx.x;
    int tid = threadIdx.x;
    const long long base = (long long)row * C;
    float sum = 0.0f, ss = 0.0f;
    for (int c = tid; c < C; c += blockDim.x) { float v = x[base+c] + y[base+c]; res[base+c] = v; sum += v; ss += v*v; }
    __shared__ float sh_sum[32];
    __shared__ float sh_ss[32];
    __shared__ float sh_mean;
    __shared__ float sh_inv;
    warp_reduce_pair(sum, ss);
    if ((tid & 31) == 0) { sh_sum[tid >> 5] = sum; sh_ss[tid >> 5] = ss; }
    __syncthreads();
    sum = (tid < (blockDim.x >> 5)) ? sh_sum[tid] : 0.0f;
    ss = (tid < (blockDim.x >> 5)) ? sh_ss[tid] : 0.0f;
    if (tid < 32) warp_reduce_pair(sum, ss);
    if (tid == 0) {
        float mean = sum / (float)C;
        float var = ss / (float)C - mean * mean;
        sh_mean = mean;
        sh_inv = rsqrtf(var + eps);
    }
    __syncthreads();
    float mean = sh_mean, inv = sh_inv;
    for (int c = tid; c < C; c += blockDim.x) { float vv = res[base+c]; out[base+c] = (vv - mean) * inv * gamma[c] + beta[c]; }
}

torch::Tensor layernorm_cuda(torch::Tensor x, torch::Tensor gamma, torch::Tensor beta, double eps) {
    auto out = torch::empty_like(x);
    int C = (int)x.size(-1); int rows = (int)(x.numel() / C);
    layernorm_kernel<<<rows, 256, 0, at::cuda::getCurrentCUDAStream()>>>(x.data_ptr<float>(), gamma.data_ptr<float>(), beta.data_ptr<float>(), out.data_ptr<float>(), rows, C, (float)eps);
    return out;
}

std::vector<torch::Tensor> add_layernorm_cuda(torch::Tensor x, torch::Tensor y, torch::Tensor gamma, torch::Tensor beta, double eps) {
    auto res = torch::empty_like(x); auto out = torch::empty_like(x);
    int C = (int)x.size(-1); int rows = (int)(x.numel() / C);
    add_layernorm_kernel<<<rows, 256, 0, at::cuda::getCurrentCUDAStream()>>>(x.data_ptr<float>(), y.data_ptr<float>(), gamma.data_ptr<float>(), beta.data_ptr<float>(), res.data_ptr<float>(), out.data_ptr<float>(), rows, C, (float)eps);
    return {res, out};
}
"""

_ext = load_inline(name="min_gpt_block_ext_final", cpp_sources=_cpp_src, cuda_sources=_cuda_src, functions=["gelu_tanh_cuda", "layernorm_cuda", "add_layernorm_cuda"], extra_cuda_cflags=["-O3", "--use_fast_math"], verbose=False)

class CausalSelfAttentionNew(nn.Module):
    def __init__(self, n_embd, n_head, attn_pdrop, resid_pdrop, max_seqlen):
        super().__init__()
        assert n_embd % n_head == 0
        self.c_attn = nn.Linear(n_embd, 3 * n_embd)
        self.c_proj = nn.Linear(n_embd, n_embd)
        self.attn_dropout = nn.Dropout(attn_pdrop)
        self.resid_dropout = nn.Dropout(resid_pdrop)
        self.register_buffer("bias", torch.tril(torch.ones(max_seqlen, max_seqlen)).view(1, 1, max_seqlen, max_seqlen))
        self.n_head = n_head
        self.n_embd = n_embd

class NewGELU(nn.Module):
    def __init__(self):
        super().__init__()
    def forward(self, x):
        return _ext.gelu_tanh_cuda(x)

class ModelNew(nn.Module):
    def __init__(self, n_embd, n_head, attn_pdrop, resid_pdrop, max_seqlen):
        super().__init__()
        self.ln_1 = nn.LayerNorm(n_embd)
        self.attn = CausalSelfAttentionNew(n_embd, n_head, attn_pdrop, resid_pdrop, max_seqlen)
        self.ln_2 = nn.LayerNorm(n_embd)
        self.mlp = nn.ModuleDict(dict(
            c_fc=nn.Linear(n_embd, 4*n_embd),
            c_proj=nn.Linear(4*n_embd, n_embd),
            act=NewGELU(),
            dropout=nn.Dropout(resid_pdrop),
        ))

    def forward(self, x):
        B, T, C = x.shape
        h = _ext.layernorm_cuda(x, self.ln_1.weight, self.ln_1.bias, self.ln_1.eps)
        qkv = self.attn.c_attn(h)
        q, k, v = qkv.split(C, dim=2)
        nh = self.attn.n_head
        hs = C // nh
        q = q.view(B, T, nh, hs).transpose(1, 2)
        k = k.view(B, T, nh, hs).transpose(1, 2)
        v = v.view(B, T, nh, hs).transpose(1, 2)
        y = F.scaled_dot_product_attention(q, k, v, attn_mask=None, dropout_p=0.0, is_causal=True)
        y = y.transpose(1, 2).contiguous().view(B, T, C)
        y = self.attn.c_proj(y)
        pair = _ext.add_layernorm_cuda(x, y, self.ln_2.weight, self.ln_2.bias, self.ln_2.eps)
        x2 = pair[0]
        m = pair[1]
        m = self.mlp.c_fc(m)
        m = _ext.gelu_tanh_cuda(m)
        m = self.mlp.c_proj(m)
        return x2 + m