# 1. ต้อง import spaces เป็นบรรทัดแรกสุดสำหรับ Hugging Face ZeroGPU
try:
    import spaces
except ImportError:
    spaces = None

import os
import re
import torch
import torch.nn as nn
import torch.nn.functional as F
from pythainlp.tokenize import word_tokenize
import gradio as gr

# Setup Device
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

# 2. โหลด Vocab และ Weights
vocab_file = "vocab_v6.pt"
model_file = "my_llm_v6_big_brain.pth"

if not os.path.exists(vocab_file) or not os.path.exists(model_file):
    raise FileNotFoundError("❌ ไม่พบไฟล์ vocab_v6.pt หรือ my_llm_v6_big_brain.pth ในโฟลเดอร์!")

vocab_data = torch.load(vocab_file, map_location="cpu")
word2idx = vocab_data['word2idx']
idx2word = vocab_data['idx2word']
vocab_size = len(word2idx)
SPECIAL_TOKENS = ["<|user|>", "<|bot|>", "<|end|>"]

# 3. Helper Functions
def preprocess_text(text):
    text = re.sub(r'([\+\-\*/=])', r' \1 ', text)
    text = re.sub(r'\s+', ' ', text)
    return text.strip()

def tokenize_input(text):
    text = preprocess_text(text)
    sub_tokens = word_tokenize(text, engine="newmm")
    return [w for w in sub_tokens if w.strip()]

def encode(tokens):
    return [word2idx.get(w, 1) for w in tokens]

def decode(token_ids):
    words = []
    for i in token_ids:
        w = idx2word.get(i, '')
        if w not in SPECIAL_TOKENS and i not in [0, 1]:
            words.append(w)
    return "".join(words)

# 4. Model Architecture (V6 Big Brain)
MAX_LEN = 64

class BigBrainMiniLLM(nn.Module):
    def __init__(self, vocab_size, d_model=512, nhead=8, num_layers=6, dropout=0.15):
        super().__init__()
        self.embedding = nn.Embedding(vocab_size, d_model)
        self.pos_embedding = nn.Embedding(MAX_LEN, d_model)
        self.drop = nn.Dropout(dropout)
        
        encoder_layer = nn.TransformerEncoderLayer(
            d_model=d_model, nhead=nhead, dim_feedforward=2048, dropout=dropout, batch_first=True
        )
        self.transformer = nn.TransformerEncoder(encoder_layer, num_layers=num_layers)
        self.fc_out = nn.Linear(d_model, vocab_size)

    def forward(self, x):
        seq_len = x.size(1)
        positions = torch.arange(0, seq_len, device=x.device).unsqueeze(0)
        out = self.drop(self.embedding(x) + self.pos_embedding(positions))
        mask = torch.triu(torch.full((seq_len, seq_len), float('-inf'), device=x.device), diagonal=1)
        out = self.transformer(out, mask=mask)
        return self.fc_out(out)

model = BigBrainMiniLLM(vocab_size=vocab_size, d_model=512, nhead=8, num_layers=6, dropout=0.15).to(device)
model.load_state_dict(torch.load(model_file, map_location=device))
model.eval()

# 5. ZeroGPU Decorator Wrapper
def gpu_decorator(func):
    if spaces is not None and hasattr(spaces, "GPU"):
        return spaces.GPU(func)
    return func

@gpu_decorator
def bot_response(message, history, temp, top_k, rep_penalty):
    if not message.strip():
        return ""
        
    tokens = ["<|user|>"] + tokenize_input(message) + ["<|bot|>"]
    encoded = encode(tokens)
    input_tensor = torch.tensor(encoded, dtype=torch.long).unsqueeze(0).to(device)
    
    generated = list(encoded)
    bot_start_len = len(encoded)
    
    with torch.no_grad():
        for _ in range(40):
            if input_tensor.size(1) >= MAX_LEN - 1:
                break
            outputs = model(input_tensor)
            logits = outputs[0, -1, :] / temp
            
            # Repetition Penalty
            generated_bot = generated[bot_start_len:]
            for t in set(generated_bot):
                if logits[t] < 0:
                    logits[t] *= rep_penalty
                else:
                    logits[t] /= rep_penalty
            
            if top_k > 0:
                v, _ = torch.topk(logits, min(top_k, logits.size(-1)))
                logits[logits < v[-1]] = -float('Inf')
                
            probs = F.softmax(logits, dim=-1)
            next_token = torch.multinomial(probs, num_samples=1).item()
            
            if next_token == 0 or idx2word.get(next_token) == "<|end|>":
                break
                
            generated.append(next_token)
            input_tensor = torch.tensor(generated, dtype=torch.long).unsqueeze(0).to(device)
            
    return decode(generated[bot_start_len:])

# 6. Build Gradio Interface (ย้าย theme ไปไว้ที่ launch() ตาม Gradio 6.0)
with gr.Blocks() as demo:
    gr.Markdown(
        """
        # 🧠 Thai Mini-LLM V6 (Big Brain 23.7M)
        ### โมเดลภาษาไทยขนาดจิ๋ว เทรนจาก 0 ด้วยสถาปัตยกรรม Transformer (Word-Level)
        """
    )
    
    with gr.Row():
        with gr.Column(scale=3):
            gr.ChatInterface(
                fn=bot_response,
                additional_inputs=[
                    gr.Slider(0.1, 1.5, value=0.5, step=0.1, label="Temperature (ความมั่ว/ความคิดสร้างสรรค์)"),
                    gr.Slider(1, 50, value=10, step=1, label="Top-K (การจำกัดขอบเขตคลังคำ)"),
                    gr.Slider(1.0, 2.5, value=1.3, step=0.1, label="Repetition Penalty (บทลงโทษคำซ้ำ)"),
                ],
            )
            
    gr.Markdown("--- \n *Created with ❤️ | Custom PyTorch Model trained from Scratch*")

if __name__ == "__main__":
    demo.launch(theme=gr.themes.Soft())