#!/usr/bin/env python
# coding: utf-8

# # Deep Past Initiative – Machine Translation (Inference Notebook)
#
# This notebook is a **starter / baseline** for this Kaggle competition.
#
# Training Code is [here](https://www.kaggle.com/code/takamichitoda/dpc-starter-train).

# In[1]:


import os

# Workaround: 新しいhuggingface_hubがローカルパスをrepo IDとして検証するバグの回避
import huggingface_hub.utils._validators
_orig_validate = huggingface_hub.utils._validators.validate_repo_id
def _patched_validate(repo_id, *args, **kwargs):
    if isinstance(repo_id, str) and (repo_id.startswith('/') or os.path.isdir(repo_id)):
        return
    _orig_validate(repo_id, *args, **kwargs)
huggingface_hub.utils._validators.validate_repo_id = _patched_validate

import json
import re
import shutil
import pandas as pd
import torch
from torch.utils.data import DataLoader, Dataset
from transformers import AutoTokenizer, AutoModelForSeq2SeqLM
from tqdm.auto import tqdm


# In[2]:


def postprocess_output(text):
    """上位勢の後処理スタック（Community Knowledge Synthesis準拠）"""
    if not isinstance(text, str) or not text.strip():
        return ""

    # 1. ḫ/Ḫ → h/H
    text = text.replace('ḫ', 'h').replace('Ḫ', 'H')

    # 2. 下付き数字 → 通常数字
    subscript_map = str.maketrans("₀₁₂₃₄₅₆₇₈₉", "0123456789")
    text = text.translate(subscript_map)

    # 3. gap/big_gap正規化
    text = re.sub(r'(\[x\]|\(x\)|\bx\b)', '<gap>', text, flags=re.I)
    text = re.sub(r'(\.{3,}|…)', '<big_gap>', text)

    # 4. scribal annotation除去
    text = re.sub(r'\((fem|plur|pl|sing|\?|!)\.?\)', '', text, flags=re.I)

    # 5. 禁止文字除去（gapマーカー保護）
    text = text.replace('<gap>', '\x00GAP\x00')
    text = text.replace('<big_gap>', '\x00BIG\x00')
    forbidden = '!?()\"—–<>⌈⌋⌊[]+ʾ/;'
    text = text.translate(str.maketrans('', '', forbidden))
    text = text.replace('\x00GAP\x00', ' <gap> ')
    text = text.replace('\x00BIG\x00', ' <big_gap> ')

    # 6. Unicode分数変換
    text = re.sub(r'(\d+)\.5\b', r'\1 ½', text)
    text = re.sub(r'(\d+)\.25\b', r'\1 ¼', text)
    text = re.sub(r'(\d+)\.75\b', r'\1 ¾', text)
    text = re.sub(r'(\d+)\.333+\b', r'\1 ⅓', text)
    text = re.sub(r'(\d+)\.666+\b', r'\1 ⅔', text)

    # 7. 単語レベル重複除去
    text = re.sub(r'\b(\w+)(?:\s+\1\b)+', r'\1', text)

    # 8. 空白正規化
    text = re.sub(r'\s+', ' ', text).strip()

    return text


MODEL_PATH = "/kaggle/input/exp-40-train/byt5-akkadian-model/"
TRAIN_OUTPUT_DIR = "/kaggle/input/exp-40-train/"
# kernels push経由の場合、kernel_sourcesは /kaggle/input/notebooks/{user}/{slug}/ にマウントされる
if not os.path.exists(MODEL_PATH):
    MODEL_PATH = "/kaggle/input/notebooks/raipachi0704/exp-40-train/byt5-akkadian-model/"
    TRAIN_OUTPUT_DIR = "/kaggle/input/notebooks/raipachi0704/exp-40-train/"


# In[3]:


