class MultiHeadCausalAttention(nn.Module):
    def __init__(self, config):
        super(MultiHeadCausalAttention, self).__init__()
        self.n_heads = config.n_head
        self.n_embd = config.n_embd
        self.n_positions = config.n_positions
        assert self.n_embd % self.n_heads == 0
        self.head_dim = self.n_embd // self.n_heads

        self.c_attn = nn.Linear(self.n_embd, self.n_embd * 3)
        self.c_proj = nn.Linear(self.n_embd, self.n_embd)

        self.attn_dropout = nn.Dropout(config.attn_pdrop)
        self.resid_dropout = nn.Dropout(config.resid_pdrop)
        self.is_causal = True

    def forward(self, hidden_states, use_cache=True, kv_past=None):
        B, L, D = hidden_states.shape
        qkv = self.c_attn(hidden_states).view(B, L, 3, self.n_heads, self.head_dim)
        queries, keys, values = torch.unbind(qkv, dim=2)  # 3*(B,L,h,D/h)
        queries = queries.permute(0, 2, 1, 3)  # (B,h,L,D/h)
        keys = keys.permute(0, 2, 1, 3)  # (B,h,L,D/h)
        values = values.permute(0, 2, 1, 3)  # (B,h,L,D/h)
        ##################################################################
        if kv_past is not None:
            k_past, v_past = kv_past
            keys = torch.cat([k_past, keys], dim=-2)
            values = torch.cat([v_past, values], dim=-2)

        if use_cache:
            kv_present = keys, values
        else:
            kv_present = None
        ##################################################################
        attn_weights = (
            torch.matmul(queries, keys.transpose(-1, -2)) / values.size(-1) ** 0.5
        )
        # create the mask
        q_len, k_len = queries.size(-2), keys.size(-2)
        causal_mask = torch.tril(
            torch.ones((self.n_positions, self.n_positions), dtype=torch.bool)
        ).view(1, 1, self.n_positions, self.n_positions)
        causal_mask = causal_mask[
            :, :, k_len - q_len : k_len, :k_len
        ]  # shape=(1,1,q_len, kv_len)
        mask_value = torch.finfo(attn_weights.dtype).min
        mask_value = torch.tensor(mask_value, dtype=attn_weights.dtype).to(
            attn_weights.device
        )
        attn_weights = torch.where(causal_mask, attn_weights, mask_value)  # (B,h,L,L)
        attn_weights = F.softmax(attn_weights, dim=-1)
        attn_weights = self.attn_dropout(attn_weights)
        attn_out = torch.matmul(attn_weights, values)  # (B,h,L,D/h)
        attn_out = attn_out.permute(0, 2, 1, 3).reshape(B, L, D).contiguous()
        attn_out = self.c_proj(attn_out)
        attn_out = self.resid_dropout(attn_out)
        return attn_out, kv_present

import math
import json
from easydict import EasyDict as edict

import torch
import torch.nn as nn


def gelu_new(input):
    return (
        0.5
        * input
        * (
            1.0
            + torch.tanh(
                math.sqrt(2.0 / math.pi) * (input + 0.044715 * torch.pow(input, 3.0))
            )
        )
    )

class FeedForward(nn.Module):
    def __init__(self, config):
        super(FeedForward, self).__init__()
        n_embd = config.n_embd
        self.c_fc = nn.Linear(n_embd, 4 * n_embd)
        self.c_proj = nn.Linear(4 * n_embd, n_embd)
        self.act = gelu_new
        self.dropout = nn.Dropout(config.resid_pdrop)

    def forward(self, hidden_states) -> torch.FloatTensor:
        hidden_states = self.c_fc(hidden_states)
        hidden_states = self.act(hidden_states)
        hidden_states = self.c_proj(hidden_states)
        hidden_states = self.dropout(hidden_states)
        return hidden_states

## add the same arguments
class DecoderBlock(nn.Module):
    def __init__(self, config):
        super(DecoderBlock, self).__init__()
        self.ln_1 = nn.LayerNorm(config.n_embd, eps=config.layer_norm_epsilon)
        self.attn = MultiHeadCausalAttention(config)
        self.ln_2 = nn.LayerNorm(config.n_embd, eps=config.layer_norm_epsilon)
        self.mlp = FeedForward(config)

    def forward(self, hidden_states, use_cache=True, kv_past=None):
        residual = hidden_states
        hidden_states = self.ln_1(hidden_states)
        attn_out, kv_present = self.attn(
            hidden_states, use_cache=use_cache, kv_past=kv_past
        )
        hidden_states = attn_out + residual

        residual = hidden_states
        hidden_states = self.ln_2(hidden_states)
        ff_out = self.mlp(hidden_states)
        hidden_states = residual + ff_out
        return hidden_states, kv_present


