"""Transformer AttRes model definition for Vietnamese to VSL translation.

Usage:
    from model import Transformer, TransformerConfig, Tokenizer, load_model
    model, src_tok, tgt_tok = load_model('path/to/model', device='cuda')
    translation = translate(model, src_tok, tgt_tok, "Xin chào")
"""
import os
import math
import json
import pickle
import torch
import torch.nn as nn
import torch.nn.functional as F


class Tokenizer:
    """Tokenizer word-level với các special tokens."""
    PAD, UNK, START, END = 0, 1, 2, 3

    def __init__(self):
        self.word2idx = {}
        self.idx2word = {}
        self.vocab_size = 0

    def fit(self, sentences):
        special = ['<pad>', '<unk>', '<start>', '<end>']
        vocab = sorted(set(w for s in sentences for w in s.split()) - set(special))
        all_tokens = special + vocab
        self.word2idx = {w: i for i, w in enumerate(all_tokens)}
        self.idx2word = {i: w for w, i in self.word2idx.items()}
        self.vocab_size = len(self.word2idx)
        return self

    def encode(self, sentence):
        return [self.word2idx.get(w, self.UNK) for w in sentence.split()]

    def decode(self, ids, skip_special=True):
        special_ids = {self.PAD, self.START, self.END}
        tokens = []
        for i in ids:
            if skip_special and i in special_ids:
                continue
            tokens.append(self.idx2word.get(i, '<unk>'))
        return ' '.join(tokens)


class PositionalEncoding(nn.Module):
    def __init__(self, d_model, max_len, dropout):
        super().__init__()
        self.dropout = nn.Dropout(dropout)
        pe = torch.zeros(max_len, d_model)
        pos = torch.arange(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)
        pe = pe.unsqueeze(0)
        self.register_buffer('pe', pe)

    def forward(self, x):
        x = x + self.pe[:, :x.size(1)]
        return self.dropout(x)


def make_pad_mask(seq, pad_idx=0):
    return (seq == pad_idx).unsqueeze(1).unsqueeze(2)


def make_causal_mask(seq_len, device):
    mask = torch.triu(torch.ones(seq_len, seq_len, device=device), diagonal=1).bool()
    return mask.unsqueeze(0).unsqueeze(0)


class MultiHeadAttention(nn.Module):
    def __init__(self, d_model, num_heads, dropout):
        super().__init__()
        assert d_model % num_heads == 0
        self.num_heads = num_heads
        self.d_k = d_model // num_heads
        self.w_q = nn.Linear(d_model, d_model, bias=False)
        self.w_k = nn.Linear(d_model, d_model, bias=False)
        self.w_v = nn.Linear(d_model, d_model, bias=False)
        self.w_o = nn.Linear(d_model, d_model, bias=False)
        self.dropout = nn.Dropout(dropout)

    def split_heads(self, x):
        B, S, _ = x.shape
        return x.view(B, S, self.num_heads, self.d_k).transpose(1, 2)

    def forward(self, q, k, v, mask=None):
        q = self.split_heads(self.w_q(q))
        k = self.split_heads(self.w_k(k))
        v = self.split_heads(self.w_v(v))
        scores = torch.matmul(q, k.transpose(-2, -1)) / math.sqrt(self.d_k)
        if mask is not None:
            scores = scores.masked_fill(mask, float('-inf'))
        weights = self.dropout(F.softmax(scores, dim=-1))
        out = torch.matmul(weights, v)
        out = out.transpose(1, 2).contiguous()
        B, S, _, _ = out.shape
        out = out.view(B, S, -1)
        return self.w_o(out)


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

    def forward(self, x):
        rms = x.float().pow(2).mean(-1, keepdim=True).add(self.eps).rsqrt()
        return (x.float() * rms).to(x.dtype) * self.weight


def block_attn_res(blocks, partial_block, proj, norm):
    V = torch.stack(blocks + [partial_block])
    K = norm(V)
    logits = torch.einsum('d, n b t d -> n b t', proj.weight.squeeze(), K)
    h = torch.einsum('n b t, n b t d -> b t d', logits.softmax(0), V)
    return h


class EncoderLayer(nn.Module):
    def __init__(self, d_model, num_heads, dff, dropout, layer_idx, block_size):
        super().__init__()
        self.layer_idx = layer_idx
        self.block_size = block_size
        self.self_attn = MultiHeadAttention(d_model, num_heads, dropout)
        self.ffn = nn.Sequential(
            nn.Linear(d_model, dff), nn.ReLU(), nn.Dropout(dropout), nn.Linear(dff, d_model)
        )
        self.norm1 = nn.LayerNorm(d_model)
        self.norm2 = nn.LayerNorm(d_model)
        self.drop1 = nn.Dropout(dropout)
        self.drop2 = nn.Dropout(dropout)
        self.attn_res_proj = nn.Linear(d_model, 1, bias=False)
        self.attn_res_norm = RMSNorm(d_model)
        self.mlp_res_proj = nn.Linear(d_model, 1, bias=False)
        self.mlp_res_norm = RMSNorm(d_model)

    def forward(self, blocks, partial_block, src_mask=None):
        h = block_attn_res(blocks, partial_block, self.attn_res_proj, self.attn_res_norm)
        h_norm = self.norm1(h)
        partial_block = partial_block + self.drop1(self.self_attn(h_norm, h_norm, h_norm, src_mask))
        h = block_attn_res(blocks, partial_block, self.mlp_res_proj, self.mlp_res_norm)
        partial_block = partial_block + self.drop2(self.ffn(self.norm2(h)))
        if self.layer_idx % self.block_size == 0:
            blocks = blocks + [partial_block]
            partial_block = None
        return blocks, partial_block


