import os
import torch
import torch.nn as nn
import torch.nn.functional as F
from tokenizers import Tokenizer, pre_tokenizers, decoders

MODEL_PATH = "model.pt"
TOKENIZER_PATH = "tokenizer.json"

D_MODEL = 512
NUM_HEADS = 8
NUM_KV_HEADS = 8
NUM_LAYERS = 18
D_FF = 2048
MAX_SEQ_LENGTH = 768
VOCAB_SIZE_LIMIT = 32768
USE_LAYER_SCALE = False

class BPETokenizer:
    def __init__(self, vocab_size=15000):
        self.tokenizer = None
        self.vocab = {}
        self.token_to_id = {}
        self.id_to_token = {}

    def load(self, filepath):
        self.tokenizer = Tokenizer.from_file(filepath)
        self.tokenizer.pre_tokenizer = pre_tokenizers.ByteLevel(add_prefix_space=False)
        self.tokenizer.decoder = decoders.ByteLevel()
        
        self.vocab = self.tokenizer.get_vocab()
        self.token_to_id = self.vocab
        self.id_to_token = {v: k for k, v in self.vocab.items()}
        print(f"Tokenizer loaded. Vocabulary size: {self.tokenizer.get_vocab_size()}")

    def encode(self, text):
        return self.tokenizer.encode(text).ids

    def decode(self, ids):
        return self.tokenizer.decode(ids, skip_special_tokens=False)

class FastRMSNorm(nn.Module):
    def __init__(self, d_model, eps=1e-5):
        super().__init__()
        self.eps = eps
        self.weight = nn.Parameter(torch.ones(d_model))

    def forward(self, x):
        return x * torch.rsqrt(x.pow(2).mean(-1, keepdim=True) + self.eps) * self.weight

class RotaryEmbedding(nn.Module):
    def __init__(self, dim, max_seq_len=4096, theta=10000.0):
        super().__init__()
        self.dim = dim
        self.theta = theta
        self.max_seq_len = max_seq_len
        
        inv_freq = 1.0 / (theta ** (torch.arange(0, dim, 2).float() / dim))
        self.register_buffer("inv_freq", inv_freq)
        self._update_cache(max_seq_len)

    def _update_cache(self, seq_len):
        self.max_seq_len = seq_len
        t = torch.arange(self.max_seq_len, device=self.inv_freq.device, dtype=self.inv_freq.dtype)
        freqs = torch.outer(t, self.inv_freq)
        emb = torch.cat((freqs, freqs), dim=-1)
        
        self.register_buffer("cos_cached", emb.cos()[None, None, :, :], persistent=False)
        self.register_buffer("sin_cached", emb.sin()[None, None, :, :], persistent=False)

    def forward(self, x, seq_len):
        if seq_len > self.max_seq_len:
            self._update_cache(seq_len)
        return (
            self.cos_cached[:, :, :seq_len, ...].to(dtype=x.dtype),
            self.sin_cached[:, :, :seq_len, ...].to(dtype=x.dtype)
        )

def apply_rotary_pos_emb(x, cos, sin):
    def rotate_half(x):
        x1, x2 = x.chunk(2, dim=-1)
        return torch.cat((-x2, x1), dim=-1)
    return (x * cos) + (rotate_half(x) * sin)

def repeat_kv(x: torch.Tensor, n_rep: int) -> torch.Tensor:
    batch, n_kv_heads, seqlen, head_dim = x.shape
    if n_rep == 1:
        return x
    return (
        x[:, :, None, :, :]
        .expand(batch, n_kv_heads, n_rep, seqlen, head_dim)
        .reshape(batch, n_kv_heads * n_rep, seqlen, head_dim)
    )

def apply_swiglu(gate_up):
    a, b = gate_up.chunk(2, dim=-1)
    return F.silu(a) * b

class TokenEmbedding(nn.Embedding):
    def __init__(self, vocab_size, d_model, pad_id=None):
        padding_idx = pad_id if pad_id is not None and pad_id >= 0 else None
        super().__init__(num_embeddings=vocab_size, embedding_dim=d_model, padding_idx=padding_idx)

