"""
Deep Past - ByT5 Single Model Inference (BACKUP)

Minimal inference using just manwithacat/byt5-combined
No model soup, no ensemble - just one proven model.

Expected BLEU: ~37.2
"""

import os
import re
import pandas as pd
from pathlib import Path

# Install dependencies
os.system("pip install -q transformers sentencepiece accelerate")

import torch
from transformers import AutoModelForSeq2SeqLM, AutoTokenizer

print("=" * 60)
print("Deep Past - ByT5 Simple Inference (BACKUP)")
print("=" * 60)

# ============================================================
# CONFIGURATION
# ============================================================

DATASET_NAME = "byt5-combined-finetuned"  # pascalledesma/byt5-combined-finetuned
TASK_PREFIX = "translate Akkadian to English: "

# ============================================================
# 1. LOCATE MODEL (KAGGLE DATASET - NO INTERNET)
# ============================================================

print(f"\nLoading model from dataset: {DATASET_NAME}")
print("(Using Kaggle dataset input - offline mode)")

# CRITICAL: Model uploaded as dataset (API workaround)
# Models must be datasets because Kaggle API ignores model_sources

def find_model_path():
    """Find model directory in /kaggle/input datasets"""
    base = Path("/kaggle/input")

    # Primary path - dataset uploaded via kaggle_model_to_dataset.py
    primary = base / DATASET_NAME
    if primary.exists():
        # Find config.json recursively
        for cfg in primary.rglob("config.json"):
            return str(cfg.parent)
        # If no config.json, return base path (model files might be flat)
        return str(primary)

    # Search all inputs for model files
    try:
        for item in base.iterdir():
            if DATASET_NAME.replace("-", "") in item.name.replace("-", ""):
                for cfg in item.rglob("config.json"):
                    return str(cfg.parent)
    except Exception:
        pass

    return None

model_path = find_model_path()

if model_path:
    print(f"Found model at: {model_path}")
    try:
        tokenizer = AutoTokenizer.from_pretrained(model_path)
        model = AutoModelForSeq2SeqLM.from_pretrained(model_path)
        print("Model loaded successfully from Kaggle dataset")
    except Exception as e:
        print(f"Error loading model: {e}")
        raise RuntimeError(f"Failed to load model from {model_path}")
else:
    print(f"CRITICAL ERROR: Dataset not found in /kaggle/input")
    print(f"\nExpected dataset: pascalledesma/{DATASET_NAME}")
    print(f"This should be auto-attached via kernel-metadata.json")
    raise RuntimeError("Model dataset not found.")

device = "cuda" if torch.cuda.is_available() else "cpu"
model = model.to(device)
model.eval()
print(f"Device: {device}")

# ============================================================
# 2. LOAD TEST DATA
# ============================================================

print("\nLoading test data...")

test_path = "/kaggle/input/deep-past-initiative-machine-translation/test.csv"
if not Path(test_path).exists():
    # Search for test.csv
    for p in Path("/kaggle/input").rglob("test.csv"):
        test_path = str(p)
        break

test_df = pd.read_csv(test_path)
print(f"Loaded {len(test_df)} test samples")

# Find source column
source_col = "transliteration"
if source_col not in test_df.columns:
    for col in test_df.columns:
        if col.lower() not in ["id", "target", "translation"]:
            source_col = col
            break
print(f"Source column: {source_col}")

# ============================================================
# 3. PREPROCESSING
# ============================================================

def preprocess(text):
    """Normalize Akkadian input."""
    if pd.isna(text) or not str(text).strip():
        return ""

    text = str(text)

    # Normalize gap markers
    text = re.sub(r'\.\.\.+', '<gap>', text)
    text = re.sub(r'\bxx?\b', '<gap>', text)
    text = re.sub(r'\[\.+\]', '<gap>', text)

    return text.strip()

# ============================================================
# 4. POSTPROCESSING
# ============================================================

def postprocess(text):
    """Clean and format output."""
    if not text:
        return ""

    # Remove non-ASCII
    text = re.sub(r'[^\x20-\x7E\s]', '', text)

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

    # Remove word repetitions
    text = re.sub(r'\b(\w+)\s+\1\b', r'\1', text, flags=re.IGNORECASE)

    # Fix fractions
    text = re.sub(r'(\d)\s*/\s*(\d)', r'\1/\2', text)

    # Clean whitespace
    text = re.sub(r'\s+', ' ', text).strip()

    # Capitalize first letter
    if text:
        text = text[0].upper() + text[1:]

    return text

# ============================================================
# 5. TRANSLATE
# ============================================================

def translate(text, max_length=256):
    """Translate Akkadian to English."""
    text = preprocess(text)
    if not text:
        return ""

    input_text = TASK_PREFIX + text

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

    with torch.no_grad():
        outputs = model.generate(
            **inputs,
            max_length=max_length,
            num_beams=10,  # Research: 7-10 optimal for ancient languages
            early_stopping=True,
            no_repeat_ngram_size=3,
            repetition_penalty=1.5,  # Stronger anti-repetition
            length_penalty=1.2,  # Allow longer, more detailed outputs
        )

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

# ============================================================
# 6. RUN INFERENCE
# ============================================================

print("\nTranslating test data...")

translations = []
for i, row in test_df.iterrows():
    source = row[source_col]
    trans = translate(source)
    translations.append(trans)

    # Progress indicator
    if i < 5:
        print(f"\n[{i}] Source: {str(source)[:70]}...")
        print(f"    Output: {trans[:70]}..." if trans else "    Output: (empty)")
    elif i % 100 == 0:
        print(f"  [{i}/{len(test_df)}] Progress...")

# ============================================================
# 7. CREATE SUBMISSION
# ============================================================

print("\nCreating submission...")

submission = pd.DataFrame({
    "id": test_df["id"],
    "translation": translations  # CRITICAL: must be "translation" not "target"
})

# Validate
empty_count = (submission["translation"] == "").sum()
if empty_count > 0:
    print(f"⚠ Warning: {empty_count} empty translations")

total_chars = submission["translation"].str.len().sum()
print(f"Total characters: {total_chars}")

submission.to_csv("submission.csv", index=False)
print(f"\n✓ Saved: submission.csv")
print(f"  Rows: {len(submission)}")
print(f"  Column names: {list(submission.columns)}")
print(f"\nFirst 5 rows:")
print(submission.head())

print("\n" + "=" * 60)
print("ByT5 Simple Inference Complete!")
print(f"  Dataset: pascalledesma/{DATASET_NAME}")
print(f"  Expected BLEU: ~37.2")
print(f"  Status: BACKUP (use if model soup fails)")
print("=" * 60)