class DecoderLayer(nn.Module):
    def __init__(self, d_model, num_heads, dff, dropout, layer_idx, block_size):
        super().__init__()
        self.layer_idx = layer_idx
        self.block_size = block_size
        self.self_attn = MultiHeadAttention(d_model, num_heads, dropout)
        self.cross_attn = MultiHeadAttention(d_model, num_heads, dropout)
        self.ffn = nn.Sequential(
            nn.Linear(d_model, dff), nn.ReLU(), nn.Dropout(dropout), nn.Linear(dff, d_model)
        )
        self.norm1 = nn.LayerNorm(d_model)
        self.norm2 = nn.LayerNorm(d_model)
        self.norm3 = nn.LayerNorm(d_model)
        self.drop1 = nn.Dropout(dropout)
        self.drop2 = nn.Dropout(dropout)
        self.drop3 = nn.Dropout(dropout)
        self.self_attn_res_proj = nn.Linear(d_model, 1, bias=False)
        self.self_attn_res_norm = RMSNorm(d_model)
        self.cross_attn_res_proj = nn.Linear(d_model, 1, bias=False)
        self.cross_attn_res_norm = RMSNorm(d_model)
        self.mlp_res_proj = nn.Linear(d_model, 1, bias=False)
        self.mlp_res_norm = RMSNorm(d_model)

    def forward(self, blocks, partial_block, enc_out, tgt_mask=None, src_mask=None):
        h = block_attn_res(blocks, partial_block, self.self_attn_res_proj, self.self_attn_res_norm)
        h_norm = self.norm1(h)
        partial_block = partial_block + self.drop1(self.self_attn(h_norm, h_norm, h_norm, tgt_mask))
        h = block_attn_res(blocks, partial_block, self.cross_attn_res_proj, self.cross_attn_res_norm)
        partial_block = partial_block + self.drop2(self.cross_attn(self.norm2(h), enc_out, enc_out, src_mask))
        h = block_attn_res(blocks, partial_block, self.mlp_res_proj, self.mlp_res_norm)
        partial_block = partial_block + self.drop3(self.ffn(self.norm3(h)))
        if self.layer_idx % self.block_size == 0:
            blocks = blocks + [partial_block]
            partial_block = None
        return blocks, partial_block


class Encoder(nn.Module):
    def __init__(self, vocab_size, d_model, num_layers, num_heads, dff, max_len, dropout, block_size):
        super().__init__()
        self.embedding = nn.Embedding(vocab_size, d_model, padding_idx=0)
        self.pos_enc = PositionalEncoding(d_model, max_len, dropout)
        self.layers = nn.ModuleList([
            EncoderLayer(d_model, num_heads, dff, dropout, layer_idx=i+1, block_size=block_size)
            for i in range(num_layers)
        ])
        self.norm = nn.LayerNorm(d_model)

    def forward(self, src, src_mask=None):
        x = self.pos_enc(self.embedding(src) * math.sqrt(self.embedding.embedding_dim))
        blocks = [x]
        partial_block = x
        for layer in self.layers:
            blocks, partial_block = layer(blocks, partial_block, src_mask)
            if partial_block is None:
                partial_block = blocks[-1]
        return self.norm(partial_block)


class Decoder(nn.Module):
    def __init__(self, vocab_size, d_model, num_layers, num_heads, dff, max_len, dropout, block_size):
        super().__init__()
        self.embedding = nn.Embedding(vocab_size, d_model, padding_idx=0)
        self.pos_enc = PositionalEncoding(d_model, max_len, dropout)
        self.layers = nn.ModuleList([
            DecoderLayer(d_model, num_heads, dff, dropout, layer_idx=i+1, block_size=block_size)
            for i in range(num_layers)
        ])
        self.norm = nn.LayerNorm(d_model)

    def forward(self, tgt, enc_out, tgt_mask=None, src_mask=None):
        x = self.pos_enc(self.embedding(tgt) * math.sqrt(self.embedding.embedding_dim))
        blocks = [x]
        partial_block = x
        for layer in self.layers:
            blocks, partial_block = layer(blocks, partial_block, enc_out, tgt_mask, src_mask)
            if partial_block is None:
                partial_block = blocks[-1]
        return self.norm(partial_block)


