# ---
# jupyter:
#   jupytext:
#     text_representation:
#       extension: .py
#       format_name: percent
#   kernelspec:
#     display_name: Python 3
#     language: python
#     name: python3
# ---

# %% [markdown]
# # Akkadian T5 Translation - Best Inference
#
# This notebook demonstrates best practices for Akkadian-to-English translation
# using our fine-tuned T5-Base model for the Deep Past Initiative competition.
#
# ## Model Details
# - **Model**: T5-Base (220M parameters)
# - **Training**: 1000 steps on enriched Akkadian dataset
# - **Handle**: `manwithacat/akkadian-t5-enriched-v2/pytorch/transformers`

# %% [markdown]
# ## 1. Configuration & Reproducibility

# %%
# Reproducibility header (per Kaggle best practices)
MODEL_HANDLE = "manwithacat/akkadian-t5-enriched-v2/pytorch/transformers"
MODEL_VERSION = 2
DATASET_HANDLE = "deep-past-initiative-machine-translation"

CONFIG = {
    # Generation parameters
    "max_src_len": 512,
    "max_tgt_len": 192,
    "num_beams": 4,
    "length_penalty": 1.0,
    "repetition_penalty": 2.5,
    "no_repeat_ngram_size": 4,
}

# Print environment info
import sys
import platform

print("Environment Info:")
print(f"  Python: {sys.version}")
print(f"  Platform: {platform.platform()}")

# %%
import os
os.environ["PROTOCOL_BUFFERS_PYTHON_IMPLEMENTATION"] = "python"

import re
import pandas as pd
from pathlib import Path

import torch
from transformers import AutoTokenizer, AutoModelForSeq2SeqLM

print(f"\nPyTorch: {torch.__version__}")
import transformers
print(f"Transformers: {transformers.__version__}")

# Device setup
DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
print(f"\nDevice: {DEVICE}")
if DEVICE == "cuda":
    print(f"GPU: {torch.cuda.get_device_name(0)}")

# %% [markdown]
# ## 2. Paths
#
# Model is attached via `model_sources` in kernel-metadata.json.
# Competition data is attached via `competition_sources`.

# %%
# Paths (Kaggle environment)
INPUT_DIR = Path("/kaggle/input/deep-past-initiative-machine-translation")
MODEL_PATH = Path(f"/kaggle/input/akkadian-t5-enriched-v2/pytorch/transformers/{MODEL_VERSION}")

print(f"Competition data: {INPUT_DIR}")
print(f"Model path: {MODEL_PATH}")

# %% [markdown]
# ## 3. Character Normalization
#
# Akkadian transliteration uses various Unicode conventions. We normalize to a
# canonical form matching our training data.

# %%
# Character mapping: normalize CDLI/ORACC/Unicode variants to canonical form
CHAR_NORMALIZE_MAP = {
    # Accented vowels -> subscript notation
    'á': 'a₂', 'à': 'a₃',
    'é': 'e₂', 'è': 'e₃',
    'í': 'i₂', 'ì': 'i₃',
    'ú': 'u₂', 'ù': 'u₃',
    # Special consonants (ensure Unicode form)
    'ṣ': 'ṣ', 'Ṣ': 'Ṣ',
    'ṭ': 'ṭ', 'Ṭ': 'Ṭ',
    'ḫ': 'ḫ', 'Ḫ': 'Ḫ',
    # Apostrophe variants -> canonical
    ''': 'ʾ', "'": 'ʾ', ''': 'ʾ',
}

MULTI_CHAR_PATTERNS = [('sz', 'š'), ('SZ', 'Š')]


def normalize_transliteration(text: str) -> str:
    """Normalize transliteration characters to canonical Unicode form."""
    if pd.isna(text):
        return ""
    text = str(text)

    # Apply multi-char patterns first
    for pattern, replacement in MULTI_CHAR_PATTERNS:
        text = text.replace(pattern, replacement)

    # Apply single-char mappings
    result = []
    for char in text:
        result.append(CHAR_NORMALIZE_MAP.get(char, char))

    return ''.join(result)


# Test normalization
test_text = "á-na szar-ri-im"
print(f"Normalization test: '{test_text}' -> '{normalize_transliteration(test_text)}'")

# %% [markdown]
# ## 4. Preprocessing

# %%
def clean_whitespace(text: str) -> str:
    """Normalize whitespace in text."""
    if pd.isna(text):
        return ""
    return re.sub(r"\s+", " ", str(text)).strip()


