"""
ByT5 Ensemble v1.0.2 - Optimized Weight Averaging
==================================================
Based on jiexusheng20bz's 33.9 kernel with analysis and optimization.

ENSEMBLE RATIONALE:
-------------------
Weight averaging combines multiple models by averaging their parameters.
Unlike output ensembling, this produces a SINGLE model for inference,
reducing memory and compute by ~3x while capturing diverse training signals.

MODEL ANALYSIS:
--------------
1. byt5-base-big-data2 (jeanjean111) - Score: ~32.6 alone
   - Trained on large diverse Akkadian corpus
   - Strong general translation capability
   - May overfit to common patterns

2. byt5-akkadian-model (llkh0a) - Score: unknown
   - Different training data/hyperparameters
   - Provides model diversity for ensemble
   - Likely captures different linguistic patterns

3. byt5-base-akkadian_gap_setence2 (qifeihhh666) - Score: unknown
   - Specifically trained to handle gap markers (damaged text)
   - Important for real-world cuneiform translation
   - Lower weight due to specialized nature

WEIGHT DERIVATION:
-----------------
Original weights (perf1=0.98, perf2=1.00, perf3=0.55) appear to be
rough estimates based on validation performance. These can be
optimized using:
1. Grid search over weight combinations
2. Optuna Bayesian optimization
3. Cross-validation on holdout set

TODO: Run Optuna optimization locally with:
- Objective: maximize BLEU/chrF on validation set
- Search space: w1, w2, w3 in [0.1, 2.0]
- Constraint: w1 + w2 + w3 = constant (normalized)

Built with Akkadian CLI - https://github.com/manwithacat/akkadian
AI-first tooling for ML Competitions on Kaggle
"""

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
# ============================================================

# Model paths
MODEL1_PATH = "/kaggle/input/byt5-base-big-data2"  # jeanjean's 32.6 model
MODEL2_PATH = "/kaggle/input/byt5-akkadian-model"  # llkh0a's model
MODEL3_PATH = (
    "/kaggle/input/train-gap-all-2/byt5-base-akkadian_gap_setence2"  # gap-trained
)

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

# Generation params - optimized from baseline
BATCH_SIZE = 8
MAX_LENGTH = 512
NUM_BEAMS = 5  # Increased from 4 for better search
EARLY_STOPPING = True
REPETITION_PENALTY = 1.1  # Light penalty to prevent loops
NO_REPEAT_NGRAM_SIZE = 3  # Prevent 3-gram repetition

# ============================================================
# ENSEMBLE WEIGHTS - TO BE OPTIMIZED
# ============================================================
# These weights determine model contribution to the ensemble.
# Higher weight = more influence on final model parameters.
#
# Current values based on estimated relative performance:
# - Model 1 (big-data2): Strong baseline, well-generalized
# - Model 2 (akkadian): Similar architecture, provides diversity
# - Model 3 (gap): Specialized for damaged text, lower weight

WEIGHT_MODEL1 = 0.95  # jeanjean's 32.6 model
WEIGHT_MODEL2 = 1.05  # llkh0a's model (slightly boosted)
WEIGHT_MODEL3 = 0.50  # gap-specialized model

# Normalize weights
TOTAL_WEIGHT = WEIGHT_MODEL1 + WEIGHT_MODEL2 + WEIGHT_MODEL3
W1 = WEIGHT_MODEL1 / TOTAL_WEIGHT
W2 = WEIGHT_MODEL2 / TOTAL_WEIGHT
W3 = WEIGHT_MODEL3 / TOTAL_WEIGHT

DEVICE = "cuda" if torch.cuda.is_available() else "cpu"


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


def replace_gaps(text):
    """
    Normalize gap markers in transliteration.

    Gap markers indicate damaged/missing text in cuneiform tablets.
    Standardizing these helps the model handle them consistently.

    Patterns:
    - 'x' alone → <gap> (single sign missing)
    - 'x x x' or 'xx' → <big_gap> (multiple signs missing)
    - '...' or '…' → <big_gap> (extended damage)
    """
    if pd.isna(text):
        return text

    # Multiple x's or ellipsis → big_gap
    text = re.sub(r"\bx\s+x\s+x\b", "<big_gap>", text)
    text = re.sub(r"\bxx+\b", "<big_gap>", text)
    text = re.sub(r"\.{3,}", "<big_gap>", text)
    text = re.sub(r"…+", "<big_gap>", text)

    # Single x → gap
    text = re.sub(r"\bx\b", "<gap>", text)

    # Clean up consecutive gaps
    text = re.sub(r"(<gap>\s*){2,}", "<big_gap> ", text)
    text = re.sub(r"<gap>\s*<big_gap>", "<big_gap>", text)
    text = re.sub(r"<big_gap>\s*<gap>", "<big_gap>", text)

    return text.strip()


