#!/usr/bin/env python3
"""
ByT5 RAG Inference v1.0.0

Tests partial RAG checkpoint (2/6 epochs, 1434 steps) with dictionary augmentation.
Based on jeanjean111's byt5-base-big-data2 model (32.6 LB).

Uses simple beam search (as proven effective by top scorers).
"""

import os
import json
import re
import pandas as pd
import torch
from torch.utils.data import DataLoader, Dataset
from transformers import AutoTokenizer, AutoModelForSeq2SeqLM
from tqdm.auto import tqdm

# ============================================================================
# Configuration
# ============================================================================
CONFIG = {
    # Model - RAG fine-tuned checkpoint (2 epochs)
    "model_path": "/kaggle/input/byt5-rag-akkadian-v1",
    # Dictionary for RAG augmentation
    "dict_path": "/kaggle/input/akkadian-dictionary/unified_akkadian_dict.json",
    # Inference settings
    "batch_size": 4,
    "max_source_length": 320,  # Extra space for DICT prefix
    "max_target_length": 256,
    # Generation - simple beam search (proven best by top scorers)
    "num_beams": 4,
    "early_stopping": True,
    # NO repetition_penalty - hurts scores
}


# ============================================================================
# RAG Augmenter (embedded for Kaggle)
# ============================================================================
class AkkadianRAG:
    """RAG augmenter for Akkadian translation."""

    CASE_ENDINGS = [
        "um",
        "im",
        "am",
        "tum",
        "tim",
        "tam",
        "ūtum",
        "ūtim",
        "ūtam",
        "ānū",
        "ānī",
        "ānē",
        "ūm",
        "ī",
        "ē",
        "ātu",
        "āti",
        "āt",
        "ū",
        "ī",
        "u",
        "i",
        "a",
        "ma",
        "mi",
    ]

    LOGOGRAM_PATTERN = re.compile(r"^[A-ZŠṢṬ][A-ZŠṢṬ0-9.]+$")

    def __init__(self, dict_path: str):
        """Load the unified dictionary."""
        with open(dict_path) as f:
            data = json.load(f)

        self.form_to_entries = data["form_to_entries"]
        self.root_to_meanings = data["root_to_meanings"]
        self.logogram_to_akkadian = data["logogram_to_akkadian"]

        print(
            f"Loaded dictionary: {len(self.form_to_entries)} forms, "
            f"{len(self.root_to_meanings)} roots, "
            f"{len(self.logogram_to_akkadian)} logograms"
        )

    def normalize(self, text: str) -> str:
        return text.lower().replace("-", "").strip()

    def strip_case_ending(self, word: str) -> str:
        normalized = self.normalize(word)
        for ending in sorted(self.CASE_ENDINGS, key=len, reverse=True):
            if normalized.endswith(ending) and len(normalized) > len(ending) + 1:
                return normalized[: -len(ending)]
        return normalized

    def lookup_word(self, word: str) -> list:
        # Check if logogram
        if self.LOGOGRAM_PATTERN.match(word):
            entries = self.logogram_to_akkadian.get(word.upper(), [])
            if entries:
                return [
                    {"root": word.upper(), "meaning": e["meaning"], "pos": "LOGO"}
                    for e in entries[:2]
                ]

        # Try exact form lookup
        normalized = self.normalize(word)
        entries = self.form_to_entries.get(normalized, [])
        if entries:
            return entries[:2]

        # Try with hyphen-removed
        no_hyphen = word.replace("-", "").lower()
        entries = self.form_to_entries.get(no_hyphen, [])
        if entries:
            return entries[:2]

        # Try stripping case endings
        stem = self.strip_case_ending(word)
        if stem != normalized:
            meanings = self.root_to_meanings.get(stem, [])
            if meanings:
                return [{"root": stem, "meaning": meanings[0], "pos": "?"}]

            for ending in ["u", "ū", ""]:
                root_guess = stem + ending
                meanings = self.root_to_meanings.get(root_guess, [])
                if meanings:
                    return [{"root": root_guess, "meaning": meanings[0], "pos": "?"}]

        return []

    def augment_text(self, text: str, max_entries: int = 10) -> str:
        tokens = text.split()
        dict_entries = []
        seen_roots = set()

        for token in tokens:
            clean = re.sub(r"[.,;:!?()[\]{}]", "", token)
            if len(clean) < 2:
                continue

            entries = self.lookup_word(clean)
            for entry in entries:
                root = entry["root"]
                if root not in seen_roots:
                    seen_roots.add(root)
                    dict_entries.append(f"{root}={entry['meaning']}")

                    if len(dict_entries) >= max_entries:
                        break

            if len(dict_entries) >= max_entries:
                break

        if dict_entries:
            dict_str = "; ".join(dict_entries)
            return f"[DICT: {dict_str}] {text}"
        else:
            return f"[DICT: ?] {text}"


