# ---
# jupyter:
#   kernelspec:
#     display_name: Python 3
#     language: python
#     name: python3
#   kaggle:
#     accelerator: gpu
#     dataSources:
#       - type: dataset
#         name: manwithacat/oracc-akkadian-english-parallel-corpus
#       - type: competition
#         name: deep-past-initiative-machine-translation
#     docker_image: gcr.io/kaggle-gpu-images/python
#     isGpuEnabled: true
#     isInternetEnabled: true
#     language: python
#     sourceType: script
# ---

# %% [markdown]
# # byt5-akkadian-combined
#
# ByT5-base fine-tuned on competition + ORACC data
#
# **Generated from**: `byt5-training.toml`
# **Version**: 1.0.7
# **Generated at**: 2025-12-24T03:18:10.075Z

# %%
import os
import sys
import time
import json
import glob
import subprocess
import shutil

os.environ["TOKENIZERS_PARALLELISM"] = "false"
os.environ["PYTORCH_CUDA_ALLOC_CONF"] = "expandable_segments:True"
os.environ["TQDM_DISABLE"] = "1"  # Disable progress bars for clean Kaggle logs
os.environ["HF_HUB_DISABLE_PROGRESS_BARS"] = "1"  # Disable HuggingFace download bars

import torch
import pandas as pd
import numpy as np
from datetime import datetime
import warnings

# Suppress internal transformers deprecation warnings (past_key_values tuple format)
# These are library-internal and will be fixed in future transformers versions
warnings.filterwarnings("ignore", message=".*past_key_values.*")
warnings.filterwarnings("ignore", category=FutureWarning, module="transformers")

if torch.cuda.is_available():
    torch.cuda.empty_cache()

print(f"Python: {sys.version}")
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)}")
    print(f"GPU Memory: {torch.cuda.get_device_properties(0).total_memory / 1e9:.1f} GB")

# %% [markdown]
# ## Install Dependencies

# %%
# sacrebleu for chrF metric - use directly instead of evaluate.load() to avoid internet download
packages = ["sentencepiece", "sacrebleu"]
# Note: Don't use --no-deps as it skips transitive dependencies (sacrebleu needs portalocker)
subprocess.run([sys.executable, "-m", "pip", "install", "-q"] + packages, check=True)
subprocess.run([sys.executable, "-m", "pip", "install", "-q", "transformers>=4.35.0", "accelerate"], check=True)
print("Packages installed")

# %%
from datasets import Dataset
from transformers import (
    AutoTokenizer,
    AutoModelForSeq2SeqLM,
    Seq2SeqTrainer,
    Seq2SeqTrainingArguments,
    DataCollatorForSeq2Seq,
    EarlyStoppingCallback,
)
import sacrebleu  # Direct usage instead of evaluate.load() - works offline

# %% [markdown]
# ## Configuration
#
# Auto-generated from `byt5-training.toml` - DO NOT EDIT MANUALLY
# To modify, edit the TOML file and regenerate with: akk notebook build byt5-training.toml

# %%
# === AUTO-GENERATED CONFIG ===
# Source: byt5-training.toml
# Generated: 2025-12-24T03:18:10.075Z
# Version: 1.0.7

