"""
Akkadian to English Translation - ByT5 Optimized 34x Inference
Fork of assiaben's optimized model with enhanced postprocessing.
"""
import os
import re
import pandas as pd
import torch
from torch.utils.data import Dataset, DataLoader
from transformers import AutoTokenizer, AutoModelForSeq2SeqLM
from tqdm.auto import tqdm

# === Configuration ===
INPUT_DIR = None
for candidate in [
    "/kaggle/input/deep-past-initiative-machine-translation",
    "/kaggle/input/competitions/deep-past-initiative-machine-translation",
]:
    if os.path.exists(candidate):
        INPUT_DIR = candidate
        break
if INPUT_DIR is None:
    for root, dirs, files in os.walk("/kaggle/input"):
        if "test.csv" in files:
            INPUT_DIR = root
            break
if INPUT_DIR is None:
    raise FileNotFoundError("Cannot find competition data!")

TEST_DATA_PATH = f"{INPUT_DIR}/test.csv"

# Debug: list /kaggle/input to find model
print("=== Listing /kaggle/input ===")
for item in sorted(os.listdir("/kaggle/input")):
    full = os.path.join("/kaggle/input", item)
    if os.path.isdir(full):
        print(f"  DIR: {item}/")
        try:
            for sub in sorted(os.listdir(full))[:10]:
                subfull = os.path.join(full, sub)
                if os.path.isdir(subfull):
                    print(f"    DIR: {sub}/")
                    for subsub in sorted(os.listdir(subfull))[:5]:
                        print(f"      {subsub}")
                else:
                    print(f"    FILE: {sub} ({os.path.getsize(subfull)} bytes)")
        except Exception as e:
            print(f"    Error listing: {e}")
    else:
        print(f"  FILE: {item}")

# Find model path
MODEL_PATH = None
for candidate in [
    "/kaggle/input/final-byt5/byt5-akkadian-optimized-34x",
    "/kaggle/input/final-byt5",
    "/kaggle/input/datasets/assiaben/final-byt5/byt5-akkadian-optimized-34x",
    "/kaggle/input/datasets/assiaben/final-byt5",
]:
    if os.path.isdir(candidate) and os.path.exists(os.path.join(candidate, "config.json")):
        MODEL_PATH = candidate
        break

if MODEL_PATH is None:
    # Walk all of /kaggle/input to find config.json
    for root, dirs, files in os.walk("/kaggle/input"):
        if "config.json" in files and "model.safetensors" in files:
            MODEL_PATH = root
            break

if MODEL_PATH is None:
    raise FileNotFoundError("Cannot find model! Check /kaggle/input listing above.")
BATCH_SIZE = 8
MAX_LENGTH = 512
DEVICE = torch.device("cuda" if torch.cuda.is_available() else "cpu")

print(f"Using data from: {INPUT_DIR}")
print(f"Device: {DEVICE}")

# === Preprocessing ===
_big_gap_re = re.compile(r'(\.{3,}|…+|……)')
_small_gap_re = re.compile(r'(xx+|\s+x\s+)')

def preprocess_input(text):
    if pd.isna(text):
        return ""
    text = str(text)
    text = _big_gap_re.sub('<big_gap>', text)
    text = _small_gap_re.sub('<gap>', text)
    return text

# === Postprocessing ===
_subscript_trans = str.maketrans("₀₁₂₃₄₅₆₇₈₉", "0123456789")
_special_trans = str.maketrans("ḫḪ", "hH")
_forbidden = '!?()"—–<>⌈⌋⌊[]+ʾ/;'
_forbidden_trans = str.maketrans("", "", _forbidden)