TEST_DATA_PATH = "/kaggle/input/deep-past-initiative-machine-translation/test.csv"
if not os.path.exists(TEST_DATA_PATH):
    TEST_DATA_PATH = "/kaggle/input/competitions/deep-past-initiative-machine-translation/test.csv"
BATCH_SIZE = 8  # 16→8: beam数増加に伴うメモリ対策
MAX_LENGTH = 512
NUM_BEAMS = 8  # 上位勢設定（baseline: 4）
LENGTH_PENALTY = 1.3  # 上位勢推奨（baseline: 1.0）
DEVICE = "cuda" if torch.cuda.is_available() else "cpu"

# --- Model Loading ---
print(f"Loading model from {MODEL_PATH}...")
tokenizer = AutoTokenizer.from_pretrained(MODEL_PATH)
model = AutoModelForSeq2SeqLM.from_pretrained(MODEL_PATH).to(DEVICE)
model.eval()

# --- Data Preparation ---
test_df = pd.read_csv(TEST_DATA_PATH)


# In[4]:


PREFIX = "translate Akkadian to English: "

class InferenceDataset(Dataset):
    def __init__(self, df, tokenizer):
        self.texts = df['transliteration'].astype(str).tolist()
        self.texts = [PREFIX + i for i in self.texts]
        self.tokenizer = tokenizer

    def __len__(self):
        return len(self.texts)

    def __getitem__(self, idx):
        text = self.texts[idx]
        inputs = self.tokenizer(
            text,
            max_length=MAX_LENGTH,
            padding="max_length",
            truncation=True,
            return_tensors="pt"
        )
        return {
            "input_ids": inputs["input_ids"].squeeze(0),
            "attention_mask": inputs["attention_mask"].squeeze(0)
        }

test_dataset = InferenceDataset(test_df, tokenizer)
test_loader = DataLoader(test_dataset, batch_size=BATCH_SIZE, shuffle=False)

# --- Inference Loop ---
print("Starting Inference...")
all_predictions = []


# In[5]:


with torch.no_grad():
    for batch in tqdm(test_loader):
        input_ids = batch["input_ids"].to(DEVICE)
        attention_mask = batch["attention_mask"].to(DEVICE)

        outputs = model.generate(
            input_ids=input_ids,
            attention_mask=attention_mask,
            max_length=MAX_LENGTH,
            num_beams=NUM_BEAMS,
            length_penalty=LENGTH_PENALTY,
            early_stopping=True,
        )

        decoded = tokenizer.batch_decode(outputs, skip_special_tokens=True)
        all_predictions.extend([postprocess_output(d) for d in decoded])


# In[6]:


# --- Submission & Predictions Output ---
submission = pd.DataFrame({
    "id": test_df["id"],
    "translation": all_predictions
})

submission["translation"] = submission["translation"].apply(lambda x: x if len(x) > 0 else "broken text")

submission.to_csv("submission.csv", index=False)
print("submission.csv saved successfully!")

submission.to_csv("predictions.csv", index=False)
print("predictions.csv saved successfully!")

submission.head()


# In[7]:


# --- Copy Metrics from Training Output ---
train_metrics_path = os.path.join(TRAIN_OUTPUT_DIR, "metrics.json")

if os.path.exists(train_metrics_path):
    shutil.copy2(train_metrics_path, "metrics.json")
    with open("metrics.json", "r") as f:
        metrics = json.load(f)
    print(f"metrics.json copied from training output.")
    print(f"Local CV ({metrics.get('metric_name', 'unknown')}): {metrics.get('local_cv')}")
else:
    print(f"WARNING: Training metrics not found at {train_metrics_path}")
    print("Creating placeholder metrics.json...")
    metrics = {
        "local_cv": None,
        "local_cv_std": None,
        "metric_name": "chrF++",
        "details": {
            "note": "Training metrics not available. Run training notebook first."
        }
    }
    with open("metrics.json", "w") as f:
        json.dump(metrics, f, indent=2)
    print("Placeholder metrics.json saved.")