#!/usr/bin/env python3
"""
ByT5 Skeleton Inference v1.0.2

Uses skeleton-trained model with consonantal prefix format.
Input: "[SKL:šrm dnm] šar-rum dan-num" -> "the mighty king"

v1.0.1: Added repetition_penalty and no_repeat_ngram_size to fix repetition issues.
v1.0.2: Increased max_source_length to 512, max_target_length to 384 per jeanjean111 feedback.
"""

import os

# Disable progress bars for cleaner Kaggle logs
os.environ["TQDM_DISABLE"] = "1"
os.environ["HF_HUB_DISABLE_PROGRESS_BARS"] = "1"

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 - skeleton-trained checkpoint
    "model_path": "/kaggle/input/byt5-skeleton-akkadian/transformers/transformers/1",
    # Inference settings - increased lengths per jeanjean111 feedback
    "batch_size": 2,  # Reduced to fit longer sequences in memory
    "max_source_length": 512,  # Was 256
    "max_target_length": 384,  # Was 192
    # Generation - beam search with repetition control
    "num_beams": 4,
    "early_stopping": True,
    "repetition_penalty": 1.2,
    "no_repeat_ngram_size": 3,
    "prefix": "translate Akkadian to English: ",
}

# ============================================================================
# Skeleton Preprocessing (from training)
# ============================================================================
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.

    Format: [SKL:consonantal_skeleton] original_text
    """
    if not text:
        return ""
    processed = preprocess_source(text)
    skeleton = to_skeleton(processed)
    return f"[SKL:{skeleton}] {processed}"


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

        # Pre-process with skeleton format
        print("Creating skeleton inputs...")
        self.processed = []
        for t in tqdm(self.texts):
            skeleton_input = create_skeleton_input(t)
            self.processed.append(CONFIG["prefix"] + skeleton_input)

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

    def __getitem__(self, idx):
        text = self.processed[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 Skeleton Inference v1.0.2")
    print("Using consonantal skeleton prefix format")
    print("With repetition penalty + extended context lengths")
    print("=" * 60)

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

    # Load model
    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()
    print(f"Parameters: {model.num_parameters():,}")

    # 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 = SkeletonInferenceDataset(
        test_df, tokenizer, CONFIG["max_source_length"]
    )
    test_loader = DataLoader(
        test_dataset, batch_size=CONFIG["batch_size"], shuffle=False
    )

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

    # Show generation config
    print(f"\nGeneration config:")
    print(f"  max_source_length: {CONFIG['max_source_length']}")
    print(f"  max_target_length: {CONFIG['max_target_length']}")
    print(f"  batch_size: {CONFIG['batch_size']}")
    print(f"  num_beams: {CONFIG['num_beams']}")
    print(f"  repetition_penalty: {CONFIG['repetition_penalty']}")
    print(f"  no_repeat_ngram_size: {CONFIG['no_repeat_ngram_size']}")

    # 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"],
                repetition_penalty=CONFIG["repetition_penalty"],
                no_repeat_ngram_size=CONFIG["no_repeat_ngram_size"],
            )

            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})

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

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

    print(f"\nStats:")
    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()