# ---
# jupyter:
#   jupytext:
#     text_representation:
#       extension: .py
#       format_name: percent
#   kernelspec:
#     display_name: Python 3
#     language: python
#     name: python3
#   kaggle:
#     accelerator: gpu
#     dataSources:
#       - 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]
# # NLLB-600M Training on Kaggle P100
#
# This notebook trains NLLB-200-distilled-600M on Akkadian-English translation.
#
# **Target**: Replicate v1 config that achieved 23.1 BLEU
#
# **Config from successful v1 run**:
# - Model: facebook/nllb-200-distilled-600M
# - Language code: zho_Hans (Chinese placeholder since Akkadian not in NLLB)
# - Epochs: 10
# - Batch size: 8, Gradient accumulation: 2 (effective batch 16)
# - Learning rate: 5e-5
# - Weight decay: 0.01, Warmup ratio: 0.1
# - predict_with_generate: False
# - metric_for_best_model: eval_loss
# - Max length: 256, fp16: True

# %%
import os
import sys

# Set environment variables before importing other libraries
os.environ["TOKENIZERS_PARALLELISM"] = "false"
os.environ["PYTORCH_CUDA_ALLOC_CONF"] = "expandable_segments:True"

import time
import json
import torch
import pandas as pd
import numpy as np
from datetime import datetime

# Clear any cached memory
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

# %%
# Install required packages - use --no-deps to avoid pyarrow conflicts
import subprocess

# Only install what's not already in Kaggle, avoid upgrading pyarrow
packages = [
    "sentencepiece",
    "sacrebleu",
    "evaluate",
]
subprocess.run([sys.executable, "-m", "pip", "install", "-q", "--no-deps"] + packages, check=True)

# These are already in Kaggle but we need specific versions
subprocess.run([sys.executable, "-m", "pip", "install", "-q",
                "transformers>=4.35.0", "accelerate"], check=True)

print("Packages installed successfully")

# %%
from datasets import Dataset
from transformers import (
    AutoTokenizer,
    AutoModelForSeq2SeqLM,
    Seq2SeqTrainer,
    Seq2SeqTrainingArguments,
    DataCollatorForSeq2Seq,
    EarlyStoppingCallback,
)
import evaluate

# %% [markdown]
# ## Configuration
#
# Hyperparameters based on successful v1 run (23.1 BLEU)

# %%
# =============================================================================
# CONFIGURATION - Modify these to experiment
# =============================================================================

CONFIG = {
    # Model
    "model_name": "facebook/nllb-200-distilled-600M",

    # Language codes (zho_Hans worked best in v1 - achieved 23.1 BLEU)
    # Options: "zho_Hans" (Chinese), "arb_Arab" (Arabic), "heb_Hebr" (Hebrew)
    "src_lang": "zho_Hans",  # Source language placeholder for Akkadian
    "tgt_lang": "eng_Latn",  # English

    # Training - reduced batch for P100 16GB memory
    "num_epochs": 10,
    "batch_size": 2,  # Reduced from 8 for P100 memory
    "gradient_accumulation_steps": 8,  # Effective batch = 16 (same as before)
    "learning_rate": 5e-5,
    "weight_decay": 0.01,
    "warmup_ratio": 0.1,

    # Sequence lengths - reduced for memory
    "max_src_len": 192,
    "max_tgt_len": 192,

    # Evaluation
    "eval_steps": 100,
    "save_steps": 100,
    "logging_steps": 25,

    # Optimization
    "fp16": True,
    "predict_with_generate": False,  # Faster training, use eval_loss for selection
    "metric_for_best_model": "eval_loss",
    "greater_is_better": False,

    # Early stopping
    "early_stopping_patience": 5,  # Stop if no improvement for 5 evals

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

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

# %% [markdown]
# ## Load Data

# %%
import glob

# Find training data
train_files = glob.glob("/kaggle/input/**/train.csv", recursive=True)
print(f"Found train files: {train_files}")

if not train_files:
    raise RuntimeError("train.csv not found! Add competition as data source.")

train_df = pd.read_csv(train_files[0])
print(f"\nTraining samples: {len(train_df)}")
print(f"Columns: {train_df.columns.tolist()}")
print(train_df.head(2))

# %%
# Check data quality
print(f"\nData quality check:")
print(f"  Null transliterations: {train_df['transliteration'].isna().sum()}")
print(f"  Null translations: {train_df['translation'].isna().sum()}")
print(f"  Avg source length: {train_df['transliteration'].str.len().mean():.1f} chars")
print(f"  Avg target length: {train_df['translation'].str.len().mean():.1f} chars")

# Remove any null rows
train_df = train_df.dropna(subset=['transliteration', 'translation'])
print(f"  After cleaning: {len(train_df)} samples")

# %% [markdown]
# ## Initialize Model & Tokenizer

# %%
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()):,}")

