# model.py
import re
import math
import json
import torch
import torch.nn as nn
import torch.nn.functional as F

# === Constants (must match training) ===
EMBED_DIM = 256
NUM_HEADS = 2
ENC_LAYERS = 2
DEC_LAYERS = 2
FFN_DIM = 512
DROPOUT = 0.1
MAX_LEN = 64

PAD_TOKEN = "<pad>"
SOS_TOKEN = "<sos>"
EOS_TOKEN = "<eos>"
UNK_TOKEN = "<unk>"

# === Text Preprocessing ===
def normalize_urdu(text):
    text = re.sub(r'[\u0610-\u061A\u064B-\u065F\u0670\u06D6-\u06ED]', '', text)
    text = re.sub('[إأآ]', 'ا', text)
    text = re.sub('[يى]', 'ی', text)
    text = re.sub('ؤ', 'و', text)
    text = re.sub(r'\s+', ' ', text).strip()
    return text

def simple_tokenize(text):
    text = re.sub(r'([،۔؟!,:؛\.\?\!\(\)\—\-\“\”\"\'«»])', r' \1 ', text)
    return text.split()

# === Vocab (for encoding/decoding) ===
class Vocab:
    def __init__(self, word2idx, idx2word):
        self.word2idx = word2idx
        self.idx2word = idx2word

    def encode(self, toks):
        ids = [self.word2idx[SOS_TOKEN]]
        for t in toks:
            ids.append(self.word2idx.get(t, self.word2idx[UNK_TOKEN]))
            if len(ids) >= MAX_LEN - 1:
                break
        ids.append(self.word2idx[EOS_TOKEN])
        ids += [self.word2idx[PAD_TOKEN]] * (MAX_LEN - len(ids))
        return ids[:MAX_LEN]

    def decode(self, ids):
        out = []
        for i in ids:
            if i >= len(self.idx2word):
                w = UNK_TOKEN
            else:
                w = self.idx2word[i]
            if w == EOS_TOKEN:
                break
            if w not in (SOS_TOKEN, PAD_TOKEN):
                out.append(w)
        return " ".join(out)

# === Model Components ===
class MultiHeadAttention(nn.Module):
    def __init__(self, d_model, num_heads):
        super().__init__()
        assert d_model % num_heads == 0
        self.num_heads = num_heads
        self.head_dim = d_model // num_heads
        self.w_q = nn.Linear(d_model, d_model)
        self.w_k = nn.Linear(d_model, d_model)
        self.w_v = nn.Linear(d_model, d_model)
        self.fc = nn.Linear(d_model, d_model)

    def forward(self, q, k, v, mask=None):
        B = q.size(0)
        q, k, v = self.w_q(q), self.w_k(k), self.w_v(v)
        def split(x):
            return x.view(B, -1, self.num_heads, self.head_dim).transpose(1, 2)
        q, k, v = split(q), split(k), split(v)
        attn = torch.matmul(q, k.transpose(-2, -1)) / math.sqrt(self.head_dim)
        if mask is not None:
            if mask.dim() == 2:
                mask = mask.unsqueeze(1).unsqueeze(1)
            elif mask.dim() == 3:
                mask = mask.unsqueeze(1)
            if mask.size(-1) != attn.size(-1):
                mask = mask[..., :attn.size(-1)]
            attn = attn.masked_fill(~mask, float('-inf'))
        w = F.softmax(attn, dim=-1)
        out = torch.matmul(w, v)
        out = out.transpose(1, 2).contiguous().view(B, -1, self.num_heads * self.head_dim)
        return self.fc(out)

class FeedForward(nn.Module):
    def __init__(self, d_model, d_ff):
        super().__init__()
        self.fc = nn.Sequential(
            nn.Linear(d_model, d_ff),
            nn.ReLU(),
            nn.Linear(d_ff, d_model)
        )
    def forward(self, x):
        return self.fc(x)

class PositionalEncoding(nn.Module):
    def __init__(self, d_model, max_len=5000):
        super().__init__()
        pe = torch.zeros(max_len, d_model)
        pos = torch.arange(0, max_len).unsqueeze(1).float()
        div = torch.exp(torch.arange(0, d_model, 2).float() * (-math.log(10000.0) / d_model))
        pe[:, 0::2] = torch.sin(pos * div)
        pe[:, 1::2] = torch.cos(pos * div)
        self.register_buffer("pe", pe.unsqueeze(0))
    def forward(self, x):
        return x + self.pe[:, :x.size(1)]

class EncoderLayer(nn.Module):
    def __init__(self, d_model, heads, d_ff, dropout=0.1):
        super().__init__()
        self.attn = MultiHeadAttention(d_model, heads)
        self.ff = FeedForward(d_model, d_ff)
        self.norm1 = nn.LayerNorm(d_model)
        self.norm2 = nn.LayerNorm(d_model)
        self.drop = nn.Dropout(dropout)
    def forward(self, x, mask):
        _x = self.attn(x, x, x, mask)
        x = self.norm1(x + self.drop(_x))
        _x = self.ff(x)
        x = self.norm2(x + self.drop(_x))
        return x

