"""Deep Past Challenge — Inference Notebook
Hybrid TF-IDF + ByT5 MBR v2 pipeline.

Approach (see writeup in discussion):
- TF-IDF word+char retrieval on train.csv (score >= 1.5 → copy translation)
- ByT5 MBR v2 (mattiaangeli/byt5-akkadian-mbr-v2) for low-confidence cases
- Regex post-processing for normalization

This solution scores Combined 37.18 on local 156-text cross-val.
"""

import math
import os
import re
import sys
import time

import numpy as np
import pandas as pd
import torch
from scipy.sparse import hstack
from sklearn.feature_extraction.text import TfidfVectorizer
from transformers import AutoModelForSeq2SeqLM, AutoTokenizer

# ============ CONFIG ============
INPUT_DIR = "/kaggle/input/deep-past-initiative-machine-translation"
MODEL_DIR = "/kaggle/input/byt5-akkadian-mbr-v2/pytorch/default/1"
OUTPUT_PATH = "/kaggle/working/submission.csv"

# Debug: list all input files
import glob
print("=== Kaggle input files ===")
for f in sorted(glob.glob("/kaggle/input/**/*", recursive=True))[:50]:
    if os.path.isfile(f):
        print(f"  {f}")
print("===")

# Auto-detect paths
_test_candidates = glob.glob("/kaggle/input/**/test.csv", recursive=True)
if _test_candidates:
    INPUT_DIR = os.path.dirname(_test_candidates[0])
    print(f"Found test.csv at: {INPUT_DIR}")

_model_candidates = glob.glob("/kaggle/input/**/model.safetensors", recursive=True)
if _model_candidates:
    MODEL_DIR = os.path.dirname(_model_candidates[0])
    print(f"Found model at: {MODEL_DIR}")

TFIDF_THRESHOLD = 1.5  # TF-IDF score >= this → copy, else ByT5
NUM_BEAMS = 4
MAX_LENGTH = 384
BATCH_SIZE = 4  # for ByT5 batched generation


def tokenize_tfidf(text):
    return str(text).lower().replace("-", " ").replace(".", " ").replace("_", " ")


def build_tfidf_index(pairs):
    sources_tok = [tokenize_tfidf(s) for s, _ in pairs]
    vec_word = TfidfVectorizer(
        analyzer="word", ngram_range=(1, 2), max_features=40000, sublinear_tf=True
    )
    vec_char = TfidfVectorizer(
        analyzer="char_wb", ngram_range=(2, 6), max_features=40000, sublinear_tf=True
    )
    m_word = vec_word.fit_transform(sources_tok)
    m_char = vec_char.fit_transform(sources_tok)
    matrix = hstack([m_word, m_char])
    return {"vec_word": vec_word, "vec_char": vec_char, "matrix": matrix, "pairs": pairs}


def tfidf_retrieve(index, queries):
    q_tok = [tokenize_tfidf(q) for q in queries]
    q_word = index["vec_word"].transform(q_tok)
    q_char = index["vec_char"].transform(q_tok)
    q_matrix = hstack([q_word, q_char])
    scores = (q_matrix @ index["matrix"].T).toarray()
    results = []
    for i in range(len(queries)):
        best_idx = np.argmax(scores[i])
        results.append((index["pairs"][best_idx][1], float(scores[i, best_idx])))
    return results


def postprocess(text):
    """Regex corrections (from top competition notebooks)."""
    text = text.replace("\u201c", '"').replace("\u201d", '"')
    text = text.replace("\u2018", "'").replace("\u2019", "'")
    text = text.replace("\u2013", "-").replace("\u2014", "-").replace("\u2212", "-")
    text = re.sub(r"\s+", " ", text).strip()
    text = re.sub(r"\b(\w+)(?:\s+\1\b)+", r"\1", text)
    text = re.sub(r"\[x\]|\(x\)|\bx\b(?=\s|$)", "<gap>", text, flags=re.I)
    text = re.sub(r"\.{3,}|…|\[\.+\]", "<gap>", text)
    text = re.sub(r"\s+([,.;:!?])", r"\1", text)
    text = re.sub(r"([,.;:!?])([A-Za-z])", r"\1 \2", text)
    if text.count("[") > text.count("]"):
        text += "]" * (text.count("[") - text.count("]"))
    if text.count("(") > text.count(")"):
        text += ")" * (text.count("(") - text.count(")"))
    text = re.sub(r"\.\.", ".", text)
    text = re.sub(r"<<>>|<>", "", text)
    if text and text[0].islower():
        text = text[0].upper() + text[1:]
    return text.strip()