# ============================================================
# POST-PROCESSING
# ============================================================


def postprocess(text):
    """Clean up output artifacts for submission."""
    if not text or pd.isna(text):
        return "broken text"

    text = str(text).strip()

    # Remove repeated punctuation
    text = re.sub(r"([.!?])\1+", r"\1", text)

    # Fix spacing around punctuation
    text = re.sub(r"\s+([.,;:!?])", r"\1", text)
    text = re.sub(r"\s+", " ", text)

    # Capitalize first letter
    if text and text[0].islower():
        text = text[0].upper() + text[1:]

    return text.strip() if text else "broken text"


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

print("\n" + "=" * 60)
print("ENSEMBLE CONFIGURATION")
print("=" * 60)
print(f"Model 1 weight: {W1:.3f} ({WEIGHT_MODEL1:.2f} raw)")
print(f"Model 2 weight: {W2:.3f} ({WEIGHT_MODEL2:.2f} raw)")
print(f"Model 3 weight: {W3:.3f} ({WEIGHT_MODEL3:.2f} raw)")

print("\n" + "=" * 60)
print("LOADING MODELS")
print("=" * 60)

print("\nLoading Model 1 (byt5-base-big-data2)...")
model1 = AutoModelForSeq2SeqLM.from_pretrained(MODEL1_PATH)
sd1 = model1.state_dict()
print(f"  Parameters: {sum(p.numel() for p in model1.parameters()):,}")

print("Loading Model 2 (byt5-akkadian-model)...")
model2 = AutoModelForSeq2SeqLM.from_pretrained(MODEL2_PATH)
sd2 = model2.state_dict()
print(f"  Parameters: {sum(p.numel() for p in model2.parameters()):,}")

print("Loading Model 3 (gap-specialized)...")
model3 = AutoModelForSeq2SeqLM.from_pretrained(MODEL3_PATH)
sd3 = model3.state_dict()
print(f"  Parameters: {sum(p.numel() for p in model3.parameters()):,}")

print("\nAveraging model weights...")
final_sd = sd2.copy()

matched_keys = 0
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]
        matched_keys += 1
    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]

print(f"  Averaged {matched_keys} parameter tensors")

# Free memory
del model1, model2, model3, sd1, sd2, sd3
torch.cuda.empty_cache() if torch.cuda.is_available() else None

# Load averaged 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"\nEnsemble model ready: {sum(p.numel() for p in model.parameters()):,} parameters"
)


# ============================================================
# LOAD AND PREPROCESS DATA
# ============================================================

print("\n" + "=" * 60)
print("LOADING DATA")
print("=" * 60)

test_df = pd.read_csv(TEST_DATA_PATH)
print(f"Test samples: {len(test_df)}")

test_df["transliteration_processed"] = test_df["transliteration"].apply(replace_gaps)

print("\nPreprocessing examples:")
for i in range(min(3, len(test_df))):
    orig = test_df["transliteration"].iloc[i][:60]
    proc = test_df["transliteration_processed"].iloc[i][:60]
    print(f"  [{i + 1}] {orig}...")
    print(f"      → {proc}...")


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

PREFIX = "translate Akkadian to English: "


class InferenceDataset(Dataset):
    def __init__(self, df, tokenizer):
        self.texts = [PREFIX + str(t) for t in df["transliteration_processed"].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("\n" + "=" * 60)
print("GENERATING TRANSLATIONS")
print(
    f"num_beams={NUM_BEAMS}, rep_penalty={REPETITION_PENALTY}, no_repeat_ngram={NO_REPEAT_NGRAM_SIZE}"
)
print("=" * 60)

all_predictions = []

torch.set_grad_enabled(False)

with torch.inference_mode():
    for batch in tqdm(test_loader, desc="Translating"):
        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,
            early_stopping=EARLY_STOPPING,
            repetition_penalty=REPETITION_PENALTY,
            no_repeat_ngram_size=NO_REPEAT_NGRAM_SIZE,
        )

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


# ============================================================
# SAMPLE OUTPUTS
# ============================================================

print("\n" + "=" * 60)
print("SAMPLE TRANSLATIONS")
print("=" * 60)

for i in range(min(5, len(all_predictions))):
    src = test_df["transliteration"].iloc[i][:50]
    tgt = all_predictions[i][:80]
    print(f"\n[{i + 1}] {src}...")
    print(f"    → {tgt}")


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

submission = pd.DataFrame(
    {
        "id": test_df["id"],
        "translation": all_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"\nSubmission saved: {len(submission)} rows")
print(submission.head())

print("\n" + "=" * 60)
print("INFERENCE COMPLETE")
print(f"Ensemble weights: {W1:.3f}, {W2:.3f}, {W3:.3f}")
print(f"Generation: beams={NUM_BEAMS}, rep_penalty={REPETITION_PENALTY}")
print("=" * 60)