class DecoderLayer(nn.Module):
    def __init__(self, d_model, heads, d_ff, dropout=0.1):
        super().__init__()
        self.self_attn = MultiHeadAttention(d_model, heads)
        self.cross_attn = MultiHeadAttention(d_model, heads)
        self.ff = FeedForward(d_model, d_ff)
        self.norm1 = nn.LayerNorm(d_model)
        self.norm2 = nn.LayerNorm(d_model)
        self.norm3 = nn.LayerNorm(d_model)
        self.drop = nn.Dropout(dropout)
    def forward(self, x, enc_out, src_mask, tgt_mask):
        _x = self.self_attn(x, x, x, tgt_mask)
        x = self.norm1(x + self.drop(_x))
        _x = self.cross_attn(x, enc_out, enc_out, src_mask)
        x = self.norm2(x + self.drop(_x))
        _x = self.ff(x)
        x = self.norm3(x + self.drop(_x))
        return x

class SimpleTransformer(nn.Module):
    def __init__(self, vocab_size):
        super().__init__()
        self.embed = nn.Embedding(vocab_size, EMBED_DIM)
        self.pos = PositionalEncoding(EMBED_DIM)
        self.encoder = nn.ModuleList([
            EncoderLayer(EMBED_DIM, NUM_HEADS, FFN_DIM, DROPOUT)
            for _ in range(ENC_LAYERS)
        ])
        self.decoder = nn.ModuleList([
            DecoderLayer(EMBED_DIM, NUM_HEADS, FFN_DIM, DROPOUT)
            for _ in range(DEC_LAYERS)
        ])
        self.fc = nn.Linear(EMBED_DIM, vocab_size)

    def encode(self, src, src_mask):
        src_emb = self.embed(src) * math.sqrt(EMBED_DIM)
        src_pos = self.pos(src_emb)
        x = src_pos
        for layer in self.encoder:
            x = layer(x, src_mask)
        return x

    def decode(self, tgt, memory, src_mask, tgt_mask):
        tgt_emb = self.embed(tgt) * math.sqrt(EMBED_DIM)
        tgt_pos = self.pos(tgt_emb)
        x = tgt_pos
        for layer in self.decoder:
            x = layer(x, memory, src_mask, tgt_mask)
        return x

    def forward(self, src, tgt):
        src_padding_mask = (src != 0).unsqueeze(1).unsqueeze(2)
        tgt_padding_mask = (tgt != 0)
        tgt_len = tgt.size(1)
        causal_mask = torch.tril(torch.ones(tgt_len, tgt_len, device=tgt.device)).bool()
        tgt_mask = causal_mask & tgt_padding_mask.unsqueeze(1)
        tgt_mask = tgt_mask.unsqueeze(1)
        memory = self.encode(src, src_padding_mask)
        dec_out = self.decode(tgt, memory, src_padding_mask, tgt_mask)
        return self.fc(dec_out)

# === Inference Function ===
def greedy_decode(model, src, vocab, max_len=50):
    model.eval()
    DEVICE = next(model.parameters()).device
    src = src.to(DEVICE)
    assert src.size(0) == 1

    src_padding_mask = (src != vocab.word2idx[PAD_TOKEN]).unsqueeze(1).unsqueeze(2)
    with torch.no_grad():
        memory = model.encode(src, src_padding_mask)

    ys = torch.full((1, 1), vocab.word2idx[SOS_TOKEN], dtype=torch.long, device=DEVICE)

    for _ in range(max_len - 1):
        tgt_len = ys.size(1)
        causal_mask = torch.tril(torch.ones(tgt_len, tgt_len, device=DEVICE)).bool().unsqueeze(0).unsqueeze(0)
        with torch.no_grad():
            dec_out = model.decode(ys, memory, src_padding_mask, causal_mask)
            logits = model.fc(dec_out[:, -1])
            next_word = logits.argmax(dim=-1).item()
        ys = torch.cat([ys, torch.tensor([[next_word]], device=DEVICE)], dim=1)
        if next_word == vocab.word2idx[EOS_TOKEN]:
            break

    tokens = []
    for idx in ys[0, 1:]:
        if idx == vocab.word2idx[EOS_TOKEN]:
            break
        tokens.append(vocab.idx2word[idx.item()])
    return " ".join(tokens)

# === Loader ===
def load_model_and_vocab(model_path, vocab_path, device="cpu"):
    # Load vocab
    with open(vocab_path, "r", encoding="utf-8") as f:
        vocab_data = json.load(f)
    vocab = Vocab(vocab_data["word2idx"], vocab_data["idx2word"])

    # Load model
    model = SimpleTransformer(len(vocab.idx2word))
    checkpoint = torch.load(model_path, map_location=device)
    model.load_state_dict(checkpoint["model_state_dict"])
    model.eval()
    return model, vocab