#!/usr/bin/env python3
"""
ByT5 Optuna-Tuned Inference v1.0.0

Based on anthonytherrien/byt-ensemble-script (35.1 LB) with hyperparameters
optimized via local Optuna sweep (30 trials on training holdout).

Key changes from original:
  - num_beams: 8 -> 9
  - length_penalty: 1.3 -> 1.0
  - repetition_penalty: 1.25 (new, mild dedup)
  - gap_strategy: keep <gap>/<big_gap> markers (don't convert to ...)

Model: assiaben/final-byt5 (byt5-akkadian-optimized-34x)
"""

# ---------------------------
# 1) System-Level Optimization
# ---------------------------

import os

os.environ["OMP_NUM_THREADS"] = "4"
os.environ["MKL_NUM_THREADS"] = "4"
os.environ["CUDA_LAUNCH_BLOCKING"] = "0"
os.environ["TORCH_CUDNN_V8_API_ENABLED"] = "1"
os.environ["TOKENIZERS_PARALLELISM"] = "true"

# ---------------------------
# 2) Imports & Setup
# ---------------------------

import re
import json
import random
import logging
import warnings
from pathlib import Path
from dataclasses import dataclass
from typing import List

import pandas as pd
import numpy as np
import torch
from torch.utils.data import Dataset, DataLoader, Sampler
from torch.cuda.amp import autocast
from transformers import AutoTokenizer, AutoModelForSeq2SeqLM
from tqdm.auto import tqdm

warnings.filterwarnings("ignore")


# ---------------------------
# 3) Configuration (Optuna-tuned)
# ---------------------------


@dataclass
class UltraConfig:
    # ============ PATHS ============
    test_data_path: str = (
        "/kaggle/input/deep-past-initiative-machine-translation/test.csv"
    )
    model_path: str = "/kaggle/input/final-byt5/byt5-akkadian-optimized-34x"
    output_dir: str = "/kaggle/working/"

    # ============ PROCESSING ============
    max_length: int = 512
    batch_size: int = 8
    num_workers: int = 4

    # ============ GENERATION (Optuna-tuned) ============
    num_beams: int = 9  # was 8 -> +1 beam improves search
    max_new_tokens: int = 512
    length_penalty: float = 1.0  # was 1.3 -> neutral is better
    repetition_penalty: float = 1.25  # was 1.0 -> mild dedup helps
    early_stopping: bool = True
    no_repeat_ngram_size: int = 0  # confirmed: >0 always hurts

    # ============ OPTIMIZATIONS ============
    use_mixed_precision: bool = True
    use_better_transformer: bool = True
    use_bucket_batching: bool = True
    use_vectorized_postproc: bool = True
    use_adaptive_beams: bool = True
    use_auto_batch_size: bool = False

    # ============ OTHER ============
    aggressive_postprocessing: bool = True
    checkpoint_freq: int = 100
    num_buckets: int = 4

    def __post_init__(self):
        self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
        Path(self.output_dir).mkdir(exist_ok=True, parents=True)
        if not torch.cuda.is_available():
            self.use_mixed_precision = False
            self.use_better_transformer = False


# ---------------------------
# 4) Logging Setup
# ---------------------------


def setup_logging(output_dir: str) -> logging.Logger:
    Path(output_dir).mkdir(exist_ok=True, parents=True)
    log_file = Path(output_dir) / "inference_ultra.log"

    for handler in logging.root.handlers[:]:
        logging.root.removeHandler(handler)

    logging.basicConfig(
        level=logging.INFO,
        format="%(asctime)s - %(levelname)s - %(message)s",
        handlers=[
            logging.StreamHandler(),
            logging.FileHandler(log_file),
        ],
    )
    return logging.getLogger(__name__)


# ---------------------------
# 5) Optimized Text Preprocessor
# ---------------------------


class OptimizedPreprocessor:
    def __init__(self):
        self.patterns = {
            "big_gap": re.compile(r"(\.{3,}|…+|……)"),
            "small_gap": re.compile(r"(xx+|\s+x\s+)"),
        }

    def preprocess_batch(self, texts: List[str]) -> List[str]:
        s = pd.Series(texts).fillna("")
        s = s.astype(str)
        s = s.str.replace(self.patterns["big_gap"], "<big_gap>", regex=True)
        s = s.str.replace(self.patterns["small_gap"], "<gap>", regex=True)
        return s.tolist()