class TransformerConfig:
    def __init__(self, src_vocab, tgt_vocab, d_model, num_layers,
                 num_heads, dff, max_len, dropout, block_size):
        self.src_vocab = src_vocab
        self.tgt_vocab = tgt_vocab
        self.d_model = d_model
        self.num_layers = num_layers
        self.num_heads = num_heads
        self.dff = dff
        self.max_len = max_len
        self.dropout = dropout
        self.block_size = block_size

    def to_dict(self):
        return self.__dict__.copy()

    @classmethod
    def from_dict(cls, d):
        return cls(**d)

    def save(self, path):
        with open(path, "w") as f:
            json.dump(self.to_dict(), f, indent=2)

    @classmethod
    def load(cls, path):
        with open(path) as f:
            return cls.from_dict(json.load(f))


class Transformer(nn.Module):
    def __init__(self, src_vocab, tgt_vocab, d_model, num_layers,
                 num_heads, dff, max_len, dropout, block_size):
        super().__init__()
        self.config = TransformerConfig(
            src_vocab=src_vocab, tgt_vocab=tgt_vocab, d_model=d_model,
            num_layers=num_layers, num_heads=num_heads, dff=dff,
            max_len=max_len, dropout=dropout, block_size=block_size
        )
        self.encoder = Encoder(src_vocab, d_model, num_layers, num_heads, dff, max_len, dropout, block_size)
        self.decoder = Decoder(tgt_vocab, d_model, num_layers, num_heads, dff, max_len, dropout, block_size)
        self.final_layer = nn.Linear(d_model, tgt_vocab)

    def forward(self, src, tgt):
        src_pad_mask = make_pad_mask(src).to(src.device)
        tgt_pad_mask = make_pad_mask(tgt).to(tgt.device)
        causal_mask = make_causal_mask(tgt.size(1), tgt.device)
        tgt_mask = tgt_pad_mask | causal_mask
        enc_out = self.encoder(src, src_pad_mask)
        dec_out = self.decoder(tgt, enc_out, tgt_mask, src_pad_mask)
        return self.final_layer(dec_out)

    def save_pretrained(self, save_directory):
        os.makedirs(save_directory, exist_ok=True)
        self.config.save(os.path.join(save_directory, "config.json"))
        torch.save(self.state_dict(), os.path.join(save_directory, "pytorch_model.bin"))

    @classmethod
    def from_pretrained(cls, load_directory, device="cpu"):
        config = TransformerConfig.load(os.path.join(load_directory, "config.json"))
        model = cls(**config.to_dict())
        state_dict = torch.load(
            os.path.join(load_directory, "pytorch_model.bin"),
            map_location=device
        )
        model.load_state_dict(state_dict)
        model.to(device)
        return model


def preprocess_sentence(sentence):
    """Tiền xử lý câu: lowercase, loại bỏ ký tự đặc biệt."""
    import re
    sentence = str(sentence).lower().strip()
    sentence = re.sub(r"[^\w\s]", "", sentence)
    return sentence


def load_model(model_dir, device="cpu"):
    """Load model và tokenizers từ thư mục.

    Returns:
        model, src_tokenizer, tgt_tokenizer
    """
    model = Transformer.from_pretrained(model_dir, device=device)
    with open(os.path.join(model_dir, "src_tokenizer.pkl"), "rb") as f:
        src_tokenizer = pickle.load(f)
    with open(os.path.join(model_dir, "tgt_tokenizer.pkl"), "rb") as f:
        tgt_tokenizer = pickle.load(f)
    return model, src_tokenizer, tgt_tokenizer


@torch.no_grad()
def translate(model, src_tokenizer, tgt_tokenizer, sentence, device="cpu", max_len=50):
    """Dịch một câu tiếng Việt → ngôn ngữ ký hiệu Việt Nam."""
    model.eval()
    src_ids = src_tokenizer.encode(preprocess_sentence(sentence))
    src_ids = src_ids[:max_len] + [src_tokenizer.PAD] * max(0, max_len - len(src_ids))
    src_tensor = torch.tensor([src_ids], dtype=torch.long).to(device)
    src_mask = make_pad_mask(src_tensor).to(device)
    enc_out = model.encoder(src_tensor, src_mask)
    tgt_ids = [tgt_tokenizer.START]
    for _ in range(max_len):
        tgt_tensor = torch.tensor([tgt_ids], dtype=torch.long).to(device)
        tgt_mask = make_causal_mask(len(tgt_ids), device)
        dec_out = model.decoder(tgt_tensor, enc_out, tgt_mask, src_mask)
        logits = model.final_layer(dec_out)
        next_id = logits[0, -1].argmax().item()
        if next_id == tgt_tokenizer.END:
            break
        tgt_ids.append(next_id)
    return tgt_tokenizer.decode(tgt_ids[1:])