CONFIG = {
    # Model
    "model_name": "google/byt5-base",

    # Language codes
    "src_lang": "akk",
    "tgt_lang": "eng",

    # Precision
    "training_precision": "fp32",
    "checkpoint_precision": "fp32",
    "fp16": False,

    # Training
    "num_epochs": 8,
    "batch_size": 2,
    "gradient_accumulation_steps": 4,
    "learning_rate": 0.0002,
    "weight_decay": 0.01,
    "warmup_ratio": 0.06,
    "lr_scheduler_type": "cosine",

    # Sequence lengths
    "max_src_len": 512,
    "max_tgt_len": 256,

    # Evaluation
    "eval_steps": 200,
    "save_steps": 200,
    "logging_steps": 50,
    "predict_with_generate": True,
    "metric_for_best_model": "chrf",
    "greater_is_better": True,

    # Checkpoints
    "save_total_limit": 1,
    "save_only_model": True,  # True = skip optimizer.pt (saves ~4GB per checkpoint)
    "load_best_at_end": True,

    # Early stopping
    "early_stopping_enabled": True,
    "early_stopping_patience": 3,

    # Validation split
    "val_split": 0.1,
    "seed": 42,

    # Generation
    "num_beams": 4,
    "max_new_tokens": 256,
    "repetition_penalty": 1,
    "no_repeat_ngram_size": 0,

    # Platform
    "clear_hf_cache": True,

    # Output
    "output_dir": "./byt5-akkadian",
}

print("Configuration:")
for k, v in CONFIG.items():
    print(f"  {k}: {v}")

# %% [markdown]
# ## Load Dataset

# %%
# Dataset sources (in priority order)
DATASET_SOURCES = [
    "/kaggle/input/oracc-akkadian-english-parallel-corpus/*.csv",
    "/kaggle/input/deep-past-initiative-machine-translation/*.csv",
]

train_df = None
dataset_type = None

for source in DATASET_SOURCES:
    files = glob.glob(source, recursive=True)
    if files:
        print(f"Using dataset: {files[0]}")
        train_df = pd.read_csv(files[0])
        dataset_type = source.split("/")[-2] if "/" in source else "local"
        break

if train_df is None:
    raise RuntimeError(f"No training data found! Searched: {DATASET_SOURCES}")

print(f"\nTraining samples: {len(train_df)}")
print(f"Columns: {train_df.columns.tolist()}")

# %%
# Data quality check
print(f"\nData quality:")
print(f"  Null sources: {train_df['akkadian'].isna().sum()}")
print(f"  Null targets: {train_df['english'].isna().sum()}")
print(f"  Avg source length: {train_df['akkadian'].str.len().mean():.1f} chars")
print(f"  Avg target length: {train_df['english'].str.len().mean():.1f} chars")

train_df = train_df.dropna(subset=['akkadian', 'english'])
print(f"  After cleaning: {len(train_df)} samples")

# %% [markdown]
# ## Initialize Model

# %%
print(f"\nLoading model: {CONFIG['model_name']}")
start = time.time()

tokenizer = AutoTokenizer.from_pretrained(CONFIG["model_name"])
model = AutoModelForSeq2SeqLM.from_pretrained(CONFIG["model_name"])

print(f"Model loaded in {time.time()-start:.1f}s")
print(f"Parameters: {sum(p.numel() for p in model.parameters()):,}")

# %%
# Clear HuggingFace cache to save disk space
if CONFIG["clear_hf_cache"]:
    hf_cache = os.path.expanduser("~/.cache/huggingface")
    if os.path.exists(hf_cache):
        cache_size = sum(os.path.getsize(os.path.join(dp, f)) for dp, dn, fn in os.walk(hf_cache) for f in fn)
        print(f"Clearing HF cache: {cache_size / 1e9:.2f} GB")
        shutil.rmtree(hf_cache, ignore_errors=True)
        print("HF cache cleared to save disk space")

# %%
# T5/ByT5 uses task prefix for translation
TASK_PREFIX = "translate Akkadian to English: "
print(f"Using T5 task prefix: {TASK_PREFIX}")

# %% [markdown]
# ## Preprocessing

# %%
def preprocess_function(examples):
    # T5/ByT5: prepend task prefix to inputs
    inputs = [TASK_PREFIX + str(x) for x in examples["akkadian"]]
    targets = [str(x) for x in examples["english"]]

    model_inputs = tokenizer(
        inputs,
        max_length=CONFIG["max_src_len"],
        truncation=True,
        padding=False,
    )

    # Tokenize targets
    labels = tokenizer(
        targets,
        max_length=CONFIG["max_tgt_len"],
        truncation=True,
        padding=False,
    )

    model_inputs["labels"] = labels["input_ids"]
    return model_inputs

