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

# %%
"""
ByT5 Train-Set Inference

Kernel: byt5-trainset-inference-v1-0-0
Model: byt5-skeleton-akkadian

Generates ByT5 translations for the TRAINING set.
Output will be used to train the philologist refinement model.
"""

import os
import re

# Suppress progress bars and warnings
os.environ["TQDM_DISABLE"] = "1"
os.environ["HF_HUB_DISABLE_PROGRESS_BARS"] = "1"

# Monkey-patch huggingface_hub to allow local paths with multiple slashes
import huggingface_hub.utils._validators as hf_validators

_original_validate_repo_id = hf_validators.validate_repo_id


def _patched_validate_repo_id(repo_id: str) -> None:
    """Allow local paths (starting with /) to bypass repo_id validation."""
    if repo_id and repo_id.startswith("/"):
        return
    return _original_validate_repo_id(repo_id)


hf_validators.validate_repo_id = _patched_validate_repo_id

import pandas as pd
import torch
from transformers import AutoTokenizer, AutoModelForSeq2SeqLM

print(f"PyTorch: {torch.__version__}")
print(f"CUDA available: {torch.cuda.is_available()}")
if torch.cuda.is_available():
    print(f"GPU: {torch.cuda.get_device_name(0)}")
    torch.cuda.empty_cache()

DEVICE = torch.device("cuda" if torch.cuda.is_available() else "cpu")

# %%
# Configuration
CONFIG = {
    "model_path": "/kaggle/input/byt5-skeleton-akkadian/transformers/transformers/1",
    # Use train.csv for philologist training data generation
    "data_source": "/kaggle/input/deep-past-initiative-machine-translation/train.csv",
    "source_column": "transliteration",
    "target_column": "translation",  # Gold reference
    "id_column": "oare_id",
    "batch_size": 8,  # Larger batch for inference-only
    "max_source_length": 512,
    "max_target_length": 384,
    "num_beams": 4,
    "repetition_penalty": 1.2,
    "no_repeat_ngram_size": 3,
    "prefix": "translate Akkadian to English: ",
}
print(f"Config: {CONFIG}")

# %%
# Load Model
print(f"\nLoading model from: {CONFIG['model_path']}")
tokenizer = AutoTokenizer.from_pretrained(CONFIG["model_path"], local_files_only=True)
model = AutoModelForSeq2SeqLM.from_pretrained(CONFIG["model_path"], local_files_only=True).to(DEVICE)
model.eval()
print(f"Parameters: {sum(p.numel() for p in model.parameters()):,}")
print(f"Model on: {DEVICE}")

# %%
# Load Training Data
train_df = pd.read_csv(CONFIG["data_source"])
print(f"\nTraining samples: {len(train_df)}")
print(f"Columns: {list(train_df.columns)}")
print(train_df.head())


# %%
# Skeleton Preprocessing (must match training preprocessing exactly)
SUBSCRIPT_MAP = {
    "₀": "0",
    "₁": "1",
    "₂": "2",
    "₃": "3",
    "₄": "4",
    "₅": "5",
    "₆": "6",
    "₇": "7",
    "₈": "8",
    "₉": "9",
}
VOWEL_PATTERN = re.compile(r"[aeiuāēīūâêîûàèìùáéíúÀÈÌÙÁÉÍÚÂÊÎÛ]", re.IGNORECASE)
SUBSCRIPTS = "₀₁₂₃₄₅₆₇₈₉ₓ"


def normalize_subscripts(text: str) -> str:
    for subscript, digit in SUBSCRIPT_MAP.items():
        text = text.replace(subscript, digit)
    return text


def normalize_gaps(text: str) -> str:
    text = re.sub(r"(?:x\s+){2,}x", "<big_gap>", text)
    text = re.sub(r"\bx\b", "<gap>", text)
    text = re.sub(r"(?:<gap>\s*)+", "<gap> ", text)
    text = re.sub(r"(?:<big_gap>\s*)+", "<big_gap> ", text)
    return text


def preprocess_source(text: str) -> str:
    text = str(text)
    text = normalize_subscripts(text)
    text = normalize_gaps(text)
    text = re.sub(r"\s+", " ", text).strip()
    return text


def is_logogram(token):
    clean = "".join(c for c in token if c not in ".₀₁₂₃₄₅₆₇₈₉ₓ")
    return clean.isupper() and len(clean) > 0


def is_number(token):
    try:
        float(token.replace(",", "."))
        return True
    except ValueError:
        return bool(re.match(r"^[\d.,/]+$", token))


def strip_determinatives(token):
    return re.sub(r"^\([a-zA-Z]{1,3}\)", "", token)