def preprocess_source(text: str, prefix: str = "translate Akkadian to English: ") -> str:
    """Preprocess source transliteration for T5 input."""
    text = str(text)

    # Normalize character variants
    text = normalize_transliteration(text)

    # Clean punctuation but keep structure
    text = re.sub(r'[!?„"""\'\[\]<>˹˺…+]', '', text)
    text = text.replace('/', ' ')

    # Normalize spacing
    text = ' '.join(text.split())

    # Add T5-style prefix
    return prefix + text


# Test preprocessing
test_input = "um-ma kà-ru-um [kà-ni-ia-ma]"
print(f"Preprocessing test:")
print(f"  Input:  '{test_input}'")
print(f"  Output: '{preprocess_source(test_input)}'")

# %% [markdown]
# ## 5. Post-Processing
#
# Post-processing normalizes model output to match expected translation style.

# %%
class TranslationPostProcessor:
    """Post-process translations to match expected style."""

    def __init__(self):
        self.akkadian_terms = {
            'waklum': 'waklum',
            'karum': 'kārum',
            'kdrum': 'kārum',
            'nibum': 'nibum',
            'tamkarum': 'tamkārum',
        }

    def normalize_opening(self, text: str, transliteration: str) -> str:
        """Normalize letter opening to training data style."""
        is_letter = bool(re.search(r'um-ma.*qí-bi|a-na.*qí-bi', transliteration, re.IGNORECASE))

        if not is_letter:
            return text

        match = re.match(
            r'^\\?n?\s*\[?Thus\]?\s*\(?(?:speaks?|says?)\)?\s*(?:the\s+)?(\w[\w\s-]*?)(?:\s*:\s*|\s+)(?:say|tell)\s+(?:to\s+)?(.+?)(?:\.|:|$)',
            text, re.IGNORECASE
        )
        if match:
            sender = match.group(1).strip()
            recipient = match.group(2).strip()
            recipient = re.sub(r'\s*[,;].*$', '', recipient)
            rest = text[match.end():].strip()
            return f"From {sender} to {recipient}: {rest}"

        return text

    def normalize_parentheses(self, text: str) -> str:
        """Remove unnecessary parenthetical content."""
        text = re.sub(r'\((?:speaks?|says?|laden|loaded|thus)\)', '', text)
        text = re.sub(r'\(\s*(?:l\.\s*)?\d+(?:-\d+)?\s*\)', '', text)
        text = re.sub(r'\((\w+)\)', r'\1', text)
        text = re.sub(r'\s+', ' ', text).strip()
        return text

    def normalize_colons(self, text: str) -> str:
        """Normalize multiple colons in text."""
        parts = text.split(':')
        if len(parts) <= 2:
            return text

        result = parts[0] + ':'
        for i, part in enumerate(parts[1:], 1):
            if i == 1:
                result += part
            elif part.strip() and part.strip()[0].isupper():
                result += '. ' + part.strip()
            else:
                result += '; ' + part.strip()

        return result

    def normalize_numbers(self, text: str) -> str:
        """Standardize number + unit phrases."""
        text = re.sub(r'(\d+)\s+shekels?\s+(?:of\s+)?silver', r'\1 shekels of silver', text)
        text = re.sub(r'(\d+)\s+talents?\s+(?:of\s+)?tin', r'\1 talents of tin', text)
        text = re.sub(r'(\d+)\s+minas?\s+(?:of\s+)?copper', r'\1 minas of copper', text)
        return text

    def normalize_akkadian_terms(self, text: str) -> str:
        """Standardize Akkadian loanwords."""
        for variant, standard in self.akkadian_terms.items():
            text = re.sub(rf'\b{variant}\b', standard, text, flags=re.IGNORECASE)
        return text

    def postprocess(self, text: str, transliteration: str = '') -> str:
        """Apply all post-processing steps."""
        text = text.replace('\\n', ' ')
        text = self.normalize_opening(text, transliteration)
        text = self.normalize_parentheses(text)
        text = self.normalize_colons(text)
        text = self.normalize_numbers(text)
        text = self.normalize_akkadian_terms(text)
        text = re.sub(r'\s+', ' ', text).strip()

        if text and text[-1] not in '.!?':
            text += '.'

        return text


postprocessor = TranslationPostProcessor()