class GroupedQueryAttention(nn.Module):
    def __init__(self, d_model, num_heads, num_kv_heads):
        super().__init__()
        self.d_model = d_model
        self.num_heads = num_heads
        self.num_kv_heads = num_kv_heads
        self.head_dim = d_model // num_heads
        self.n_rep = self.num_heads // self.num_kv_heads

        self.W_q = nn.Linear(d_model, num_heads * self.head_dim, bias=False)
        self.W_k = nn.Linear(d_model, num_kv_heads * self.head_dim, bias=False)
        self.W_v = nn.Linear(d_model, num_kv_heads * self.head_dim, bias=False)
        self.W_o = nn.Linear(num_heads * self.head_dim, d_model, bias=False)
        
        self.rope = RotaryEmbedding(self.head_dim)
        
        self.q_norm = FastRMSNorm(self.head_dim)
        self.k_norm = FastRMSNorm(self.head_dim)

    def forward(self, q_in, k_in, v_in):
        batch_size, seq_len, _ = q_in.shape
        
        q = self.W_q(q_in)
        k = self.W_k(k_in)
        v = self.W_v(v_in)
        
        q = q.view(batch_size, seq_len, self.num_heads, self.head_dim)
        k = k.view(batch_size, seq_len, self.num_kv_heads, self.head_dim)
        v = v.view(batch_size, seq_len, self.num_kv_heads, self.head_dim)
        
        q = self.q_norm(q)
        k = self.k_norm(k)

        q = q.transpose(1, 2)
        k = k.transpose(1, 2)
        v = v.transpose(1, 2)
        
        cos_q, sin_q = self.rope(q, q.shape[2])
        cos_k, sin_k = self.rope(k, k.shape[2])
        
        q = apply_rotary_pos_emb(q, cos_q, sin_q)
        k = apply_rotary_pos_emb(k, cos_k, sin_k)

        k = repeat_kv(k, self.n_rep)
        v = repeat_kv(v, self.n_rep)

        context = F.scaled_dot_product_attention(
            q, k, v, attn_mask=None, dropout_p=0.0, is_causal=True 
        )
        context = context.transpose(1, 2).contiguous().view(batch_size, seq_len, self.d_model)
        return self.W_o(context)

class PositionWiseFeedForward(nn.Module):
    def __init__(self, d_model, d_ff, bias=False):
        super().__init__()
        self.gate_up_proj = nn.Linear(d_model, d_ff * 2, bias=bias)
        self.down_proj = nn.Linear(d_ff, d_model, bias=bias)

    def forward(self, x):
        gate_up = self.gate_up_proj(x)
        activated = apply_swiglu(gate_up)
        return self.down_proj(activated)

class DecoderLayer(nn.Module):
    def __init__(self, d_model, num_heads, num_kv_heads, d_ff, use_layer_scale=False):
        super().__init__()
        self.norm1 = FastRMSNorm(d_model)
        self.norm2 = FastRMSNorm(d_model)

        self.attn = GroupedQueryAttention(d_model, num_heads, num_kv_heads)
        self.ffn = PositionWiseFeedForward(d_model, d_ff)
        
        self.use_layer_scale = use_layer_scale
        if self.use_layer_scale:
            init_value = 1e-4
            self.ls_1 = nn.Parameter(torch.ones(d_model) * init_value)
            self.ls_2 = nn.Parameter(torch.ones(d_model) * init_value)

    def forward(self, x):
        x_norm = self.norm1(x)
        attn_out = self.attn(x_norm, x_norm, x_norm)
        x = x + (attn_out * self.ls_1 if self.use_layer_scale else attn_out)
        
        x_norm = self.norm2(x)
        ffn_out = self.ffn(x_norm)
        x = x + (ffn_out * self.ls_2 if self.use_layer_scale else ffn_out)
        
        return x