def to_skeleton(text):
    """Convert text to consonantal skeleton."""
    if not text:
        return ""
    result = []
    for token in str(text).split():
        if is_logogram(token) or is_number(token):
            result.append(token)
            continue
        normalized = strip_determinatives(token)
        normalized = "".join(c for c in normalized if c not in SUBSCRIPTS)
        skeleton = VOWEL_PATTERN.sub("", normalized)
        skeleton = re.sub(r"-+", "", skeleton)
        result.append(skeleton if skeleton else "_")
    return " ".join(result)


def create_skeleton_input(text: str) -> str:
    """Create input with skeleton prefix: [SKL:skeleton] original_text"""
    if not text:
        return ""
    processed = preprocess_source(text)
    skeleton = to_skeleton(processed)
    return f"[SKL:{skeleton}] {processed}"


def clean_translation(text):
    """Clean up translation output."""
    text = re.sub(r"\s*\(\s*s\s*\)\s*", "s ", text)
    text = re.sub(r"\s*\(\s*(or|and)\s*\)\s*", r" \1 ", text)
    text = re.sub(r"\s*\(\s*(ki|d|m|f|lu|munus)\s*\)\s*", " ", text)
    text = " ".join(text.split())
    return text.strip()


# %%
# Generate Translations for Training Set
print("\n" + "=" * 60)
print("GENERATING TRANSLATIONS FOR TRAINING SET")
print("=" * 60)

texts = train_df[CONFIG["source_column"]].tolist()
gold_refs = train_df[CONFIG["target_column"]].tolist()
ids = train_df[CONFIG["id_column"]].tolist()

# Apply skeleton preprocessing and prefix
prefix = CONFIG["prefix"]
prefixed_texts = [prefix + create_skeleton_input(t) for t in texts]

print(f"Total samples: {len(texts)}")
print(f"Sample input: {prefixed_texts[0][:150]}...")

translations = []
batch_size = CONFIG["batch_size"]
total_batches = (len(prefixed_texts) + batch_size - 1) // batch_size

for i in range(0, len(prefixed_texts), batch_size):
    batch = prefixed_texts[i : i + batch_size]

    inputs = tokenizer(
        batch,
        return_tensors="pt",
        padding=True,
        truncation=True,
        max_length=CONFIG["max_source_length"],
    ).to(DEVICE)

    with torch.no_grad():
        outputs = model.generate(
            **inputs,
            num_beams=CONFIG["num_beams"],
            max_new_tokens=CONFIG["max_target_length"],
            early_stopping=True,
            repetition_penalty=CONFIG["repetition_penalty"],
            no_repeat_ngram_size=CONFIG["no_repeat_ngram_size"],
        )

    decoded = tokenizer.batch_decode(outputs, skip_special_tokens=True)
    cleaned = [clean_translation(t) for t in decoded]
    translations.extend(cleaned)

    batch_num = i // batch_size + 1
    if batch_num % 50 == 0 or batch_num == total_batches:
        print(f"  Batch {batch_num}/{total_batches} ({100 * batch_num / total_batches:.1f}%)")

print(f"\nGenerated {len(translations)} translations")

# %%
# Sample Comparisons (ByT5 vs Gold)
print("\n" + "=" * 60)
print("SAMPLE COMPARISONS (ByT5 vs Gold)")
print("=" * 60)
for i in range(min(5, len(texts))):
    print(f"\n[{i + 1}] Source: {texts[i][:80]}...")
    print(f"    ByT5:  {translations[i][:100]}...")
    print(f"    Gold:  {gold_refs[i][:100]}...")

# %%
# Create Output for Philologist Training
# This CSV contains: id, source (transliteration), byt5_translation, gold_translation
output_df = pd.DataFrame(
    {
        "id": ids,
        "source": texts,
        "byt5_translation": translations,
        "gold_translation": gold_refs,
    }
)

print(f"\nOutput shape: {output_df.shape}")
print(output_df.head())

# Save for philologist training
output_df.to_csv("byt5_trainset_translations.csv", index=False)
print(f"\nSaved byt5_trainset_translations.csv")
print(f"File size: {os.path.getsize('byt5_trainset_translations.csv'):,} bytes")

# Also save just the translations for quick reference
translations_only = pd.DataFrame(
    {
        "id": ids,
        "translation": translations,
    }
)
translations_only.to_csv("submission.csv", index=False)
print(f"Saved submission.csv (translations only)")

# %%
print("\n" + "=" * 60)
print("TRAIN-SET INFERENCE COMPLETE")
print("=" * 60)
print(f"Output files:")
print(f"  - byt5_trainset_translations.csv (full data for philologist)")
print(f"  - submission.csv (translations only)")