#!/usr/bin/env python3
"""
Deep Past Challenge: Akkadian to English Translation
Inference notebook for Kaggle submission

Uses pre-trained byt5-base model from assiaben/final-byt5
"""

import os
import re
import sys
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):
    """Preprocess Akkadian transliteration for the model."""
    if pd.isna(text) or text is None:
        return ""
    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
    # [x], [y] etc -> <gap>
    text = re.sub(r'\[[^\]]*\]', '<gap>', text)
    # ... or … -> <big_gap>
    text = text.replace('…', '<big_gap>').replace('...', '<big_gap>')

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

    return text


def postprocess_translation(text):
    """Clean up model output."""
    if not text:
        return ""

    # Fix spaces before punctuation
    text = re.sub(r'\s+([.,;:!?])', r'\1', text)
    text = re.sub(r'(["\'])\s+', r'\1 ', text)

    # Remove duplicate words (simple ngram dedup)
    words = text.split()
    if len(words) > 3:
        seen_bigrams = set()
        result = [words[0]]
        for i in range(1, len(words)):
            bigram = (words[i-1].lower(), words[i].lower())
            if bigram not in seen_bigrams or words[i] in '.,:;':
                result.append(words[i])
                seen_bigrams.add(bigram)
        text = ' '.join(result)

    return text.strip()


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

def find_model_path():
    """Find the pre-trained model in Kaggle dataset locations."""
    # Kaggle input structure (from debug output):
    # /kaggle/input/datasets/assiaben/final-byt5/byt5-akkadian-optimized-34x/
    # /kaggle/input/competitions/deep-past-initiative-machine-translation/
    possible_paths = [
        "/kaggle/input/datasets/assiaben/final-byt5/byt5-akkadian-optimized-34x",
        "/kaggle/input/datasets/assiaben/final-byt5",
        "/kaggle/input/final-byt5/byt5-akkadian-optimized-34x",
        "/kaggle/input/final-byt5",
    ]

    # Check for model files
    for path in possible_paths:
        if os.path.exists(path):
            # Check if it has model files
            files = os.listdir(path)
            print(f"Checking {path}: {files}")
            if any(f.endswith('.bin') or f.endswith('.safetensors') or f == 'config.json' for f in files):
                print(f"Found model at: {path}")
                return path

    # Try to find recursively in common dataset locations
    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
                ):
                    print(f"Found model at: {root}")
                    return root

    raise FileNotFoundError(f"Could not find model. Searched: {possible_paths}")


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

def run_inference(test_file, model_path, output_file,
                  batch_size=8, num_beams=8, max_length=512,
                  prefix="translate Akkadian to English: "):
    """Run inference on test data 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(test_df.head())

    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,
        torch_dtype=torch.float32,
    )
    model = model.to(device)
    model.eval()

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

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

    print(f"\nRunning inference with num_beams={num_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  Example {j+1}:")
                print(f"  Input:  {src[:80]}...")
                print(f"  Output: {tgt[:80]}")

    # Create submission
    submission = pd.DataFrame({
        'id': test_df['id'],
        'translation': translations
    })

    submission.to_csv(output_file, index=False)
    print(f"\nSubmission saved to {output_file}")
    print(submission.head())

    return submission


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

if __name__ == "__main__":
    # Kaggle paths - use both old and new path conventions
    test_file_paths = [
        "/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"

    # Debug: list available inputs (2 levels deep)
    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):
                    sub_path = f"{item_path}/{sub}"
                    print(f"    -> {sub}")
                    if os.path.isdir(sub_path):
                        for subsub in os.listdir(sub_path)[:5]:
                            print(f"      -> {subsub}")
    print("========================")

    # Find test file
    test_file = None
    for path in test_file_paths:
        if os.path.exists(path):
            test_file = path
            break
    if test_file is None:
        raise FileNotFoundError(f"Could not find test.csv in {test_file_paths}")

    # Find model
    model_path = find_model_path()

    # Run inference
    submission = run_inference(
        test_file=test_file,
        model_path=model_path,
        output_file=output_file,
        batch_size=8,
        num_beams=8,
        max_length=512,
    )

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