# %%
# Set language codes
tokenizer.src_lang = CONFIG["src_lang"]
tokenizer.tgt_lang = CONFIG["tgt_lang"]

print(f"\nLanguage configuration:")
print(f"  Source: {CONFIG['src_lang']} (placeholder for Akkadian)")
print(f"  Target: {CONFIG['tgt_lang']} (English)")

# Get forced_bos_token_id for English output
eng_token_id = tokenizer.convert_tokens_to_ids(CONFIG["tgt_lang"])
print(f"  English token ID: {eng_token_id}")

# %% [markdown]
# ## Preprocessing

# %%
def preprocess_function(examples):
    """Tokenize source and target for NLLB."""
    inputs = examples["transliteration"]
    targets = examples["translation"]

    # Tokenize inputs (Akkadian transliteration)
    model_inputs = tokenizer(
        inputs,
        max_length=CONFIG["max_src_len"],
        truncation=True,
        padding=False,  # Dynamic padding via DataCollator
    )

    # Tokenize targets (English translation)
    with tokenizer.as_target_tokenizer():
        labels = tokenizer(
            targets,
            max_length=CONFIG["max_tgt_len"],
            truncation=True,
            padding=False,
        )

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

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

dataset = Dataset.from_pandas(train_df[["transliteration", "translation"]])

# Train/val split
split = dataset.train_test_split(test_size=CONFIG["val_split"], seed=CONFIG["seed"])
train_dataset = split["train"]
val_dataset = split["test"]

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

# %%
# Tokenize
print("\nTokenizing...")
train_dataset = train_dataset.map(
    preprocess_function,
    batched=True,
    remove_columns=["transliteration", "translation"],
    desc="Tokenizing train"
)
val_dataset = val_dataset.map(
    preprocess_function,
    batched=True,
    remove_columns=["transliteration", "translation"],
    desc="Tokenizing val"
)

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

# %% [markdown]
# ## Training Setup

# %%
OUTPUT_DIR = "/kaggle/working/nllb-akkadian"
os.makedirs(OUTPUT_DIR, exist_ok=True)

# Clean up HuggingFace cache to save disk space
import shutil
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"HuggingFace cache size: {cache_size / 1e9:.2f} GB")
    # Don't delete - model is loaded from cache

# Estimate training time
steps_per_epoch = len(train_dataset) // (CONFIG["batch_size"] * CONFIG["gradient_accumulation_steps"])
total_steps = steps_per_epoch * CONFIG["num_epochs"]
print(f"\nTraining estimation:")
print(f"  Steps per epoch: {steps_per_epoch}")
print(f"  Total steps: {total_steps}")
print(f"  Eval checkpoints: ~{total_steps // CONFIG['eval_steps']}")

# %%
training_args = Seq2SeqTrainingArguments(
    output_dir=OUTPUT_DIR,

    # Training
    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"],

    # Optimizer
    learning_rate=CONFIG["learning_rate"],
    weight_decay=CONFIG["weight_decay"],
    warmup_ratio=CONFIG["warmup_ratio"],
    lr_scheduler_type="linear",

    # Evaluation & Saving - minimal disk usage for Kaggle
    eval_strategy="steps",
    eval_steps=CONFIG["eval_steps"],
    save_strategy="steps",
    save_steps=CONFIG["save_steps"] * 2,  # Save less frequently
    save_total_limit=1,  # Only keep best checkpoint to save disk
    load_best_model_at_end=True,
    metric_for_best_model=CONFIG["metric_for_best_model"],
    greater_is_better=CONFIG["greater_is_better"],

    # Generation during eval (disabled for speed)
    predict_with_generate=CONFIG["predict_with_generate"],

    # Performance
    fp16=CONFIG["fp16"],
    dataloader_num_workers=0,  # Reduce disk I/O

    # Logging
    logging_dir=f"{OUTPUT_DIR}/logs",
    logging_steps=CONFIG["logging_steps"],
    report_to="none",

    # Misc
    seed=CONFIG["seed"],
    push_to_hub=False,
)

print("\nTraining arguments configured")

# %%
# Data collator for dynamic padding
data_collator = DataCollatorForSeq2Seq(
    tokenizer=tokenizer,
    model=model,
    padding=True,
    label_pad_token_id=-100,
)

# %%
# Create trainer
trainer = Seq2SeqTrainer(
    model=model,
    args=training_args,
    train_dataset=train_dataset,
    eval_dataset=val_dataset,
    tokenizer=tokenizer,
    data_collator=data_collator,
    callbacks=[EarlyStoppingCallback(early_stopping_patience=CONFIG["early_stopping_patience"])],
)

print("Trainer initialized")

# %% [markdown]
# ## Train!

# %%
print("\n" + "="*60)
print("STARTING TRAINING")
print("="*60)

train_start = time.time()

