import torch
from config import Config
from model import make_model
from tokenizer import Tokenizer

class Translator:
    """翻译器: 加载模型进行推理"""
    def __init__(self, config):
        self.config = config
        self.device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
        
        # 1. 加载分词器
        self.src_tokenizer = Tokenizer(config.src_tokenizer_path)
        self.tgt_tokenizer = Tokenizer(config.tgt_tokenizer_path)
        
        # 2. 构建模型
        self.model = make_model(
            src_vocab=config.src_vocab_size,
            tgt_vocab=config.tgt_vocab_size,
            N=config.n_layers,
            d_model=config.d_model,
            d_ff=config.d_ff,
            h=config.n_heads,
            dropout=0  # 推理时关闭dropout
        ).to(self.device)
        
        # 3. 加载权重
        checkpoint = torch.load(config.model_path, map_location=self.device)
        self.model.load_state_dict(checkpoint)
        self.model.eval()  # 推理模式
    
    @torch.no_grad()
    def translate(self, text, max_len=50):
        """
        翻译函数 (贪心解码)
        text: 源语言文本
        max_len: 最大生成长度
        return: 目标语言文本
        """
        # 1. 编码源语言
        src_ids = self.src_tokenizer.encode(text)
        src_tensor = torch.LongTensor([src_ids]).to(self.device)
        src_mask = self.make_src_mask(src_tensor)
        
        # 2. 编码器前向传播 (只计算一次)
        memory = self.model.encode(src_tensor, src_mask)
        
        # 3. 自回归生成
        # 初始输入: <S> (起始符)
        tgt_ids = [self.config.bos_idx]
        
        for _ in range(max_len):
            tgt_tensor = torch.LongTensor([tgt_ids]).to(self.device)
            tgt_mask = self.make_tgt_mask(tgt_tensor)
            
            # 解码器前向传播
            output = self.model.decode(
                memory, src_mask, tgt_tensor, tgt_mask
            )
            
            # 生成器: 取最后一个位置的预测
            logits = self.model.generator(output[:, -1, :])
            next_token = logits.argmax(dim=-1).item()
            
            # 添加到序列
            tgt_ids.append(next_token)
            
            # 遇到结束符停止
            if next_token == self.config.eos_idx:
                break
        
        # 4. 解码为目标语言文本
        # 去掉起始符和结束符
        result_ids = tgt_ids[1:-1] if tgt_ids[-1] == self.config.eos_idx else tgt_ids[1:]
        result = self.tgt_tokenizer.decode(result_ids)
        
        return result
    
    def make_src_mask(self, src):
        """源语言填充掩码"""
        return (src != self.config.pad_idx).unsqueeze(-2)
    
    def make_tgt_mask(self, tgt):
        """目标语言掩码: 填充掩码 + 因果掩码"""
        tgt_pad_mask = (tgt != self.config.pad_idx).unsqueeze(-2)
        tgt_len = tgt.size(1)
        # 因果掩码 (下三角)
        tgt_causal_mask = torch.tril(
            torch.ones(tgt_len, tgt_len, device=self.device)
        ).bool()
        # 合并两个掩码
        return tgt_pad_mask & tgt_causal_mask


# 使用示例
if __name__ == "__main__":
    config = Config()
    config.model_path = 'runs/best.pt'  # 权重路径
    
    translator = Translator(config)
    
    # 测试翻译
    test_sentences = [
        "i love you",
        "i am a dog",
        "The government has implemented a series of policies to improve the living standards of citizens"
    ]
    
    for src in test_sentences:
        result = translator.translate(src)
        print(f"英文: {src}")
        print(f"中文: {result}")
        print("-" * 50)