def batch_generate(model, tokenizer, texts, device, batch_size=BATCH_SIZE):
    """Batched beam search on ByT5."""
    results = []
    for i in range(0, len(texts), batch_size):
        batch = texts[i : i + batch_size]
        prompts = [f"translate Akkadian to English: {t}" for t in batch]
        inputs = tokenizer(
            prompts, return_tensors="pt", max_length=512, truncation=True, padding=True
        ).to(device)
        with torch.no_grad():
            outputs = model.generate(
                **inputs,
                max_new_tokens=MAX_LENGTH,
                num_beams=NUM_BEAMS,
                length_penalty=1.3,
                early_stopping=True,
            )
        results.extend(tokenizer.batch_decode(outputs, skip_special_tokens=True))
    return results


def main():
    t0 = time.time()

    # 1. Load data
    print("Loading data...")
    test = pd.read_csv(f"{INPUT_DIR}/test.csv")
    train = pd.read_csv(f"{INPUT_DIR}/train.csv")
    print(f"Test: {len(test)}, Train: {len(train)}")

    # 2. Build TF-IDF index from train
    print("Building TF-IDF index...")
    train_pairs = list(zip(train["transliteration"].tolist(), train["translation"].tolist()))
    index = build_tfidf_index(train_pairs)

    # 3. TF-IDF retrieval for all test texts
    test_sources = test["transliteration"].tolist()
    tfidf_results = tfidf_retrieve(index, test_sources)

    # Classify into tiers
    copy_indices = []
    byt5_indices = []
    for i, (_, score) in enumerate(tfidf_results):
        if score >= TFIDF_THRESHOLD:
            copy_indices.append(i)
        else:
            byt5_indices.append(i)
    print(f"TF-IDF copy: {len(copy_indices)}, ByT5: {len(byt5_indices)}")

    # 4. Load ByT5 model and generate for low-confidence texts
    hypotheses = [None] * len(test_sources)

    # Fill copy indices
    for i in copy_indices:
        hypotheses[i] = postprocess(tfidf_results[i][0])

    if byt5_indices:
        print(f"Loading ByT5 from {MODEL_DIR}...")
        # Detect GPU compatibility (P100/sm_60 is incompatible with PyTorch 2.10+)
        device = "cpu"
        if torch.cuda.is_available():
            try:
                cap = torch.cuda.get_device_capability(0)
                name = torch.cuda.get_device_name(0)
                print(f"GPU: {name}, sm_{cap[0]}{cap[1]}, torch={torch.__version__}")
                if cap >= (7, 0):
                    # Quick CUDA test
                    _test = torch.randn(2, 2, device="cuda") @ torch.randn(2, 2, device="cuda")
                    _ = _test.sum().item()
                    device = "cuda"
                    print("CUDA OK")
                else:
                    print(f"GPU sm_{cap[0]}{cap[1]} not supported — falling back to CPU")
            except Exception as e:
                print(f"CUDA test failed: {e} — falling back to CPU")

        tokenizer = AutoTokenizer.from_pretrained(MODEL_DIR)
        model = AutoModelForSeq2SeqLM.from_pretrained(MODEL_DIR, torch_dtype=torch.float32)
        model.eval().to(device)
        print(f"ByT5 loaded on {device}")

        byt5_texts = [test_sources[i] for i in byt5_indices]
        t_gen = time.time()
        byt5_hyps = batch_generate(model, tokenizer, byt5_texts, device)
        print(f"ByT5 generated {len(byt5_hyps)} in {time.time()-t_gen:.1f}s")

        for idx, hyp in zip(byt5_indices, byt5_hyps):
            hypotheses[idx] = postprocess(hyp)

    assert all(h is not None for h in hypotheses)

    # 5. Write submission
    submission = pd.DataFrame({"id": test["id"], "translation": hypotheses})
    submission.to_csv(OUTPUT_PATH, index=False)
    print(f"Submission saved to {OUTPUT_PATH}")
    print(f"Total time: {time.time()-t0:.1f}s")

    # Preview
    print("\nPreview:")
    print(submission.head(10).to_string())


if __name__ == "__main__":
    main()