"""
Two-Tier Akkadian Translation Pipeline v1.0.2

Tier 1: 33.9 Weight-Averaged ByT5 Ensemble (base translator)
Tier 2: Flan-T5 Philologist (refinement)

HYPOTHESIS:
-----------
The 33.9 ensemble produces strong raw translations but may have stylistic
inconsistencies. The philologist, trained to map raw MT output to gold
reference style, should improve fluency and accuracy.

TIER 1 MODELS (weight-averaged):
--------------------------------
1. byt5-base-big-data2 (jeanjean111) - weight 0.98
2. byt5-akkadian-model (llkh0a) - weight 1.00
3. byt5-base-akkadian_gap_setence2 (qifeihhh666) - weight 0.55

TIER 2 MODEL:
-------------
- akkadian-flan-t5-philologist (trained to refine ByT5 output)

MEMORY STRATEGY:
----------------
Models are loaded and freed sequentially to fit in P100 16GB:
1. Load 3 ByT5 models, average weights, delete originals (~6GB peak)
2. Run Tier 1 inference with averaged model (~2.2GB)
3. Delete ensemble, load philologist (~1GB)
4. Run Tier 2 inference

Note: Preflight warns about 22.5GB but that assumes parallel loading.
Sequential loading with gc keeps peak VRAM under 8GB.

Built with Akkadian CLI - https://github.com/manwithacat/akkadian
"""

import os
import re
import gc

# Suppress progress bars
os.environ["TQDM_DISABLE"] = "1"
os.environ["HF_HUB_DISABLE_PROGRESS_BARS"] = "1"

import pandas as pd
import torch
from torch.utils.data import DataLoader, Dataset
from transformers import (
    AutoTokenizer,
    AutoModelForSeq2SeqLM,
    T5Tokenizer,
    T5ForConditionalGeneration,
)
from tqdm.auto import tqdm

print(f"PyTorch: {torch.__version__}")
print(f"CUDA: {torch.cuda.is_available()}")
if torch.cuda.is_available():
    print(f"GPU: {torch.cuda.get_device_name(0)}")
    torch.cuda.empty_cache()

# ============================================================
# CONFIGURATION
# ============================================================

# Tier 1: ByT5 Ensemble Models
MODEL1_PATH = "/kaggle/input/byt5-base-big-data2"
MODEL2_PATH = "/kaggle/input/byt5-akkadian-model"
MODEL3_PATH = "/kaggle/input/train-gap-all-2/byt5-base-akkadian_gap_setence2"

# Tier 2: Philologist Model (Kaggle Model format: model-slug/framework/instance/version)
PHILOLOGIST_PATH = (
    "/kaggle/input/akkadian-flan-t5-philologist-v1-0-0/pytorch/transformers/1"
)

# Data
TEST_DATA_PATH = "/kaggle/input/deep-past-initiative-machine-translation/test.csv"
OUTPUT_DIR = "/kaggle/working"

# Generation params
BATCH_SIZE = 8
MAX_LENGTH = 512
DEVICE = "cuda" if torch.cuda.is_available() else "cpu"

# Original ensemble weights from 33.9 kernel
perf1 = 0.98
perf2 = 1.00
perf3 = 0.55

total = perf1 + perf2 + perf3
w1 = perf1 / total
w2 = perf2 / total
w3 = perf3 / total

print(f"\nEnsemble weights: w1={w1:.3f}, w2={w2:.3f}, w3={w3:.3f}")


# ============================================================
# GAP PREPROCESSING (from 33.9 kernel)
# ============================================================


def replace_gaps(text):
    """Normalize gap markers in transliteration."""
    if pd.isna(text):
        return text
    text = re.sub(r"\.{3,}|\.{2,}|x+x+|\bx\b|\s+x\s+", "<big_gap>", text)
    text = re.sub(r"\bx\b", "<gap>", text)
    return text


# ============================================================
# TIER 1: LOAD AND AVERAGE ENSEMBLE MODELS
# ============================================================


def load_tier1_ensemble():
    """Load and weight-average the 3 ByT5 models."""
    print("\n" + "=" * 60)
    print("TIER 1: Loading ByT5 Ensemble")
    print("=" * 60)

    print("Loading Model 1 (byt5-base-big-data2)...")
    model1 = AutoModelForSeq2SeqLM.from_pretrained(MODEL1_PATH)
    sd1 = model1.state_dict()

    print("Loading Model 2 (byt5-akkadian-model)...")
    model2 = AutoModelForSeq2SeqLM.from_pretrained(MODEL2_PATH)
    sd2 = model2.state_dict()

    print("Loading Model 3 (byt5-base-akkadian_gap_setence2)...")
    model3 = AutoModelForSeq2SeqLM.from_pretrained(MODEL3_PATH)
    sd3 = model3.state_dict()

    print("Averaging weights...")
    final_sd = sd2.copy()

    for key in final_sd:
        if key in sd1 and key in sd3:
            final_sd[key] = w1 * sd1[key] + w2 * sd2[key] + w3 * sd3[key]
        elif key in sd1:
            final_sd[key] = w1 * sd1[key] + (w2 + w3) * sd2[key]
        elif key in sd3:
            final_sd[key] = w3 * sd3[key] + (w1 + w2) * sd2[key]

    # Free memory
    del model1, model2, model3, sd1, sd3
    gc.collect()
    if torch.cuda.is_available():
        torch.cuda.empty_cache()

    # Load final model
    model = AutoModelForSeq2SeqLM.from_pretrained(MODEL2_PATH)
    model.load_state_dict(final_sd)
    model = model.to(DEVICE)
    model.eval()
    model.float()

    tokenizer = AutoTokenizer.from_pretrained(MODEL2_PATH)

    print(f"Ensemble loaded - dtype: {next(model.parameters()).dtype}")
    print(f"Parameters: {sum(p.numel() for p in model.parameters()):,}")

    return model, tokenizer