# %%
print("\nPreparing datasets...")

dataset = Dataset.from_pandas(train_df[["akkadian", "english"]])
split = dataset.train_test_split(test_size=CONFIG["val_split"], seed=CONFIG["seed"])
train_dataset = split["train"]
val_dataset = split["test"]

print(f"Train: {len(train_dataset)}, Val: {len(val_dataset)}")

# %%
train_tokenized = train_dataset.map(preprocess_function, batched=True, remove_columns=train_dataset.column_names)
val_tokenized = val_dataset.map(preprocess_function, batched=True, remove_columns=val_dataset.column_names)

print(f"Tokenized - Train: {len(train_tokenized)}, Val: {len(val_tokenized)}")

# %%
# Compute metrics using sacrebleu directly (no internet required)
def compute_metrics(eval_preds):
    preds, labels = eval_preds
    # Replace -100 (ignored tokens) with pad token for decoding
    labels = np.where(labels != -100, labels, tokenizer.pad_token_id)
    decoded_preds = tokenizer.batch_decode(preds, skip_special_tokens=True)
    decoded_labels = tokenizer.batch_decode(labels, skip_special_tokens=True)
    # Strip whitespace
    decoded_preds = [pred.strip() for pred in decoded_preds]
    decoded_labels = [label.strip() for label in decoded_labels]
    # Compute chrF using sacrebleu directly
    chrf = sacrebleu.corpus_chrf(decoded_preds, [decoded_labels])
    return {"chrf": chrf.score}

# %% [markdown]
# ## Training Setup

# %%
data_collator = DataCollatorForSeq2Seq(tokenizer=tokenizer, model=model, padding=True)

training_args = Seq2SeqTrainingArguments(
    output_dir=CONFIG["output_dir"],
    num_train_epochs=CONFIG["num_epochs"],
    per_device_train_batch_size=CONFIG["batch_size"],
    per_device_eval_batch_size=CONFIG["batch_size"],
    gradient_accumulation_steps=CONFIG["gradient_accumulation_steps"],
    learning_rate=CONFIG["learning_rate"],
    weight_decay=CONFIG["weight_decay"],
    warmup_ratio=CONFIG["warmup_ratio"],
    lr_scheduler_type=CONFIG.get("lr_scheduler_type", "linear"),
    fp16=CONFIG["fp16"],
    eval_strategy="steps",  # Modern API (evaluation_strategy is deprecated)
    eval_steps=CONFIG["eval_steps"],
    save_strategy="steps",
    save_steps=CONFIG["save_steps"],
    save_total_limit=CONFIG["save_total_limit"],
    save_only_model=CONFIG["save_only_model"],  # Skip optimizer.pt to save disk space
    logging_steps=CONFIG["logging_steps"],
    load_best_model_at_end=CONFIG["load_best_at_end"],
    metric_for_best_model=CONFIG["metric_for_best_model"],
    greater_is_better=CONFIG["greater_is_better"],
    predict_with_generate=CONFIG["predict_with_generate"],
    report_to="none",
    seed=CONFIG["seed"],
)

callbacks = []
if CONFIG["early_stopping_enabled"]:
    callbacks.append(EarlyStoppingCallback(early_stopping_patience=CONFIG["early_stopping_patience"]))

trainer = Seq2SeqTrainer(
    model=model,
    args=training_args,
    train_dataset=train_tokenized,
    eval_dataset=val_tokenized,
    processing_class=tokenizer,  # Modern API (tokenizer= is deprecated in transformers>=4.46)
    data_collator=data_collator,
    compute_metrics=compute_metrics,
    callbacks=callbacks,
)

# %% [markdown]
# ## Train

