import numpy as np
import torch
import torch.nn.functional as F
from transformers import MarianMTModel, MarianTokenizer
import spacy

MODEL_NAME = "staka/fugumt-en-ja"
SKIP_TOKENS = {"</s>", "<pad>", "<unk>"}
 
def load_models():
    tokenizer = MarianTokenizer.from_pretrained(MODEL_NAME)
    model = MarianMTModel.from_pretrained(MODEL_NAME, attn_implementation="eager")
    model.generation_config.max_length = None
    model.eval()
 
    nlp_en = spacy.load("en_core_web_sm")
    nlp_ja = spacy.load("ja_ginza", exclude=["compound_splitter"])

    return tokenizer, model, nlp_en, nlp_ja


def gradient_attention_matrix(text: str, tokenizer, model, translation: str = None):
    inputs = tokenizer(text, return_tensors="pt", padding=True)
 
    #obtain translation
    if translation is None:
        with torch.no_grad():
            gen_ids = model.generate(**inputs, max_new_tokens=30, num_beams=4, no_repeat_ngram_size=3, repetition_penalty=1.3 )
        translation = tokenizer.decode(gen_ids[0], skip_special_tokens=True)
 
    target_enc = tokenizer(translation, return_tensors="pt")
    target_ids = target_enc.input_ids
    decoder_input_ids = model.prepare_decoder_input_ids_from_labels(target_ids)
 
    #forward pass w gradients
    model.zero_grad()
    outputs = model(
        input_ids = inputs.input_ids,
        attention_mask = inputs.attention_mask,
        decoder_input_ids = decoder_input_ids,
        output_attentions = True,
    )
 
    #retain_grad()
    cross_attns = outputs.cross_attentions
    for layer_attn in cross_attns:
        layer_attn.retain_grad()
 
    #NLL
    logits = outputs.logits
    log_probs = F.log_softmax(logits, dim=-1)
    nll = -log_probs[0, torch.arange(target_ids.shape[1]), target_ids[0]].sum()
    nll.backward()
 
    #build matrix
    src_tokens = tokenizer.convert_ids_to_tokens(inputs["input_ids"][0])
    tgt_tokens = tokenizer.convert_ids_to_tokens(target_ids[0])
 
    layer_matrices = []
    for layer_attn in cross_attns:
        grad = layer_attn.grad
        if grad is None:
            continue
        grad_attn = (grad.abs() * layer_attn).mean(dim=1)[0]
        layer_matrices.append(grad_attn.detach().numpy())
 
    matrix = np.mean(layer_matrices, axis=0)
 
    for j, tok in enumerate(src_tokens):
        if tok in SKIP_TOKENS:
            matrix[:, j] = 0.0

    row_sums = matrix.sum(axis=1, keepdims=True)
    matrix = matrix / np.where(row_sums > 0, row_sums, 1.0)

    return {
        "translation": translation,
        "src_tokens":  src_tokens,
        "tgt_tokens":  tgt_tokens,
        "matrix":      matrix,
        "logits":      logits.detach()
    }
    
#GiNZA
def annotate(text_en: str, text_ja: str, nlp_en, nlp_ja):
    doc_en = nlp_en(text_en)
    doc_ja = nlp_ja(text_ja)
 
    en_tokens = [
        {
            "text": tok.text,
            "lemma": tok.lemma_,
            "pos": tok.pos_,
            "dep": tok.dep_,
            "ent_type": tok.ent_type_ or None,
        }
        for tok in doc_en
    ]
 
    ja_tokens = [
        {
            "text": tok.text,
            "lemma": tok.lemma_,
            "pos": tok.pos_,
            "dep": tok.dep_,
            "reading": tok.morph.get("Reading")[0] if tok.morph.get("Reading") else None,
            "ent_type": tok.ent_type_ or None,
        }
        for tok in doc_ja
    ]
 
    return {"en_tokens": en_tokens, "ja_tokens": ja_tokens}

def explain_token(tgt_idx: int, tgt_tokens: list, src_tokens: list, matrix: np.ndarray,
                annotations: dict, logits: torch.Tensor, tokenizer: MarianTokenizer, top_k: int = 3) -> dict:
    attn_row = matrix[tgt_idx]
    top_indices = np.argsort(attn_row)[::-1][:top_k]
    top_sources = [
        {"token": src_tokens[i], "weight": float(attn_row[i])}
        for i in top_indices
    ]
    
    tgt_surface = tgt_tokens[tgt_idx].replace("▁", "").replace("##", "")
    
    top5 = logits[0, tgt_idx].topk(5)
    alternatives = [
        t.replace("▁", "")
        for t in tokenizer.convert_ids_to_tokens(top5.indices.tolist())
        if t not in SKIP_TOKENS and t.replace("▁", "") != tgt_surface
    ]


    ja_anno = next(
        (t for t in annotations["ja_tokens"] if tgt_surface in t["text"]),
        None,
    )

    primary_src = top_sources[0]["token"].replace("▁", "").replace("##", "")
    en_anno = next(
        (t for t in annotations["en_tokens"] if primary_src.lower() in t["text"].lower()),
        None,
    )
 
    src_word = en_anno["text"] if en_anno else primary_src
    pos_label = en_anno["pos"].lower() if en_anno else ""
    alt_str = ", ".join([a for a in alternatives if a != tgt_surface][:3])
    ne_note = f" [{en_anno['ent_type']}]" if en_anno and en_anno.get("ent_type") else ""

    rationale = (
        f"「{tgt_surface}」← \"{src_word}\"{ne_note} ({pos_label})"
        + (f" | also considered: {alt_str}" if alt_str else "")
    )
 
    return {
        "tgt_token":     tgt_tokens[tgt_idx],
        "tgt_surface":   tgt_surface,
        "ja_annotation": ja_anno,
        "top_sources":   top_sources,
        "alternatives":  alternatives,
        "rationale":     rationale,
    }

def print_explanation(text_en: str, tokenizer, model, nlp_en, nlp_ja):
    print(f"\n{'='*70}")
    print(f"  Source:  {text_en}")

    grad_res = gradient_attention_matrix(text_en, tokenizer, model)
    print(f"  Output:  {grad_res['translation']}")

    annots = annotate(text_en, grad_res["translation"], nlp_en, nlp_ja)

    print("\n-- Per-token explanations --")
    for i in range(len(grad_res["tgt_tokens"])):
        exp = explain_token(
            i,
            grad_res["tgt_tokens"],
            grad_res["src_tokens"],
            grad_res["matrix"],
            annots,
            grad_res["logits"],
            tokenizer,
        )

        print(f"\n  [{i}] {exp['tgt_token']}")
        print(f"Top sources: {[s['token'] for s in exp['top_sources']]}")
        print(f"Weights: {[round(s['weight'], 3) for s in exp['top_sources']]}")
        print(f"Rationale: {exp['rationale']}")

if __name__ == "__main__":
    tokenizer, model, nlp_en, nlp_ja = load_models()

    
    examples = [
        "The server crashed because of too many requests.",
        "She quickly ran to the store before it closed.",
        "The children played in the park until sunset."
    ]
 
    for sentence in examples:
        print_explanation(sentence, tokenizer, model, nlp_en, nlp_ja)