# ============================================================
# TIER 2: LOAD PHILOLOGIST MODEL
# ============================================================


def load_tier2_philologist():
    """Load the Flan-T5 philologist refinement model."""
    print("\n" + "=" * 60)
    print("TIER 2: Loading Philologist")
    print("=" * 60)

    print(f"Loading from {PHILOLOGIST_PATH}...")
    # Use T5-specific classes + local_files_only to bypass HF repo ID validation
    tokenizer = T5Tokenizer.from_pretrained(PHILOLOGIST_PATH, local_files_only=True)
    model = T5ForConditionalGeneration.from_pretrained(
        PHILOLOGIST_PATH,
        torch_dtype=torch.float32,
        local_files_only=True,
    ).to(DEVICE)
    model.eval()

    print(f"Philologist loaded - dtype: {next(model.parameters()).dtype}")
    print(f"Parameters: {sum(p.numel() for p in model.parameters()):,}")

    return model, tokenizer


# ============================================================
# INFERENCE DATASET
# ============================================================


class InferenceDataset(Dataset):
    def __init__(self, texts, tokenizer, prefix=""):
        self.texts = [prefix + str(t) for t in texts]
        self.tokenizer = tokenizer

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

    def __getitem__(self, idx):
        inputs = self.tokenizer(
            self.texts[idx],
            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),
        }


# ============================================================
# TIER 1: ENSEMBLE TRANSLATION
# ============================================================


def tier1_translate(model, tokenizer, test_df):
    """Generate raw translations using the weight-averaged ensemble."""
    print("\n" + "=" * 60)
    print("TIER 1: Generating Raw Translations")
    print("=" * 60)

    prefix = "translate Akkadian to English: "
    texts = test_df["transliteration"].tolist()

    dataset = InferenceDataset(texts, tokenizer, prefix=prefix)
    loader = DataLoader(dataset, batch_size=BATCH_SIZE, shuffle=False)

    all_predictions = []
    torch.set_grad_enabled(False)

    with torch.inference_mode():
        for batch in tqdm(loader, desc="Tier 1 (Ensemble)"):
            input_ids = batch["input_ids"].to(DEVICE)
            attention_mask = batch["attention_mask"].to(DEVICE)

            # ORIGINAL GENERATION PARAMS from 33.9 kernel
            outputs = model.generate(
                input_ids=input_ids,
                attention_mask=attention_mask,
                max_length=MAX_LENGTH,
                num_beams=4,
                early_stopping=True,
            )

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

    print(f"Tier 1 complete: {len(all_predictions)} translations")
    return all_predictions


# ============================================================
# TIER 2: PHILOLOGIST REFINEMENT
# ============================================================


def tier2_refine(model, tokenizer, raw_translations):
    """Refine translations using the philologist model."""
    print("\n" + "=" * 60)
    print("TIER 2: Refining Translations")
    print("=" * 60)

    prefix = "Refine Akkadian translation: "

    dataset = InferenceDataset(raw_translations, tokenizer, prefix=prefix)
    loader = DataLoader(
        dataset, batch_size=BATCH_SIZE * 2, shuffle=False
    )  # Philologist is smaller

    refined = []
    torch.set_grad_enabled(False)

    with torch.inference_mode():
        for batch in tqdm(loader, desc="Tier 2 (Philologist)"):
            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=4,
                early_stopping=True,
                length_penalty=1.0,
            )

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

    print(f"Tier 2 complete: {len(refined)} translations refined")
    return refined


# ============================================================
# MAIN PIPELINE
# ============================================================


def main():
    print("=" * 60)
    print("Two-Tier Akkadian Translation Pipeline v1.0.2")
    print("=" * 60)

    # Load test data
    test_df = pd.read_csv(TEST_DATA_PATH)
    test_df["transliteration"] = test_df["transliteration"].apply(replace_gaps)
    print(f"\nTest samples: {len(test_df)}")

    # TIER 1: Load ensemble and generate raw translations
    ensemble_model, ensemble_tokenizer = load_tier1_ensemble()
    raw_translations = tier1_translate(ensemble_model, ensemble_tokenizer, test_df)

    # Free ensemble memory before loading philologist
    del ensemble_model, ensemble_tokenizer
    gc.collect()
    if torch.cuda.is_available():
        torch.cuda.empty_cache()
    print("\nFreed ensemble memory for Tier 2")

    # TIER 2: Load philologist and refine translations
    philologist_model, philologist_tokenizer = load_tier2_philologist()
    final_translations = tier2_refine(
        philologist_model, philologist_tokenizer, raw_translations
    )

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

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

    output_path = os.path.join(OUTPUT_DIR, "submission.csv")
    submission.to_csv(output_path, index=False)
    print(f"\nSubmission saved to {output_path}")
    print(f"File size: {os.path.getsize(output_path):,} bytes")

    # Show comparison examples
    print("\n" + "=" * 60)
    print("Sample Comparison (Raw vs Refined):")
    print("=" * 60)
    for i in range(min(5, len(test_df))):
        print(f"\n[{i + 1}] Source: {test_df.iloc[i]['transliteration'][:80]}...")
        print(f"    Tier 1 (Raw):     {raw_translations[i][:80]}...")
        print(f"    Tier 2 (Refined): {final_translations[i][:80]}...")

    print("\n" + "=" * 60)
    print("TWO-TIER INFERENCE COMPLETE")
    print("=" * 60)


if __name__ == "__main__":
    main()