# %%
print("\n" + "="*60)
print("STARTING TRAINING")
print(f"Config: byt5-akkadian-combined v1.0.7")
print(f"Dataset: {dataset_type} ({len(train_tokenized)} samples)")
print("="*60 + "\n")

train_start = time.time()
train_result = trainer.train()
train_time = time.time() - train_start

print(f"\nTraining completed in {train_time/60:.1f} minutes")
print(f"Final metrics: {train_result.metrics}")

# %% [markdown]
# ## Save Model

# %%
output_dir = CONFIG["output_dir"]
trainer.save_model(output_dir)
tokenizer.save_pretrained(output_dir)

# Save training config
config_path = os.path.join(output_dir, "training_config.json")
with open(config_path, "w") as f:
    json.dump({
        **{k: v for k, v in CONFIG.items() if not callable(v)},
        "dataset_type": dataset_type,
        "train_samples": len(train_tokenized),
        "val_samples": len(val_tokenized),
        "train_time_minutes": train_time / 60,
        "final_loss": train_result.metrics.get("train_loss"),
        "_meta": {
            "name": "byt5-akkadian-combined",
            "version": "1.0.7",
            "source_config": "byt5-training.toml",
            "generated_at": "2025-12-24T03:18:10.075Z",
        },
    }, f, indent=2)

# Save length config for inference
length_config = {
    "length_ratio": 1.5,
    "length_slack": 20,
    "global_max": 512,
    "global_min": 32,
}
with open(os.path.join(output_dir, "length_config.json"), "w") as f:
    json.dump(length_config, f)

print(f"\nModel saved to: {output_dir}")
print(f"Files: {os.listdir(output_dir)}")

# %% [markdown]
# ## Upload to Kaggle Model Registry

# %%
# Install kagglehub for model registry upload
subprocess.run([sys.executable, "-m", "pip", "install", "-q", "kagglehub>=0.2.5"], check=True)
import kagglehub

# Upload model to Kaggle registry
# Model version is auto-derived from notebook version (1.0.7)
model_version = "v1-0-7"
kaggle_handle = f"manwithacat/akkadian-byt5/transformers/{model_version}"
print(f"\nUploading model to Kaggle: {kaggle_handle}")

try:
    kagglehub.model_upload(
        handle=kaggle_handle,
        local_model_dir=output_dir,
        version_notes=f"byt5-akkadian-combined v1.0.7 - trained on {dataset_type}",
        license_name="Apache 2.0",
    )
    print(f"Model uploaded successfully to: kaggle.com/models/{kaggle_handle.rsplit('/', 2)[0]}")
except Exception as e:
    print(f"Warning: Model upload failed: {e}")
    print("Model is still saved locally and can be uploaded manually")

# %% [markdown]
# ## Quick Evaluation

# %%
print("\n" + "="*60)
print("SAMPLE TRANSLATIONS")
print("="*60)

model.eval()
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
model = model.to(device)

test_samples = val_dataset.select(range(min(5, len(val_dataset))))

for i, sample in enumerate(test_samples):
    src = sample["akkadian"]
    ref = sample["english"]

    # T5: prepend task prefix
    inputs = tokenizer(TASK_PREFIX + src, return_tensors="pt", max_length=CONFIG["max_src_len"], truncation=True)
    inputs = {k: v.to(device) for k, v in inputs.items()}

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

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

    print(f"\n[{i}] SRC: {src[:80]}...")
    print(f"    REF: {ref[:80]}...")
    print(f"    OUT: {pred[:80]}...")

# %% [markdown]
# ## Generate Competition Submission

# %%
import re