# ============================================================================
# Dataset
# ============================================================================
PREFIX = "translate Akkadian to English: "


class RAGInferenceDataset(Dataset):
    def __init__(self, df, tokenizer, rag, max_length):
        self.texts = df["transliteration"].astype(str).tolist()
        self.tokenizer = tokenizer
        self.rag = rag
        self.max_length = max_length

        # Pre-augment all texts
        print("Augmenting texts with dictionary lookups...")
        self.augmented = [PREFIX + rag.augment_text(t) for t in tqdm(self.texts)]

    def __len__(self):
        return len(self.texts)

    def __getitem__(self, idx):
        text = self.augmented[idx]
        inputs = self.tokenizer(
            text,
            max_length=self.max_length,
            padding="max_length",
            truncation=True,
            return_tensors="pt",
        )
        return {
            "input_ids": inputs["input_ids"].squeeze(0),
            "attention_mask": inputs["attention_mask"].squeeze(0),
        }


# ============================================================================
# Main
# ============================================================================
def main():
    print("=" * 60)
    print("ByT5 RAG Inference v1.0.0")
    print("Model: RAG fine-tuned checkpoint (2 epochs)")
    print("Base: jeanjean111's byt5-base-big-data2")
    print("=" * 60)

    DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
    print(f"Device: {DEVICE}")

    # Load RAG augmenter
    print("\nLoading dictionary...")
    rag = AkkadianRAG(CONFIG["dict_path"])

    # Load model and tokenizer
    print(f"\nLoading model from {CONFIG['model_path']}...")
    tokenizer = AutoTokenizer.from_pretrained(CONFIG["model_path"])
    model = AutoModelForSeq2SeqLM.from_pretrained(CONFIG["model_path"]).to(DEVICE)
    model.eval()

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

    # Create dataset
    test_dataset = RAGInferenceDataset(
        test_df, tokenizer, rag, CONFIG["max_source_length"]
    )
    test_loader = DataLoader(
        test_dataset, batch_size=CONFIG["batch_size"], shuffle=False
    )

    # Show augmentation examples
    print("\nAugmentation examples:")
    for i in range(min(3, len(test_dataset))):
        print(f"  [{i}] {test_dataset.augmented[i][:100]}...")

    # Run inference
    print("\nRunning inference...")
    all_predictions = []

    with torch.no_grad():
        for batch in tqdm(test_loader):
            input_ids = batch["input_ids"].to(DEVICE)
            attention_mask = batch["attention_mask"].to(DEVICE)

            outputs = model.generate(
                input_ids=input_ids,
                attention_mask=attention_mask,
                max_length=CONFIG["max_target_length"],
                num_beams=CONFIG["num_beams"],
                early_stopping=CONFIG["early_stopping"],
            )

            decoded = tokenizer.batch_decode(outputs, skip_special_tokens=True)
            all_predictions.extend([d.strip() for d in decoded])

    # Create submission
    submission = pd.DataFrame({"id": test_df["id"], "translation": all_predictions})

    # Handle empty predictions
    submission["translation"] = submission["translation"].apply(
        lambda x: x if len(str(x)) > 0 else "broken text"
    )

    # Save
    submission.to_csv("submission.csv", index=False)
    print("\nSubmission saved!")
    print(submission.head(10))

    # Stats
    print(f"\nPrediction stats:")
    print(f"  Total: {len(all_predictions)}")
    print(f"  Empty: {sum(1 for p in all_predictions if not p)}")
    print(
        f"  Avg length: {sum(len(p) for p in all_predictions) / len(all_predictions):.1f}"
    )


if __name__ == "__main__":
    main()