#!/usr/bin/env python3
"""
Deep Past Challenge: Akkadian to English Translation
Using top-scoring model from wukeneth (known high-performance model)
"""

import os
import re
import numpy as np
import pandas as pd
import torch
from pathlib import Path


# ========================================
# PREPROCESSING
# ========================================

def normalize_subscripts(text):
    """Convert Unicode subscript digits to regular digits."""
    subscript_map = str.maketrans("₀₁₂₃₄₅₆₇₈₉", "0123456789")
    return text.translate(subscript_map)

def preprocess_akkadian(text, prefix="translate Akkadian to English: "):
    """Preprocess Akkadian transliteration for the model."""
    if pd.isna(text) or text is None:
        return prefix
    text = str(text)

    # Normalize subscript digits
    text = normalize_subscripts(text)

    # Normalize ḫ/Ḫ to h/H (test uses h, training uses ḫ)
    text = text.replace('ḫ', 'h').replace('Ḫ', 'H')

    # Normalize gap markers
    text = re.sub(r'\[[^\]]*\]', '<gap>', text)
    text = text.replace('…', '<big_gap>').replace('...', '<big_gap>')

    # Clean up spaces
    text = ' '.join(text.split())

    return prefix + text


def postprocess_translation(text):
    """Clean up model output."""
    if not text:
        return ""
    text = re.sub(r'\s+([.,;:!?])', r'\1', text)
    return text.strip()


# ========================================
# FIND MODEL PATH
# ========================================

def find_model(search_name=None):
    """Find pre-trained model in Kaggle dataset locations."""
    # Print all available inputs for debugging
    print("=== Available inputs ===")
    if os.path.exists("/kaggle/input"):
        for item in os.listdir("/kaggle/input"):
            item_path = f"/kaggle/input/{item}"
            print(f"  /kaggle/input/{item}")
            if os.path.isdir(item_path):
                for sub in os.listdir(item_path)[:5]:
                    print(f"    -> {sub}")
    print("========================")

    # Search for model config.json recursively
    search_bases = ["/kaggle/input/datasets", "/kaggle/input"]

    for base in search_bases:
        if os.path.exists(base):
            for root, dirs, files in os.walk(base):
                if 'config.json' in files and (
                    any(f.endswith('.bin') or f.endswith('.safetensors') for f in files)
                ):
                    if search_name is None or search_name in root:
                        print(f"Found model at: {root}")
                        return root

    raise FileNotFoundError(f"No model found. search_name={search_name}")


# ========================================
# INFERENCE
# ========================================

def run_inference(test_file, model_path, output_file,
                  batch_size=4, num_beams=12, max_length=512):
    """Run inference and save submission."""
    from transformers import AutoTokenizer, AutoModelForSeq2SeqLM

    print(f"Loading test data from {test_file}")
    test_df = pd.read_csv(test_file)
    print(f"Test rows: {len(test_df)}")

    print(f"\nLoading model from {model_path}")
    tokenizer = AutoTokenizer.from_pretrained(model_path)

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

    model = AutoModelForSeq2SeqLM.from_pretrained(
        model_path,
        dtype=torch.float32,
    )
    model = model.to(device)
    model.eval()

    n_params = sum(p.numel() for p in model.parameters())
    print(f"Parameters: {n_params:,}")

    # Preprocess
    texts = test_df['transliteration'].fillna("").tolist()
    texts = [preprocess_akkadian(t) for t in texts]

    print(f"\nRunning inference: {len(texts)} texts, {num_beams} beams")

    translations = []
    for i in range(0, len(texts), batch_size):
        batch = texts[i:i+batch_size]
        batch_end = min(i+batch_size, len(texts))
        print(f"  Processing {i+1}-{batch_end}/{len(texts)}", flush=True)

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

        with torch.no_grad():
            outputs = model.generate(
                **inputs,
                max_new_tokens=max_length,
                num_beams=num_beams,
                early_stopping=True,
                length_penalty=1.0,
                no_repeat_ngram_size=4,
            )

        batch_translations = tokenizer.batch_decode(outputs, skip_special_tokens=True)
        batch_translations = [postprocess_translation(t) for t in batch_translations]
        translations.extend(batch_translations)

        # Show examples
        if i == 0:
            for j, (src, tgt) in enumerate(zip(batch[:2], batch_translations[:2])):
                print(f"\n  Input:  {src[:80]}...")
                print(f"  Output: {tgt[:80]}")

    # Save submission
    submission = pd.DataFrame({
        'id': test_df['id'],
        'translation': translations
    })
    submission.to_csv(output_file, index=False)
    print(f"\nSaved: {output_file}")
    print(submission)

    return submission


# ========================================
# MAIN
# ========================================

if __name__ == "__main__":
    # Paths
    test_file_candidates = [
        "/kaggle/input/deep-past-initiative-machine-translation/test.csv",
        "/kaggle/input/competitions/deep-past-initiative-machine-translation/test.csv",
    ]
    output_file = "/kaggle/working/submission.csv"

    # Find test file
    test_file = None
    for path in test_file_candidates:
        if os.path.exists(path):
            test_file = path
            break
    if test_file is None:
        raise FileNotFoundError(f"Test file not found in {test_file_candidates}")
    print(f"Using test file: {test_file}")

    # Find model
    model_path = find_model()

    # Run inference with 12 beams for better quality
    submission = run_inference(
        test_file=test_file,
        model_path=model_path,
        output_file=output_file,
        batch_size=4,
        num_beams=12,
        max_length=512,
    )

    print("\n=== DONE ===")
    print(f"Generated {len(submission)} translations")