class Transformer(nn.Module):
    def __init__(self, vocab_size, d_model, num_heads, num_kv_heads, num_layers, d_ff, pad_id, use_layer_scale=False):
        super().__init__()
        self.pad_id = pad_id
        self.embedding = TokenEmbedding(vocab_size, d_model, pad_id)
        self.layers = nn.ModuleList([
            DecoderLayer(d_model, num_heads, num_kv_heads, d_ff, use_layer_scale) for _ in range(num_layers)
        ])
        self.norm_f = FastRMSNorm(d_model)
        self.fc_out = nn.Linear(d_model, vocab_size, bias=False)

    def forward(self, x):
        x = self.embedding(x)
        for layer in self.layers:
            x = layer(x)
        x = self.norm_f(x)
        return self.fc_out(x)

@torch.no_grad()
def generate_response(model, tokenizer, user_input, max_seq_length, device, temperature=0.7, top_k=50, repetition_penalty=1.2):
    model.eval()

    SOS_ID = tokenizer.token_to_id.get("<sos>", None)
    EOS_ID = tokenizer.token_to_id.get("<eos>", None)

    if not user_input.strip():
        return "Please say something."

    input_ids = tokenizer.encode(user_input)
    if SOS_ID is not None:
        input_ids = [SOS_ID] + input_ids

    generated_tokens_set = set()
    generated_new_tokens = [] 

    for _ in range(max_seq_length - len(input_ids)):
        x_tensor = torch.LongTensor([input_ids]).to(device)
        output = model(x_tensor)
        
        last_logits = output[0, -1, :] / temperature

        for token_id in generated_tokens_set:
            if last_logits[token_id] > 0:
                last_logits[token_id] /= repetition_penalty
            else:
                last_logits[token_id] *= repetition_penalty
        
        if top_k > 0:
            v, _ = torch.topk(last_logits, top_k)
            last_logits[last_logits < v[-1]] = -float('Inf')

        probs = torch.softmax(last_logits, dim=-1)
        next_word_id = torch.multinomial(probs, num_samples=1).item()
        
        if next_word_id == EOS_ID:
            break
            
        input_ids.append(next_word_id)
        generated_new_tokens.append(next_word_id)
        generated_tokens_set.add(next_word_id)
        
    raw_response = tokenizer.decode(generated_new_tokens)

    for st in ["<sos>", "<eos>", "<pad>"]:
        raw_response = raw_response.replace(st, "")
        
    return raw_response.strip().capitalize()

if __name__ == '__main__':
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    print(f"Using device: {device}")

    if not os.path.exists(TOKENIZER_PATH):
        raise FileNotFoundError(f"Tokenizer not found at {TOKENIZER_PATH}")

    tokenizer = BPETokenizer(vocab_size=VOCAB_SIZE_LIMIT)
    tokenizer.load(TOKENIZER_PATH)

    vocab_size = len(tokenizer.vocab)
    PAD_ID = tokenizer.token_to_id.get("<pad>", 0)

    print("Initializing model architecture...")
    model = Transformer(
        vocab_size=vocab_size, 
        d_model=D_MODEL, 
        num_heads=NUM_HEADS, 
        num_kv_heads=NUM_KV_HEADS, 
        num_layers=NUM_LAYERS, 
        d_ff=D_FF, 
        pad_id=PAD_ID,
        use_layer_scale=USE_LAYER_SCALE
    )

    if not os.path.exists(MODEL_PATH):
        raise FileNotFoundError(f"Model file not found: {MODEL_PATH}")
        
    print(f"Loading weights from {MODEL_PATH}...")
    model.load_state_dict(torch.load(MODEL_PATH, map_location="cpu"))
    print("Weights loaded successfully!")

    dtype = torch.bfloat16 if torch.cuda.is_bf16_supported() else torch.float16
    model = model.to(dtype=dtype).to(device)

    print("Model ready! Type 'exit' to quit.\n")

    while True:
        try:
            user_msg = input("You: ")
            if user_msg.lower() in ['exit', 'quit']: 
                break
                
            response = generate_response(
                model=model, 
                tokenizer=tokenizer, 
                user_input=user_msg, 
                max_seq_length=MAX_SEQ_LENGTH, 
                device=device
            )
            print(f"Bot: {response}\n")
            
        except KeyboardInterrupt:
            print("\nExiting...")
            break