"""
ByT5 Ensemble v1.0.3 - Weight Averaged Ensemble (Original Params)
=================================================================
Based on jiexusheng20bz's 33.9 kernel - EXACT ORIGINAL PARAMETERS.

EXPERIMENT FAILED (v1-0-1/v1-0-2):
---------------------------------
Attempted optimizations that DEGRADED score from 33.9 → 19.4:
- num_beams: 4 → 5 (REVERTED)
- repetition_penalty: 1.1 (REMOVED)
- no_repeat_ngram_size: 3 (REMOVED)
- Slight weight adjustments (REVERTED)

LESSON: For ensemble weight averaging, stick with original generation params.
The averaged model is sensitive to decoding strategy changes.

v1-0-3: Reverted to EXACT original params. Documentation only, no functional changes.

ENSEMBLE APPROACH:
-----------------
Weight averaging combines 3 ByT5 models by averaging their parameters.
This produces a SINGLE model, reducing inference cost by ~3x vs output ensembling.

MODELS:
------
1. byt5-base-big-data2 (jeanjean111) - weight 0.98
   - Trained on large diverse Akkadian corpus
   - Strong general translation capability

2. byt5-akkadian-model (llkh0a) - weight 1.00
   - Different training data/hyperparameters
   - Provides model diversity

3. byt5-base-akkadian_gap_setence2 (qifeihhh666) - weight 0.55
   - Specialized for handling gap markers (damaged text)
   - Lower weight due to specialized nature

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

import re
import pandas as pd
import torch
from torch.utils.data import DataLoader, Dataset
from transformers import AutoTokenizer, AutoModelForSeq2SeqLM
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)}")

# ============================================================
# CONFIGURATION - ORIGINAL PARAMETERS (DO NOT MODIFY)
# ============================================================

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"

TEST_DATA_PATH = "/kaggle/input/deep-past-initiative-machine-translation/test.csv"

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

# Original 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"Weights: w1={w1:.3f}, w2={w2:.3f}, w3={w3:.3f}")


# ============================================================
# GAP PREPROCESSING
# ============================================================


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


# ============================================================
# LOAD AND AVERAGE MODELS
# ============================================================

print("\nLoading Model 1...")
model1 = AutoModelForSeq2SeqLM.from_pretrained(MODEL1_PATH)
sd1 = model1.state_dict()

print("Loading Model 2...")
model2 = AutoModelForSeq2SeqLM.from_pretrained(MODEL2_PATH)
sd2 = model2.state_dict()

print("Loading Model 3...")
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
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"Model dtype: {next(model.parameters()).dtype}")

# ============================================================
# 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)}")


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

PREFIX = "translate Akkadian to English: "


class InferenceDataset(Dataset):
    def __init__(self, df, tokenizer):
        self.texts = [PREFIX + str(t) for t in df["transliteration"].tolist()]
        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),
        }


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

print("\nStarting Inference...")
all_predictions = []

torch.set_grad_enabled(False)

with torch.inference_mode():
    for batch in tqdm(test_loader, desc="Predicting"):
        input_ids = batch["input_ids"].to(DEVICE)
        attention_mask = batch["attention_mask"].to(DEVICE)

        # ORIGINAL GENERATION PARAMS - DO NOT MODIFY
        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])

# ============================================================
# CREATE SUBMISSION
# ============================================================

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

# Handle empty predictions
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"\nSaved {len(submission)} predictions to submission.csv")
print(submission.head())