"""
NLLB-200 with RAG Lexicon Enhancement for Deep Past
Uses 39K entry OA_Lexicon_eBL.csv for context-aware translation
"""
import pandas as pd
import torch
from pathlib import Path
from transformers import AutoModelForSeq2SeqLM, AutoTokenizer
import re
from collections import defaultdict

# Load lexicon at module level
LEXICON_PATH = "/kaggle/input/deep-past-initiative-machine-translation/OA_Lexicon_eBL.csv"

def load_lexicon():
    """Load and index the Akkadian lexicon"""
    lexicon_df = pd.read_csv(LEXICON_PATH)

    # Build lookup dictionaries - column names are 'form' and 'norm' for Akkadian
    word_to_meaning = {}
    for _, row in lexicon_df.iterrows():
        # Try 'norm' first, then 'form' for Akkadian
        akkadian = str(row.get("norm", "")).strip().lower()
        if not akkadian:
            akkadian = str(row.get("form", "")).strip().lower()

        # Try 'eBL' or 'lexeme' for English meaning
        english = str(row.get("eBL", "")).strip()
        if not english or english == "nan":
            english = str(row.get("lexeme", "")).strip()

        if akkadian and english and english != "nan":
            word_to_meaning[akkadian] = english

    return word_to_meaning

def get_lexicon_hints(text, lexicon, max_hints=5):
    """Extract relevant lexicon entries for input text"""
    text_lower = text.lower()
    hints = []

    # Find matches
    for akk_word, eng_meaning in lexicon.items():
        if akk_word in text_lower:
            hints.append(f"{akk_word}={eng_meaning}")
        if len(hints) >= max_hints:
            break

    return " | ".join(hints) if hints else ""

def preprocess_akkadian(text):
    """Clean and normalize Akkadian text"""
    # Handle unicode ellipsis
    text = text.replace('…', '<gap>')
    text = re.sub(r'\.+', '<gap>', text)
    # Normalize whitespace
    text = re.sub(r'\s+', ' ', text).strip()
    return text

def main():
    print("Loading NLLB model and lexicon...")

    # Load model - use local_files_only=True for Kaggle dataset paths
    model_path = "/kaggle/input/nllb-deep-past"
    tokenizer = AutoTokenizer.from_pretrained(model_path, local_files_only=True)
    model = AutoModelForSeq2SeqLM.from_pretrained(model_path, local_files_only=True)

    if torch.cuda.is_available():
        model = model.cuda()
    model.eval()

    # Load lexicon
    lexicon = load_lexicon()
    print(f"Loaded {len(lexicon)} lexicon entries")

    # Load test data
    test_df = pd.read_csv("/kaggle/input/deep-past-initiative-machine-translation/test.csv")
    print(f"Test samples: {len(test_df)}")
    print(f"Test columns: {list(test_df.columns)}")

    # Determine source column
    source_col = "transliteration"
    if source_col not in test_df.columns:
        for col in test_df.columns:
            if col.lower() not in ["id", "target", "translation"]:
                source_col = col
                break

    print(f"Using source column: {source_col}")

    # Generation config optimized for chrF++
    GEN_CONFIG = {
        "max_length": 256,
        "num_beams": 10,
        "repetition_penalty": 1.5,
        "length_penalty": 1.2,
        "no_repeat_ngram_size": 3,
        "early_stopping": True,
        "forced_bos_token_id": tokenizer.convert_tokens_to_ids("eng_Latn"),
    }

    translations = []

    for idx, row in test_df.iterrows():
        # Preprocess
        akkadian = preprocess_akkadian(row[source_col])

        # Get lexicon hints
        hints = get_lexicon_hints(akkadian, lexicon)

        # Build input with context
        if hints:
            # Prepend lexicon hints as context
            input_text = f"Akkadian (context: {hints}): {akkadian}"
        else:
            input_text = f"Akkadian: {akkadian}"

        # Tokenize
        inputs = tokenizer(input_text, return_tensors="pt", max_length=512, truncation=True)
        if torch.cuda.is_available():
            inputs = {k: v.cuda() for k, v in inputs.items()}

        # Generate
        with torch.no_grad():
            outputs = model.generate(**inputs, **GEN_CONFIG)

        translation = tokenizer.decode(outputs[0], skip_special_tokens=True)
        translations.append(translation)

        if (idx + 1) % 10 == 0:
            print(f"Processed {idx + 1}/{len(test_df)}")

    # Save with correct column name
    submission = pd.DataFrame({
        "id": test_df["id"],
        "translation": translations
    })

    submission.to_csv("submission.csv", index=False)
    print(f"Saved {len(submission)} translations")

    # Print sample outputs
    print("\nSample outputs:")
    for i in range(min(3, len(submission))):
        print(f"  {submission.iloc[i]['id']}: {submission.iloc[i]['translation'][:80]}...")

if __name__ == "__main__":
    main()