# %% [code]
# Deep Past Challenge - mT5 Inference
# Uses fine-tuned mT5-base model for Akkadian to English translation

import os
import pandas as pd
import torch
from tqdm import tqdm
import warnings
warnings.filterwarnings('ignore')

print("="*60)
print("Deep Past Challenge - mT5 Inference")
print("="*60)

# ============================================================================
# Install Dependencies
# ============================================================================
import subprocess
subprocess.run(['pip', 'install', '-q', 'transformers', 'sentencepiece', 'accelerate'],
               capture_output=True)

from transformers import MT5ForConditionalGeneration, MT5Tokenizer

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

# ============================================================================
# Load Pre-trained Model from Kaggle Dataset
# ============================================================================
MODEL_PATH = '/kaggle/input/mt5-akkadian-translation'

print(f"\nLoading fine-tuned model from: {MODEL_PATH}")

try:
    tokenizer = MT5Tokenizer.from_pretrained(MODEL_PATH, legacy=True)
    model = MT5ForConditionalGeneration.from_pretrained(MODEL_PATH)
    print("Model loaded successfully!")
except Exception as e:
    print(f"Error loading model: {e}")
    raise

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

param_count = sum(p.numel() for p in model.parameters())
print(f"Model parameters: {param_count / 1e6:.1f}M")
print(f"Using device: {device}")

# ============================================================================
# Load Test Data
# ============================================================================
test = pd.read_csv('/kaggle/input/deep-past-initiative-machine-translation/test.csv')
print(f"\nTest samples: {len(test)}")

# ============================================================================
# Preprocessing
# ============================================================================
def preprocess_text(text):
    """Clean and prepare text for mT5"""
    if pd.isna(text):
        return ""
    text = str(text).strip()
    text = ' '.join(text.split())
    return text

# ============================================================================
# Generate Predictions
# ============================================================================
print("\n" + "="*60)
print("Generating Predictions...")
print("="*60)

predictions = []

for idx, row in tqdm(test.iterrows(), total=len(test)):
    try:
        src = preprocess_text(row['transliteration'])
        input_text = f"translate Akkadian to English: {src}"

        inputs = tokenizer(
            input_text,
            max_length=256,
            padding=True,
            truncation=True,
            return_tensors='pt'
        ).to(device)

        with torch.no_grad():
            outputs = model.generate(
                input_ids=inputs['input_ids'],
                attention_mask=inputs['attention_mask'],
                max_length=256,
                num_beams=5,
                length_penalty=1.0,
                early_stopping=True,
                no_repeat_ngram_size=3,
                do_sample=False
            )

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

        # Fallback if empty
        if not translation or len(translation.strip()) < 2:
            translation = "The text describes a transaction or record."

        predictions.append({
            'id': int(row['id']),
            'translation': translation
        })

        # Print sample predictions
        if idx < 5 or idx % 100 == 0:
            print(f"\n[{idx}] Source: {src[:80]}...")
            print(f"    Translation: {translation[:80]}...")

    except Exception as e:
        print(f"Error on row {idx}: {e}")
        predictions.append({
            'id': int(row['id']),
            'translation': "The text describes a transaction or record."
        })

# ============================================================================
# Create Submission
# ============================================================================
submission = pd.DataFrame(predictions)
submission = submission.sort_values('id').reset_index(drop=True)

print(f"\n" + "="*60)
print("Submission Summary")
print("="*60)
print(f"Total rows: {len(submission)}")
print(f"\nFirst 10 predictions:")
for i, row in submission.head(10).iterrows():
    print(f"  ID {row['id']}: {row['translation'][:60]}...")

submission.to_csv('submission.csv', index=False)
print("\nSubmission saved to submission.csv")
print("="*60)