class GPT2Model(nn.Module):
    def __init__(self, config):
        super(GPT2Model, self).__init__()
        self.n_layer = config.n_layer
        self.n_embd = config.n_embd

        self.wte = nn.Embedding(config.vocab_size, self.n_embd)
        self.wpe = nn.Embedding(config.n_positions, self.n_embd)
        self.drop = nn.Dropout(config.embd_pdrop)

        self.h = nn.ModuleList([DecoderBlock(config) for _ in range(config.n_layer)])
        self.ln_f = nn.LayerNorm(self.n_embd, eps=config.layer_norm_epsilon)
        self.lm_head = nn.Linear(config.n_embd, config.vocab_size, bias=False)

    def forward(self, inputs_ids, position_ids, use_cache=True, kv_cache=None):
        input_shape = inputs_ids.size()
        ##################################################################
        if kv_cache is None:
            past_length = 0
            kv_cache = tuple([None] * len(self.h))
        else:
            past_length = kv_cache[0][0].size(-2)
        ##################################################################
        if position_ids is None:
            position_ids = torch.arange(
                past_length,
                input_shape[-1] + past_length,
                dtype=torch.long,
                device=inputs_ids.device,
            )
            position_ids = position_ids.unsqueeze(0)

        inputs_embeds = self.wte(inputs_ids)
        position_embeds = self.wpe(position_ids)
        hidden_states = inputs_embeds + position_embeds
        hidden_states = self.drop(hidden_states)
        kv_present = [] if use_cache else None

        for i in range(len(self.h)):
            block, kv_past = self.h[i], kv_cache[i]
            hidden_states, present_i = block(
                hidden_states, use_cache=use_cache, kv_past=kv_past
            )
            ##################################################################
            if use_cache:
                kv_present.append(present_i)
            ##################################################################

        hidden_states = self.ln_f(hidden_states)
        logits = self.lm_head(hidden_states)
        return logits, kv_present

def model_inference(
    model, input_ids, max_length=1000, use_cache=True, top_k=50, top_p=0.95
):
    time_stamps = []

    output_ids = input_ids
    kv_cache = None if use_cache else []
    print(tokenizer.decode(input_ids[0]), end="", flush=True)

    model.eval()
    with torch.no_grad():
        for step in range(max_length):
            t_s = time.time()
            if use_cache:
                # Use only the last token for efficient generation
                position_ids = torch.tensor([[output_ids.size(-1)]])
                logits, kv_cache = model(
                    output_ids[:, -1:],
                    position_ids,
                    use_cache=True,
                    kv_cache=kv_cache,
                )
            else:
                # Use all tokens without KV cache
                position_ids = torch.arange(output_ids.size(-1)).unsqueeze(0)
                logits, _ = model(output_ids, position_ids, use_cache=False)
            time_stamps.append(time.time() - t_s)

            logits = logits[:, -1, :]

            # Top-k and top-p sampling
            logits = logits / 0.7  # Temperature scaling (1.0 = no scaling)
            logits = logits + 1e-8  # Avoid NaNs by adding small epsilon
            sorted_logits, sorted_indices = torch.sort(logits, descending=True)
            cumulative_probs = torch.cumsum(F.softmax(sorted_logits, dim=-1), dim=-1)
            sorted_indices_to_remove = cumulative_probs > top_p

            if top_k > 0:
                indices_to_remove = sorted_logits < sorted_logits[:, top_k - 1 : top_k]
                sorted_indices_to_remove = sorted_indices_to_remove | indices_to_remove

            sorted_logits[sorted_indices_to_remove] = -float("Inf")
            probs = F.softmax(sorted_logits, dim=-1)
            probs[torch.isnan(probs)] = 0  # Replace NaNs with zeros
            probs[probs < 0] = 0  # Replace negative values with zeros

            if torch.sum(probs) == 0:
                # print("Warning: Sum of probabilities is zero. Falling back to uniform sampling.")
                next_token_id = torch.randint(
                    0, logits.size(-1), (1, 1)
                )  # Fallback random token
            else:
                probs = probs / torch.sum(probs)  # Re-normalize the distribution
                next_token_id = torch.multinomial(probs, num_samples=1)
                next_token_id = torch.gather(sorted_indices, -1, next_token_id)

            output_ids = torch.cat([output_ids, next_token_id], dim=-1)

            # Print generated token incrementally
            print(tokenizer.decode(next_token_id[0]), end="", flush=True)

            # Stop if EOS token is generated
            if next_token_id.item() == tokenizer.eos_token_id:
                break

    return time_stamps

if __name__ == "__main__":
    model_name = "gpt2"
    with open("{}-config.json".format(model_name), "r") as f:
        config = json.load(f)
        config = edict(config)

    model = GPT2Model(config=config)
    w_dict = {}
    # model weights availble at: https://huggingface.co/openai-community/gpt2/resolve/main/model.safetensors
    with safetensors.safe_open("{}.safetensors".format(model_name), framework="pt") as f:
        for k in f.keys():
            w_dict[k] = f.get_tensor(k)
            # GPT-2 use Conv1D, who works same as nn.Linear but with transposed weights
            if k.startswith("h."):
                w_dict[k] = w_dict[k].T
        w_dict["lm_head.weight"] = w_dict["wte.weight"]
    # gpt2.load_state_dict(w_dict)
    model.load_state_dict(w_dict, strict=False)

    tokenizer = AutoTokenizer.from_pretrained(model_name)

    # Input prompt
    prompt = "Once upon a time in a distant kingdom, there was a wise old king who"
    inputs = tokenizer(prompt, return_tensors="pt")
    input_ids = inputs["input_ids"]

    # Generate a long sequence with kv-cache
    max_length = 1000
    output_ids = input_ids
    kv_cache = None

    time_stamps = model_inference(
        model, input_ids, max_length=500, use_cache=False, top_k=30, top_p=0.9
    )
    print("\naverage time: {:.2f}s/token".format(np.mean(time_stamps)))
    # average time: 0.13s/token

        time_stamps = model_inference(
        model, input_ids, max_length=500, use_cache=True, top_k=30, top_p=0.9
    )
    print("\naverage time: {:.2f}s/token".format(np.mean(time_stamps)))
    # average time: 0.03s/token