# ---------------------------
# 6) Vectorized Postprocessor
# ---------------------------


class VectorizedPostprocessor:
    def __init__(self, aggressive: bool = True):
        self.aggressive = aggressive

        self.patterns = {
            "gap": re.compile(r"(\[x\]|\(x\)|\bx\b)", re.I),
            "big_gap": re.compile(r"(\.{3,}|…|\[\.+\])"),
            "annotations": re.compile(
                r"\((fem|plur|pl|sing|singular|plural|\?|!)\..\s*\w*\)", re.I
            ),
            "repeated_words": re.compile(r"\b(\w+)(?:\s+\1\b)+"),
            "whitespace": re.compile(r"\s+"),
            "punct_space": re.compile(r"\s+([.,:])"),
            "repeated_punct": re.compile(r"([.,])\1+"),
        }

        self.subscript_trans = str.maketrans("₀₁₂₃₄₅₆₇₈₉", "0123456789")
        self.special_chars_trans = str.maketrans("ḫḪ", "hH")
        self.forbidden_chars = '!?()"——<>⌈⌋⌊[]+ʾ/;'
        self.forbidden_trans = str.maketrans("", "", self.forbidden_chars)

    def postprocess_batch(self, translations: List[str]) -> List[str]:
        s = pd.Series(translations)

        valid_mask = s.apply(lambda x: isinstance(x, str) and x.strip())
        if not valid_mask.all():
            s[~valid_mask] = ""

        # Basic cleaning
        s = s.str.translate(self.special_chars_trans)
        s = s.str.translate(self.subscript_trans)
        s = s.str.replace(self.patterns["whitespace"], " ", regex=True)
        s = s.str.strip()

        if self.aggressive:
            # Normalize gaps
            s = s.str.replace(self.patterns["gap"], "<gap>", regex=True)
            s = s.str.replace(self.patterns["big_gap"], "<big_gap>", regex=True)

            # Merge adjacent gaps
            s = s.str.replace("<gap> <gap>", "<big_gap>", regex=False)
            s = s.str.replace("<big_gap> <big_gap>", "<big_gap>", regex=False)

            # Remove annotations
            s = s.str.replace(self.patterns["annotations"], "", regex=True)

            # Protect gaps during forbidden-character removal
            s = s.str.replace("<gap>", "\x00GAP\x00", regex=False)
            s = s.str.replace("<big_gap>", "\x00BIG\x00", regex=False)
            s = s.str.translate(self.forbidden_trans)
            s = s.str.replace("\x00GAP\x00", " <gap> ", regex=False)
            s = s.str.replace("\x00BIG\x00", " <big_gap> ", regex=False)

            # Fractions
            s = s.str.replace(r"(\d+)\.5\b", r"\1½", regex=True)
            s = s.str.replace(r"\b0\.5\b", "½", regex=True)
            s = s.str.replace(r"(\d+)\.25\b", r"\1¼", regex=True)
            s = s.str.replace(r"\b0\.25\b", "¼", regex=True)
            s = s.str.replace(r"(\d+)\.75\b", r"\1¾", regex=True)
            s = s.str.replace(r"\b0\.75\b", "¾", regex=True)

            # Remove repeated words
            s = s.str.replace(self.patterns["repeated_words"], r"\1", regex=True)

            # Remove repeated n-grams (4 -> 2)
            for n in range(4, 1, -1):
                pattern = r"\b((?:\w+\s+){" + str(n - 1) + r"}\w+)(?:\s+\1\b)+"
                s = s.str.replace(pattern, r"\1", regex=True)

            # Fix punctuation spacing and repeats
            s = s.str.replace(self.patterns["punct_space"], r"\1", regex=True)
            s = s.str.replace(self.patterns["repeated_punct"], r"\1", regex=True)

            # Final whitespace cleanup
            s = s.str.replace(self.patterns["whitespace"], " ", regex=True)
            s = s.str.strip().str.strip("-").str.strip()

        return s.tolist()


# ---------------------------
# 7) Bucket Batch Sampler
# ---------------------------