# Train
train_result = trainer.train()

train_time = time.time() - train_start
print(f"\nTraining completed in {train_time/60:.1f} minutes")

# %%
# Log training metrics
print("\nTraining metrics:")
for key, value in train_result.metrics.items():
    print(f"  {key}: {value}")

# %% [markdown]
# ## Final Evaluation

# %%
# Evaluate best model
print("\nEvaluating best model...")
eval_results = trainer.evaluate()

print("\nEvaluation results:")
for key, value in eval_results.items():
    print(f"  {key}: {value}")

# %% [markdown]
# ## Compute BLEU Score

# %%
# Now compute BLEU with generation
print("\nComputing BLEU score (this takes longer)...")

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

# Get a sample of validation data for BLEU
val_sample = val_dataset.select(range(min(100, len(val_dataset))))

predictions = []
references = []

# Re-load original val data for references
val_indices = split["test"].select(range(min(100, len(split["test"]))))
val_texts = Dataset.from_dict({
    "transliteration": train_df["transliteration"].iloc[val_indices.indices if hasattr(val_indices, 'indices') else range(len(val_indices))].tolist(),
    "translation": train_df["translation"].iloc[val_indices.indices if hasattr(val_indices, 'indices') else range(len(val_indices))].tolist()
})

# %%
from tqdm import tqdm

# Generate predictions
for i, text in enumerate(tqdm(val_texts["transliteration"], desc="Generating")):
    inputs = tokenizer(
        text,
        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_tgt_len"],
            num_beams=5,
            forced_bos_token_id=eng_token_id,
            repetition_penalty=1.5,
            no_repeat_ngram_size=3,
        )

    pred = tokenizer.decode(outputs[0], skip_special_tokens=True)
    predictions.append(pred)
    references.append(val_texts["translation"][i])

# %%
# Compute BLEU
bleu_metric = evaluate.load("sacrebleu")
chrf_metric = evaluate.load("chrf")

bleu_result = bleu_metric.compute(predictions=predictions, references=[[r] for r in references])
chrf_result = chrf_metric.compute(predictions=predictions, references=[[r] for r in references])

print("\n" + "="*60)
print("FINAL METRICS")
print("="*60)
print(f"BLEU: {bleu_result['score']:.2f}")
print(f"chrF: {chrf_result['score']:.2f}")
print(f"Eval Loss: {eval_results.get('eval_loss', 'N/A')}")

# %% [markdown]
# ## Save Model

# %%
# Save best model
best_model_path = f"{OUTPUT_DIR}/best"
trainer.save_model(best_model_path)
tokenizer.save_pretrained(best_model_path)

# Save config and metrics
metrics = {
    "bleu": bleu_result["score"],
    "chrf": chrf_result["score"],
    "eval_loss": eval_results.get("eval_loss"),
    "training_time_minutes": train_time / 60,
    "config": CONFIG,
    "timestamp": datetime.now().isoformat(),
}

with open(f"{best_model_path}/training_metrics.json", "w") as f:
    json.dump(metrics, f, indent=2)

print(f"\nModel saved to {best_model_path}")

# %%
# Show sample predictions
print("\n" + "="*60)
print("SAMPLE PREDICTIONS")
print("="*60)

for i in range(min(5, len(predictions))):
    print(f"\n[{i}]")
    print(f"SRC: {val_texts['transliteration'][i][:80]}...")
    print(f"REF: {references[i][:80]}...")
    print(f"OUT: {predictions[i][:80]}...")

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

# %%
# Create model metadata for upload
print("\nPreparing model for upload...")

# List files to upload
print(f"\nFiles in {best_model_path}:")
for f in os.listdir(best_model_path):
    size = os.path.getsize(f"{best_model_path}/{f}") / 1e6
    print(f"  {f}: {size:.1f} MB")

# %%
# Instructions for manual upload
print("\n" + "="*60)
print("UPLOAD INSTRUCTIONS")
print("="*60)
print("""
To upload this model to Kaggle:

1. From this kernel's output, download the model files
2. Go to https://www.kaggle.com/models
3. Click 'New Model' or update existing model
4. Upload the files from /kaggle/working/nllb-akkadian/best/

Or use kagglehub (if internet enabled):
```python
import kagglehub
kagglehub.model_upload(
    handle="manwithacat/akkadian-nllb-1b3/transformers/new-version",
    local_model_dir="/kaggle/working/nllb-akkadian/best",
    license_name="Apache 2.0",
    version_notes=f"BLEU: {bleu_result['score']:.2f}, Training: Kaggle P100"
)
```
""")

# %%
print("\n" + "="*60)
print("TRAINING COMPLETE")
print("="*60)
print(f"Total time: {train_time/60:.1f} minutes")
print(f"BLEU: {bleu_result['score']:.2f}")
print(f"Model saved to: {best_model_path}")