"""
Trực quan hóa Ma trận Self-Attention của mô hình Transformer.
Chạy: python src/visualize_attention.py
Kết quả: Lưu file ảnh .png vào thư mục logs/
"""
import torch
import pickle
import os
import sys
import matplotlib
matplotlib.use('Agg')  # Non-interactive backend
import matplotlib.pyplot as plt
import seaborn as sns
import numpy as np

sys.stdout.reconfigure(encoding='utf-8')
sys.path.insert(0, os.path.join(os.path.dirname(__file__), '..'))
sys.path.insert(0, os.path.dirname(__file__))

from model import make_model

VOCAB_PATH     = 'data/processed/vocab.pkl'
MODEL_PATH     = 'models/best/model.pt'
TOKENIZER_PATH = 'data/processed/bpe_tokenizer.json'
OUTPUT_DIR     = 'logs'

os.makedirs(OUTPUT_DIR, exist_ok=True)


def load_bpe_tokenizer():
    if os.path.exists(TOKENIZER_PATH):
        from tokenizers import Tokenizer
        return Tokenizer.from_file(TOKENIZER_PATH)
    return None


def visualize_attention(question, model, device, word2idx, bpe_tok=None):
    """Trích xuất và vẽ ma trận Attention cho một câu hỏi."""

    # Tokenize
    if bpe_tok is not None:
        enc = bpe_tok.encode(question.strip())
        token_ids = enc.ids
        token_labels = enc.tokens
    else:
        import re
        text = str(question).lower()
        text = re.sub(r'([.,!?()])', r' \1 ', text)
        text = re.sub(r'\s{2,}', ' ', text)
        words = text.strip().split()
        token_ids = [word2idx.get(w, word2idx.get('<unk>', 3)) for w in words]
        token_labels = words

    sos_id = word2idx.get('<sos>', 1)
    input_ids = [sos_id] + token_ids
    token_labels = ['<sos>'] + token_labels

    x = torch.tensor([input_ids]).to(device)

    # Lấy attention maps
    attn_maps = model.get_attention_maps(x)

    n_layers = len(attn_maps)
    n_heads  = attn_maps[0].size(1)

    # Vẽ heatmap cho tất cả các layer (lấy trung bình các head)
    fig, axes = plt.subplots(1, n_layers, figsize=(6 * n_layers, 5))
    if n_layers == 1:
        axes = [axes]

    fig.suptitle(f'Self-Attention Maps\nInput: "{question}"',
                 fontsize=14, fontweight='bold', y=1.02)

    for layer_idx, (ax, attn) in enumerate(zip(axes, attn_maps)):
        # Trung bình các head: (1, n_head, seq, seq) -> (seq, seq)
        avg_attn = attn[0].mean(dim=0).cpu().numpy()

        sns.heatmap(
            avg_attn,
            xticklabels=token_labels,
            yticklabels=token_labels,
            cmap='YlOrRd',
            ax=ax,
            square=True,
            cbar_kws={'shrink': 0.6},
            vmin=0, vmax=1,
        )
        ax.set_title(f'Layer {layer_idx + 1}', fontsize=12, fontweight='bold')
        ax.set_xlabel('Key (Nhìn vào)')
        ax.set_ylabel('Query (Từ đang xét)')
        ax.tick_params(axis='x', rotation=45)
        ax.tick_params(axis='y', rotation=0)

    plt.tight_layout()
    safe_name = question[:30].replace(' ', '_').replace('?', '').replace('/', '_')
    save_path = os.path.join(OUTPUT_DIR, f'attention_{safe_name}.png')
    plt.savefig(save_path, dpi=150, bbox_inches='tight')
    plt.close()
    print(f"  Đã lưu: {save_path}")

    # Vẽ chi tiết từng head của Layer cuối
    last_attn = attn_maps[-1][0].cpu().numpy()  # (n_head, seq, seq)
    fig2, axes2 = plt.subplots(2, (n_heads + 1) // 2, figsize=(4 * ((n_heads + 1) // 2), 8))
    axes2 = axes2.flatten()

    fig2.suptitle(f'Attention Heads (Layer {n_layers})\nInput: "{question}"',
                  fontsize=14, fontweight='bold')

    for head_idx in range(n_heads):
        ax = axes2[head_idx]
        sns.heatmap(
            last_attn[head_idx],
            xticklabels=token_labels,
            yticklabels=token_labels,
            cmap='Blues',
            ax=ax,
            square=True,
            cbar=False,
            vmin=0, vmax=1,
        )
        ax.set_title(f'Head {head_idx + 1}', fontsize=10)
        ax.tick_params(axis='x', rotation=45, labelsize=7)
        ax.tick_params(axis='y', rotation=0, labelsize=7)

    # Ẩn các axes thừa
    for idx in range(n_heads, len(axes2)):
        axes2[idx].set_visible(False)

    plt.tight_layout()
    save_path2 = os.path.join(OUTPUT_DIR, f'attention_heads_{safe_name}.png')
    plt.savefig(save_path2, dpi=150, bbox_inches='tight')
    plt.close()
    print(f"  Đã lưu: {save_path2}")


def main():
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

    if not os.path.exists(MODEL_PATH):
        print(f"Không tìm thấy model: {MODEL_PATH}")
        print("Hãy chạy train_gpu.py trước.")
        return

    with open(VOCAB_PATH, 'rb') as f:
        vocab = pickle.load(f)
    word2idx = vocab['word2idx']

    ckpt = torch.load(MODEL_PATH, map_location=device)
    model = make_model(
        len(word2idx),
        n_layer=ckpt['n_l'], d_model=ckpt['d_m'],
        n_head=ckpt['n_h'], dropout=ckpt['dr']
    )
    model.load_state_dict(ckpt['model_state_dict'])
    model.to(device)
    model.eval()

    bpe_tok = load_bpe_tokenizer()
    print(f"Model: d_model={ckpt['d_m']} | layers={ckpt['n_l']} | heads={ckpt['n_h']}")
    print(f"BPE Tokenizer: {'Loaded' if bpe_tok else 'Fallback'}")
    print()

    # Danh sách câu hỏi mẫu để vẽ biểu đồ
    test_questions = [
        "Luật lũy đẳng là gì?",
        "Thuật toán Dijkstra hoạt động như thế nào?",
        "Đường đi Euler khác đường đi Hamilton chỗ nào?",
    ]

    print(f"{'='*55}")
    print(f"  VẼ BIỂU ĐỒ ATTENTION HEATMAP")
    print(f"{'='*55}")

    for q in test_questions:
        print(f"\n  Câu hỏi: \"{q}\"")
        visualize_attention(q, model, device, word2idx, bpe_tok)

    print(f"\n  Tất cả biểu đồ đã lưu vào thư mục: {OUTPUT_DIR}/")
    print(f"  Bạn có thể dán ảnh này vào Slide PowerPoint để báo cáo!")


if __name__ == "__main__":
    main()