class BucketBatchSampler(Sampler):
    def __init__(
        self,
        dataset: Dataset,
        batch_size: int,
        num_buckets: int,
        logger: logging.Logger,
        shuffle: bool = False,
    ):
        self.dataset = dataset
        self.batch_size = batch_size
        self.shuffle = shuffle
        self.logger = logger

        lengths = [len(text.split()) for _, text in dataset]
        sorted_indices = sorted(range(len(lengths)), key=lambda i: lengths[i])

        bucket_size = max(1, len(sorted_indices) // max(1, num_buckets))
        self.buckets = []

        for i in range(num_buckets):
            start = i * bucket_size
            end = None if i == num_buckets - 1 else (i + 1) * bucket_size
            self.buckets.append(sorted_indices[start:end])

        self.logger.info(f"Created {num_buckets} buckets:")
        for i, bucket in enumerate(self.buckets):
            bucket_lengths = (
                [lengths[idx] for idx in bucket] if len(bucket) > 0 else [0]
            )
            self.logger.info(
                f"  Bucket {i}: {len(bucket)} samples, "
                f"length range [{min(bucket_lengths)}, {max(bucket_lengths)}]"
            )

    def __iter__(self):
        for bucket in self.buckets:
            if self.shuffle:
                random.shuffle(bucket)
            for i in range(0, len(bucket), self.batch_size):
                yield bucket[i : i + self.batch_size]

    def __len__(self):
        total = 0
        for bucket in self.buckets:
            total += (len(bucket) + self.batch_size - 1) // self.batch_size
        return total


# ---------------------------
# 8) Dataset
# ---------------------------


class AkkadianDataset(Dataset):
    def __init__(
        self,
        dataframe: pd.DataFrame,
        preprocessor: OptimizedPreprocessor,
        logger: logging.Logger,
    ):
        self.sample_ids = dataframe["id"].tolist()
        raw_texts = dataframe["transliteration"].tolist()
        preprocessed = preprocessor.preprocess_batch(raw_texts)
        self.input_texts = [
            "translate Akkadian to English: " + text for text in preprocessed
        ]
        logger.info(f"Dataset created with {len(self.sample_ids)} samples")

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

    def __getitem__(self, index: int):
        return self.sample_ids[index], self.input_texts[index]


# ---------------------------
# 9) Inference Engine
# ---------------------------


class UltraInferenceEngine:
    def __init__(self, config: UltraConfig, logger: logging.Logger):
        self.config = config
        self.logger = logger
        self.preprocessor = OptimizedPreprocessor()
        self.postprocessor = VectorizedPostprocessor(
            aggressive=config.aggressive_postprocessing
        )
        self.results = []
        self._load_model()

    def _load_model(self):
        self.logger.info(f"Loading model from {self.config.model_path}")
        self.model = AutoModelForSeq2SeqLM.from_pretrained(self.config.model_path)
        self.model = self.model.to(self.config.device)
        self.model = self.model.eval()
        self.tokenizer = AutoTokenizer.from_pretrained(self.config.model_path)

        num_params = sum(p.numel() for p in self.model.parameters())
        self.logger.info(f"Model loaded: {num_params:,} parameters")

        if self.config.use_better_transformer and torch.cuda.is_available():
            try:
                from optimum.bettertransformer import BetterTransformer

                self.logger.info("Applying BetterTransformer...")
                self.model = BetterTransformer.transform(self.model)
                self.logger.info("BetterTransformer applied")
            except ImportError:
                self.logger.warning("optimum not installed, skipping BetterTransformer")
            except Exception as exc:
                self.logger.warning(f"BetterTransformer failed: {exc}")

    def _collate_fn(self, batch_samples):
        batch_ids = [s[0] for s in batch_samples]
        batch_texts = [s[1] for s in batch_samples]
        tokenized = self.tokenizer(
            batch_texts,
            max_length=self.config.max_length,
            padding=True,
            truncation=True,
            return_tensors="pt",
        )
        return batch_ids, tokenized

    def _get_adaptive_beam_size(self, attention_mask: torch.Tensor) -> int:
        if not self.config.use_adaptive_beams:
            return self.config.num_beams

        lengths = attention_mask.sum(dim=1)
        short_beams = max(4, self.config.num_beams // 2)
        beam_sizes = torch.where(
            lengths < 100,
            torch.tensor(short_beams, device=lengths.device),
            torch.tensor(self.config.num_beams, device=lengths.device),
        )
        return int(beam_sizes[0].item())

    def _save_checkpoint(self):
        if (
            len(self.results) > 0
            and len(self.results) % self.config.checkpoint_freq == 0
        ):
            checkpoint_path = (
                Path(self.config.output_dir) / f"checkpoint_{len(self.results)}.csv"
            )
            df = pd.DataFrame(self.results, columns=["id", "translation"])
            df.to_csv(checkpoint_path, index=False)
            self.logger.info(f"Checkpoint: {len(self.results)} translations")

    def run_inference(self, test_df: pd.DataFrame) -> pd.DataFrame:
        self.logger.info("Starting inference")

        dataset = AkkadianDataset(test_df, self.preprocessor, self.logger)

        if self.config.use_bucket_batching:
            batch_sampler = BucketBatchSampler(
                dataset=dataset,
                batch_size=self.config.batch_size,
                num_buckets=self.config.num_buckets,
                logger=self.logger,
                shuffle=False,
            )
            dataloader = DataLoader(
                dataset,
                batch_sampler=batch_sampler,
                num_workers=self.config.num_workers,
                collate_fn=self._collate_fn,
                pin_memory=True,
                prefetch_factor=2,
                persistent_workers=self.config.num_workers > 0,
            )
        else:
            dataloader = DataLoader(
                dataset,
                batch_size=self.config.batch_size,
                shuffle=False,
                num_workers=self.config.num_workers,
                collate_fn=self._collate_fn,
                pin_memory=True,
                prefetch_factor=2,
                persistent_workers=self.config.num_workers > 0,
            )

        self.logger.info(f"DataLoader created: {len(dataloader)} batches")
        self.logger.info(f"  Mixed Precision: {self.config.use_mixed_precision}")
        self.logger.info(f"  BetterTransformer: {self.config.use_better_transformer}")
        self.logger.info(f"  Bucket Batching: {self.config.use_bucket_batching}")
        self.logger.info(f"  Adaptive Beams: {self.config.use_adaptive_beams}")

        base_gen_config = {
            "max_new_tokens": self.config.max_new_tokens,
            "length_penalty": self.config.length_penalty,
            "repetition_penalty": self.config.repetition_penalty,
            "early_stopping": self.config.early_stopping,
            "use_cache": True,
        }

        if self.config.no_repeat_ngram_size > 0:
            base_gen_config["no_repeat_ngram_size"] = self.config.no_repeat_ngram_size

        self.results = []

        with torch.inference_mode():
            for batch_idx, (batch_ids, tokenized) in enumerate(
                tqdm(dataloader, desc="Translating")
            ):
                try:
                    input_ids = tokenized.input_ids.to(self.config.device)
                    attention_mask = tokenized.attention_mask.to(self.config.device)

                    beam_size = self._get_adaptive_beam_size(attention_mask)
                    gen_config = dict(base_gen_config)
                    gen_config["num_beams"] = beam_size

                    if self.config.use_mixed_precision:
                        with autocast():
                            outputs = self.model.generate(
                                input_ids=input_ids,
                                attention_mask=attention_mask,
                                **gen_config,
                            )
                    else:
                        outputs = self.model.generate(
                            input_ids=input_ids,
                            attention_mask=attention_mask,
                            **gen_config,
                        )

                    translations = self.tokenizer.batch_decode(
                        outputs, skip_special_tokens=True
                    )
                    cleaned = self.postprocessor.postprocess_batch(translations)
                    self.results.extend(list(zip(batch_ids, cleaned)))
                    self._save_checkpoint()

                    if torch.cuda.is_available() and batch_idx % 10 == 0:
                        torch.cuda.empty_cache()

                except Exception as exc:
                    self.logger.error(f"Batch {batch_idx} error: {exc}")
                    self.results.extend([(bid, "") for bid in batch_ids])

        self.logger.info("Inference completed")

        results_df = pd.DataFrame(self.results, columns=["id", "translation"])
        self._validate_results(results_df)
        return results_df

    def _validate_results(self, df: pd.DataFrame):
        print("\n" + "=" * 60)
        print("VALIDATION REPORT")
        print("=" * 60)

        empty = df["translation"].astype(str).str.strip().eq("").sum()
        print(f"\nEmpty: {empty} ({(empty / max(1, len(df))) * 100:.2f}%)")

        lengths = df["translation"].astype(str).str.len()
        print(f"\nLength stats:")
        print(f"   Mean: {lengths.mean():.1f}, Median: {lengths.median():.1f}")
        print(f"   Min: {lengths.min()}, Max: {lengths.max()}")

        short = ((lengths < 5) & (lengths > 0)).sum()
        if short > 0:
            print(f"   {short} very short translations")

        print(f"\nSample translations:")
        sample_indices = [0]
        if len(df) > 2:
            sample_indices.append(len(df) // 2)
        if len(df) > 1:
            sample_indices.append(len(df) - 1)
        for idx in sample_indices:
            row = df.iloc[idx]
            text = str(row["translation"])
            preview = text[:70] + "..." if len(text) > 70 else text
            print(f"   ID {int(row['id']):4d}: {preview}")

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


# ---------------------------
# 10) IO Helpers
# ---------------------------


def print_environment_info():
    print(f"PyTorch version: {torch.__version__}")
    print(f"CUDA available: {torch.cuda.is_available()}")

    if torch.cuda.is_available():
        print(f"GPU: {torch.cuda.get_device_name(0)}")
        total_mem_gb = torch.cuda.get_device_properties(0).total_memory / 1e9
        print(f"GPU Memory: {total_mem_gb:.2f} GB")

    try:
        from optimum.bettertransformer import BetterTransformer  # noqa: F401

        print("BetterTransformer available")
    except ImportError:
        print("BetterTransformer NOT available")


def save_outputs(
    results_df: pd.DataFrame,
    config: UltraConfig,
    output_dir: str,
    logger: logging.Logger,
):
    output_path = Path(output_dir) / "submission.csv"
    config_path = Path(output_dir) / "ultra_config.json"

    results_df.to_csv(output_path, index=False)
    logger.info(f"Submission saved to {output_path}")

    config_dict = {
        "batch_size": config.batch_size,
        "num_beams": config.num_beams,
        "length_penalty": config.length_penalty,
        "repetition_penalty": config.repetition_penalty,
        "no_repeat_ngram_size": config.no_repeat_ngram_size,
        "optimizations": {
            "mixed_precision": config.use_mixed_precision,
            "better_transformer": config.use_better_transformer,
            "bucket_batching": config.use_bucket_batching,
            "vectorized_postproc": config.use_vectorized_postproc,
            "adaptive_beams": config.use_adaptive_beams,
        },
    }

    with open(config_path, "w", encoding="utf-8") as f:
        json.dump(config_dict, f, indent=2)

    print("\n" + "=" * 60)
    print("INFERENCE COMPLETE")
    print("=" * 60)
    print(f"Submission file: {output_path}")
    print(f"Config file: {config_path}")
    print(f"Total translations: {len(results_df)}")
    print("=" * 60)


# ---------------------------
# 11) Main
# ---------------------------


def main():
    config = UltraConfig()

    logger = setup_logging(config.output_dir)
    logger.info("Logging initialized")

    print_environment_info()

    logger.info("Configuration (Optuna-tuned):")
    logger.info(f"  Device: {config.device}")
    logger.info(f"  Batch size: {config.batch_size}")
    logger.info(f"  Beams: {config.num_beams}")
    logger.info(f"  Length penalty: {config.length_penalty}")
    logger.info(f"  Repetition penalty: {config.repetition_penalty}")
    logger.info(f"  Mixed Precision: {config.use_mixed_precision}")
    logger.info(f"  BetterTransformer: {config.use_better_transformer}")
    logger.info(f"  Bucket Batching: {config.use_bucket_batching}")
    logger.info(f"  Adaptive Beams: {config.use_adaptive_beams}")

    logger.info(f"Loading test data from {config.test_data_path}")
    test_df = pd.read_csv(config.test_data_path, encoding="utf-8")
    logger.info(f"Loaded {len(test_df)} test samples")

    print("\nFirst 5 samples:")
    print(test_df.head())

    engine = UltraInferenceEngine(config, logger)
    results_df = engine.run_inference(test_df)
    save_outputs(results_df, config, config.output_dir, logger)


if __name__ == "__main__":
    main()