def clean_translation(text):
    """
    Clean ORACC academic conventions from translation output.
    These artifacts hurt scoring against clean English references.
    """
    # Remove indeterminate plural markers: "( s )" or "(s)"
    text = re.sub(r'\s*\(\s*s\s*\)', 's', text)

    # Remove parenthesized alternatives: "( or another )" -> "or another"
    text = re.sub(r'\s*\(\s*(or|and)\s+([^)]+)\s*\)\s*', r' \1 \2 ', text)

    # Remove parenthesized conjunctions alone: "( or )" "( and )" etc.
    text = re.sub(r'\s*\(\s*(or|and|i\.e\.|etc\.)\s*\)\s*', r' \1 ', text)

    # Remove Akkadian determinatives: (ki), (m), (f), (d), etc.
    text = re.sub(r'\s*\(\s*(ki|m|f|d|pl|sg|lú|giš|kur|uru)\s*\)', '', text, flags=re.IGNORECASE)

    # Remove other short parentheticals that look like annotations
    text = re.sub(r'\s*\(\s*[a-z]{1,3}\s*\)', '', text)

    # Normalize multiple spaces
    text = re.sub(r'\s+', ' ', text)

    # Fix spacing around punctuation
    text = re.sub(r'\s+([.,;:!?])', r'\1', text)
    text = re.sub(r'([.,;:!?])(?=[A-Za-z])', r'\1 ', text)

    return text.strip()

# %%
# Load test data
test_files = glob.glob("/kaggle/input/deep-past-initiative-machine-translation/test.csv", recursive=True)
print(f"Found test files: {test_files}")

if not test_files:
    print("WARNING: test.csv not found - skipping submission generation")
    test_df = None
else:
    test_df = pd.read_csv(test_files[0])
    print(f"Test samples: {len(test_df)}")
    print(test_df.head())

# %%
# Define batch translation function
def translate_batch(texts):
    """Translate a batch of texts."""
    # T5: prepend task prefix
    inputs = tokenizer(
        [TASK_PREFIX + t for t in texts],
        max_length=CONFIG["max_src_len"],
        truncation=True,
        padding=True,
        return_tensors="pt",
    )
    inputs = {k: v.to(device) for k, v in inputs.items()}

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

    return tokenizer.batch_decode(outputs, skip_special_tokens=True)

# %%
if test_df is not None:
    # Detect source column
    source_col = "transliteration" if "transliteration" in test_df.columns else "source" if "source" in test_df.columns else test_df.columns[1]
    sources = test_df[source_col].tolist()

    print(f"\nTranslating {len(sources)} test samples using column: {source_col}")

    translations = []
    batch_size = 16
    for i in range(0, len(sources), batch_size):
        batch = sources[i:i+batch_size]
        batch_translations = translate_batch(batch)
        translations.extend(batch_translations)
        if (i + batch_size) % 100 == 0 or i + batch_size >= len(sources):
            print(f"  Translated {min(i + batch_size, len(sources))}/{len(sources)}")

    print(f"Generated {len(translations)} translations")

    # Clean translations (remove ORACC artifacts that hurt scoring)
    translations = [clean_translation(t) for t in translations]
    print("Cleaned ORACC artifacts from translations")

    # Create submission
    id_col = "id" if "id" in test_df.columns else "id"
    submission = pd.DataFrame({
        id_col: test_df[id_col] if id_col in test_df.columns else range(len(translations)),
        "translation": translations
    })

    print("\nSubmission preview:")
    print(submission.head(10))

    submission.to_csv("submission.csv", index=False)
    print("\nSaved submission.csv")

    # Show sample translations
    print("\n=== Sample Test Translations ===")
    for i in range(min(5, len(sources))):
        print(f"\n[{i}]")
        print(f"SRC: {sources[i][:100]}...")
        print(f"OUT: {translations[i][:100]}...")

# %%
print("\n" + "="*60)
print("TRAINING COMPLETE")
print(f"Config: byt5-akkadian-combined v1.0.7")
print(f"Dataset: {dataset_type}")
print(f"Samples: {len(train_tokenized)} train, {len(val_tokenized)} val")
print(f"Time: {train_time/60:.1f} minutes")
print(f"Model saved to: {output_dir}")
print("="*60)