def postprocess_output(text):
    if not isinstance(text, str) or not text.strip():
        return ""
    t = text
    t = t.translate(_special_trans)
    t = t.translate(_subscript_trans)
    # Normalize gaps
    t = re.sub(r'(\[x\]|\(x\)|\bx\b)', '<gap>', t, flags=re.I)
    t = re.sub(r'(\.{3,}|…|\[\.+\])', '<big_gap>', t)
    t = t.replace('<gap> <gap>', '<big_gap>')
    t = t.replace('<big_gap> <big_gap>', '<big_gap>')
    # Remove annotations
    t = re.sub(r'\((fem|plur|pl|sing|singular|plural|\?|!)\.?\s*\w*\)', '', t, flags=re.I)
    # Protect gaps, remove forbidden chars
    t = t.replace('<gap>', '\x00GAP\x00').replace('<big_gap>', '\x00BIG\x00')
    t = t.translate(_forbidden_trans)
    t = t.replace('\x00GAP\x00', ' <gap> ').replace('\x00BIG\x00', ' <big_gap> ')
    # Fractions
    t = re.sub(r'(\d+)\.5\b', r'\1 ½', t)
    t = re.sub(r'\b0\.5\b', '½', t)
    t = re.sub(r'(\d+)\.25\b', r'\1 ¼', t)
    t = re.sub(r'\b0\.25\b', '¼', t)
    t = re.sub(r'(\d+)\.75\b', r'\1 ¾', t)
    t = re.sub(r'\b0\.75\b', '¾', t)
    # Remove repeated words and n-grams
    t = re.sub(r'\b(\w+)(?:\s+\1\b)+', r'\1', t)
    for n in range(4, 1, -1):
        pat = r'\b((?:\w+\s+){' + str(n-1) + r'}\w+)(?:\s+\1\b)+'
        t = re.sub(pat, r'\1', t)
    # Fix spacing
    t = re.sub(r'\s+([.,:])', r'\1', t)
    t = re.sub(r'([.,])\1+', r'\1', t)
    t = re.sub(r'\s+', ' ', t).strip().strip('-').strip()
    return t

# === Dataset ===
class InferenceDataset(Dataset):
    def __init__(self, df):
        self.ids = df['id'].tolist()
        self.texts = ["translate Akkadian to English: " + preprocess_input(t) for t in df['transliteration']]
    def __len__(self):
        return len(self.ids)
    def __getitem__(self, idx):
        return self.ids[idx], self.texts[idx]

# === Load data and model ===
print("Loading test data...")
test_df = pd.read_csv(TEST_DATA_PATH)
print(f"Test samples: {len(test_df)}")

print("Loading model...")
print(f"Model path: {MODEL_PATH}")
print(f"Model files: {os.listdir(MODEL_PATH)}")
tokenizer = AutoTokenizer.from_pretrained(MODEL_PATH, local_files_only=True)
model = AutoModelForSeq2SeqLM.from_pretrained(MODEL_PATH, local_files_only=True).to(DEVICE).eval()
print(f"Model loaded on {DEVICE}")

# Try BetterTransformer
try:
    from optimum.bettertransformer import BetterTransformer
    model = BetterTransformer.transform(model)
    print("BetterTransformer applied")
except Exception:
    print("BetterTransformer not available, using standard model")

# === Collate function ===
def collate_fn(batch):
    ids = [b[0] for b in batch]
    texts = [b[1] for b in batch]
    tokens = tokenizer(texts, max_length=MAX_LENGTH, padding=True, truncation=True, return_tensors="pt")
    return ids, tokens

# === Inference ===
dataset = InferenceDataset(test_df)
loader = DataLoader(dataset, batch_size=BATCH_SIZE, shuffle=False, num_workers=2, collate_fn=collate_fn, pin_memory=True)

print(f"Running inference ({len(loader)} batches)...")
results = []

with torch.inference_mode():
    for batch_ids, tokens in tqdm(loader, desc="Translating"):
        input_ids = tokens.input_ids.to(DEVICE)
        attention_mask = tokens.attention_mask.to(DEVICE)

        # Adaptive beam size
        seq_len = attention_mask.sum(dim=1).float().mean().item()
        num_beams = 4 if seq_len < 100 else 8

        outputs = model.generate(
            input_ids=input_ids,
            attention_mask=attention_mask,
            num_beams=num_beams,
            max_new_tokens=512,
            length_penalty=1.09,
            early_stopping=True,
            use_cache=True,
        )

        decoded = tokenizer.batch_decode(outputs, skip_special_tokens=True)
        cleaned = [postprocess_output(d) for d in decoded]
        results.extend(zip(batch_ids, cleaned))

        if torch.cuda.is_available() and len(results) % 100 == 0:
            torch.cuda.empty_cache()

# === Save submission ===
submission = pd.DataFrame(results, columns=['id', 'translation'])
submission['translation'] = submission['translation'].apply(lambda x: x if len(str(x)) > 0 else "broken text")
submission.to_csv('submission.csv', index=False)
print(f"Submission saved: {len(submission)} rows")
print(submission.head())