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

# %%
"""
Chunky v1.2.0 — S1: Bare competitor reproduction

byt5-akkadian-optimized-34x model, no chunking, matching competitor config.
Purpose: reproduce their 35.0 LB baseline.
"""

import os
import re

os.environ["TQDM_DISABLE"] = "1"
os.environ["HF_HUB_DISABLE_PROGRESS_BARS"] = "1"

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

DEVICE = torch.device("cuda" if torch.cuda.is_available() else "cpu")
print(
    f"PyTorch: {torch.__version__}, CUDA: {torch.cuda.is_available()}, Device: {DEVICE}"
)
if torch.cuda.is_available():
    print(f"GPU: {torch.cuda.get_device_name(0)}")

# %%
# Config — matches competitor's settings exactly
MODEL_PATH = "/kaggle/input/final-byt5/byt5-akkadian-optimized-34x"
TEST_PATH = "/kaggle/input/deep-past-initiative-machine-translation/test.csv"
PREFIX = "translate Akkadian to English: "

NUM_BEAMS = 8
LENGTH_PENALTY = 1.09
MAX_NEW_TOKENS = 512
MAX_SOURCE_LENGTH = 512

# %%
# Preprocessing — match competitor's gap normalization
PAT_BIG_GAP = re.compile(r"(\.{3,}|…+|……)")
PAT_SMALL_GAP = re.compile(r"(xx+|\s+x\s+)")


def preprocess(text):
    if pd.isna(text):
        return ""
    text = str(text)
    text = PAT_BIG_GAP.sub("<big_gap>", text)
    text = PAT_SMALL_GAP.sub("<gap>", text)
    return text


# %%
# Postprocessing — minimal, match competitor
PAT_REPEATED_WORDS = re.compile(r"\b(\w+)(?:\s+\1\b)+")
PAT_WHITESPACE = re.compile(r"\s+")


def postprocess(text):
    if not isinstance(text, str) or not text.strip():
        return ""
    # Remove repeated adjacent words
    text = PAT_REPEATED_WORDS.sub(r"\1", text)
    # Remove repeated n-grams (4 down to 2)
    for n in range(4, 1, -1):
        pattern = r"\b((?:\w+\s+){" + str(n - 1) + r"}\w+)(?:\s+\1\b)+"
        text = re.sub(pattern, r"\1", text)
    text = PAT_WHITESPACE.sub(" ", text).strip()
    return text


# %%
# Load model
print(f"\nLoading model from: {MODEL_PATH}")
tokenizer = AutoTokenizer.from_pretrained(MODEL_PATH)
model = AutoModelForSeq2SeqLM.from_pretrained(MODEL_PATH).to(DEVICE)
model.eval()
print(f"Parameters: {sum(p.numel() for p in model.parameters()):,}")

# %%
# Load test data
test_df = pd.read_csv(TEST_PATH)
print(f"\nTest samples: {len(test_df)}")
print(test_df.head())

# %%
# Translate
print("\n" + "=" * 60)
print("TRANSLATING")
print("=" * 60)

translations = []
for i, row in test_df.iterrows():
    text = preprocess(row["transliteration"])
    source = PREFIX + text

    inputs = tokenizer(
        source,
        return_tensors="pt",
        max_length=MAX_SOURCE_LENGTH,
        truncation=True,
    ).to(DEVICE)

    with torch.no_grad():
        outputs = model.generate(
            input_ids=inputs.input_ids,
            attention_mask=inputs.attention_mask,
            num_beams=NUM_BEAMS,
            max_new_tokens=MAX_NEW_TOKENS,
            length_penalty=LENGTH_PENALTY,
            early_stopping=True,
            no_repeat_ngram_size=0,
        )

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

    if not translation or len(translation.strip()) < 2:
        translation = "broken text"

    translations.append(translation)
    print(f"  [{i}] {translation[:80]}...")

# %%
# Save submission
submission = pd.DataFrame({"id": test_df["id"], "translation": translations})
print(f"\nSubmission shape: {submission.shape}")
print(submission)

submission.to_csv("submission.csv", index=False)
print(f"\nSaved submission.csv ({os.path.getsize('submission.csv'):,} bytes)")