# %%
def remove_repetitions(text: str) -> str:
    """Remove repetitive phrases from text."""
    text = re.sub(r'(,\s*(?:and|or)\s+[^,]+?)(?:\1){2,}', r'\1', text)
    text = re.sub(r'(\b\w+(?:\s+\w+){0,3}),\s*(?:\1,?\s*){2,}', r'\1', text)
    return text


def format_single_sentence(text: str) -> str:
    """Format output as single sentence (competition requirement)."""
    text = clean_whitespace(text)
    if not text:
        return "broken text."

    text = remove_repetitions(text)

    # Split on sentence endings and rejoin with semicolons
    parts = re.split(r"(?<=[.!?])\s+", text)
    parts = [p.strip() for p in parts if p.strip()]
    parts = [re.sub(r"[.!?]+\s*$", "", p).strip() for p in parts]
    parts = [p for p in parts if p]
    out = "; ".join(parts).strip()

    if not out.endswith("."):
        out += "."
    return out

# %% [markdown]
# ## 6. Load Model

# %%
print(f"\nLoading model from: {MODEL_PATH}")
print(f"Model files: {list(MODEL_PATH.iterdir())}")

tokenizer = AutoTokenizer.from_pretrained(MODEL_PATH, local_files_only=True)
model = AutoModelForSeq2SeqLM.from_pretrained(MODEL_PATH, local_files_only=True)
model.to(DEVICE)
model.eval()

print(f"\nModel loaded successfully!")
print(f"  Type: {model.config.model_type}")
print(f"  Parameters: {sum(p.numel() for p in model.parameters()):,}")

# %% [markdown]
# ## 7. Translation Functions

# %%
def generate_translation(src_text: str) -> str:
    """Generate translation using the trained model."""
    src_preprocessed = preprocess_source(src_text)

    enc = tokenizer(
        src_preprocessed,
        return_tensors="pt",
        max_length=CONFIG["max_src_len"],
        truncation=True,
    ).to(DEVICE)

    with torch.no_grad():
        out_ids = model.generate(
            **enc,
            num_beams=CONFIG["num_beams"],
            max_new_tokens=CONFIG["max_tgt_len"],
            length_penalty=CONFIG["length_penalty"],
            repetition_penalty=CONFIG["repetition_penalty"],
            no_repeat_ngram_size=CONFIG["no_repeat_ngram_size"],
            early_stopping=True,
        )

    pred = tokenizer.decode(out_ids[0], skip_special_tokens=True)
    return pred


def translate_with_postprocess(src_text: str) -> str:
    """Generate translation and apply full post-processing pipeline."""
    raw_translation = generate_translation(src_text)
    processed = postprocessor.postprocess(raw_translation, src_text)
    final = format_single_sentence(processed)
    return final

# %% [markdown]
# ## 8. Load Test Data

# %%
test_df = pd.read_csv(INPUT_DIR / "test.csv")
print(f"Test samples: {len(test_df)}")
print(f"Columns: {list(test_df.columns)}")
print(f"\nFirst few rows:")
print(test_df.head())

# %% [markdown]
# ## 9. Generate Predictions

# %%
print("\nGenerating predictions...")
predictions = []

for idx, row in test_df.iterrows():
    src_text = row["transliteration"]
    translation = translate_with_postprocess(src_text)

    predictions.append({
        "id": row["id"],
        "translation": translation
    })

    # Progress output
    if idx < 5 or idx % 10 == 0:
        print(f"[{idx+1}/{len(test_df)}] ID {row['id']}: {translation[:60]}...")

submission = pd.DataFrame(predictions)
print(f"\nGenerated {len(submission)} predictions")

# %% [markdown]
# ## 10. Save Submission

# %%
submission.to_csv("submission.csv", index=False)
print("Saved submission.csv")

print("\nSubmission preview:")
print(submission)

# Verify file
verify = pd.read_csv("submission.csv")
print(f"\nVerified: {len(verify)} rows, columns: {list(verify.columns)}")

# %% [markdown]
# ## 11. Example Translations
#
# Display a few example translations for inspection.

# %%
print("\n" + "="*80)
print("EXAMPLE TRANSLATIONS")
print("="*80)

for idx in range(min(4, len(test_df))):
    row = test_df.iloc[idx]
    src = row["transliteration"]
    translation = submission.iloc[idx]["translation"]

    print(f"\n--- Example {idx+1} (ID: {row['id']}) ---")
    print(f"Source: {src[:100]}{'...' if len(src) > 100 else ''}")
    print(f"Translation: {translation}")