# ============================================================
# Akkadian MT: Ensemble + Data Augmentation Pipeline
# ============================================================
# Improvements over base notebook:
# - Model Ensemble with self-ensembling
# - Lexicon-based data augmentation
# - Enhanced MBR with diversity
# - Test-Time Augmentation (TTA)
# - OFFLINE MODE compatible
# ============================================================

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"

import re
import json
import random
import logging
import warnings
from pathlib import Path
from dataclasses import dataclass, field
from typing import List, Dict, Tuple, Optional
from collections import Counter
from contextlib import nullcontext

import pandas as pd
import numpy as np
import torch
from torch.utils.data import Dataset, DataLoader
from transformers import AutoTokenizer, AutoModelForSeq2SeqLM
from tqdm.auto import tqdm
import sacrebleu

warnings.filterwarnings("ignore")

# Set seeds
def set_seed(seed=42):
    random.seed(seed)
    np.random.seed(seed)
    torch.manual_seed(seed)
    torch.cuda.manual_seed_all(seed)

set_seed(42)

# Setup logging
logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')
logger = logging.getLogger(__name__)

# ============================================================
# Configuration
# ============================================================

@dataclass
class EnsembleConfig:
    test_data_path: str = "/kaggle/input/deep-past-initiative-machine-translation/test.csv"
    lexicon_path: str = "/kaggle/input/deep-past-initiative-machine-translation/OA_Lexicon_eBL.csv"
    dictionary_path: str = "/kaggle/input/deep-past-initiative-machine-translation/eBL_Dictionary.csv"
    model_path: str = "/kaggle/input/models/mattiaangeli/byt5-akkadian-mbr/pytorch/default/6"
    output_dir: str = "/kaggle/working/"
    
    # Ensemble settings
    use_self_ensemble: bool = True
    self_ensemble_size: int = 3
    ensemble_temperatures: List[float] = field(default_factory=lambda: [0.6, 0.8, 1.0])
    
    # Data augmentation
    use_lexicon_hints: bool = True
    use_tta: bool = True
    tta_variants: int = 2
    
    # Enhanced MBR
    use_mbr: bool = True
    mbr_num_beam_cands: int = 6
    mbr_num_sample_cands: int = 4
    mbr_pool_cap: int = 48
    mbr_diversity_penalty: float = 0.5
    
    # Generation
    max_length: int = 512
    batch_size: int = 1
    num_beams: int = 12
    max_new_tokens: int = 512
    length_penalty: float = 1.3
    repetition_penalty: float = 1.15
    
    # Optimizations
    use_bf16: bool = True
    use_better_transformer: bool = True
    
    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_bf16 = False
            self.use_better_transformer = False

config = EnsembleConfig()
logger.info(f"Device: {config.device}")

# ============================================================
# Lexicon Loader for Hints
# ============================================================

class LexiconHelper:
    def __init__(self, lexicon_path: str, dict_path: str):
        self.akka_to_eng = {}
        self._load(lexicon_path)
        self._load(dict_path)
        logger.info(f"Loaded {len(self.akka_to_eng)} lexicon entries")
    
    def _load(self, path: str):
        if not os.path.exists(path):
            return
        try:
            df = pd.read_csv(path, encoding='utf-8')
            akk_col = next((c for c in df.columns if 'akk' in c.lower() or 'word' in c.lower()), None)
            eng_col = next((c for c in df.columns if 'eng' in c.lower() or 'mean' in c.lower()), None)
            
            if akk_col and eng_col:
                for _, row in df.iterrows():
                    if pd.notna(row[akk_col]) and pd.notna(row[eng_col]):
                        akka = str(row[akk_col]).strip().lower()
                        eng = str(row[eng_col]).strip()
                        if akka and eng:
                            if akka not in self.akka_to_eng:
                                self.akka_to_eng[akka] = []
                            self.akka_to_eng[akka].extend([m.strip() for m in eng.split(',')])
        except Exception as e:
            logger.warning(f"Failed to load {path}: {e}")
    
    def get_hints(self, text: str) -> List[str]:
        hints = []
        for word in text.lower().split():
            clean = re.sub(r'[^\w]', '', word)
            if clean in self.akka_to_eng:
                hints.extend(self.akka_to_eng[clean][:2])
        return hints

lexicon = LexiconHelper(config.lexicon_path, config.dictionary_path)

# ============================================================
# Preprocessing
# ============================================================

class Preprocessor:
    def __init__(self):
        self.patterns = {
            "big_gap": re.compile(r"(\.\.\.+|…+)"),
            "small_gap": re.compile(r"(xx+|\s+x\s+)"),
        }
    
    def preprocess(self, text: str) -> str:
        if pd.isna(text):
            return ""
        text = str(text)
        text = self.patterns["big_gap"].sub("<big_gap>", text)
        text = self.patterns["small_gap"].sub("<gap>", text)
        return f"translate Akkadian to English: {text}"
    
    def tta_variants(self, text: str) -> List[str]:
        if not config.use_tta:
            return [text]
        variants = [text]
        if config.tta_variants >= 2:
            v1 = text.replace("<gap>", "[x]").replace("<big_gap>", "[...]")
            variants.append(v1)
        return variants[:config.tta_variants]

preprocessor = Preprocessor()

# ============================================================
# Postprocessing
# ============================================================

class Postprocessor:
    def __init__(self):
        self.forbidden = str.maketrans("", "", '!?()\"——<>⌈⌋⌊[]+ʾ/;')
        self.subscripts = str.maketrans("₀₁₂₃₄₅₆₇₈₉", "0123456789")
    
    def postprocess(self, texts: List[str]) -> List[str]:
        results = []
        for text in texts:
            text = str(text).translate(self.subscripts)
            text = text.translate(self.forbidden)
            text = re.sub(r'\s+', ' ', text).strip()
            # Deduplicate repeated words
            text = re.sub(r'\b(\w+)(?:\s+\1\b)+', r'\1', text)
            results.append(text)
        return results

postprocessor = Postprocessor()

# ============================================================
# Enhanced MBR
# ============================================================

class MBRSelector:
    def __init__(self):
        self.chrf = sacrebleu.metrics.CHRF(word_order=2)
    
    def similarity(self, a: str, b: str) -> float:
        if not a or not b:
            return 0.0
        return float(self.chrf.sentence_score(a, [b]).score)
    
    def select(self, candidates: List[str]) -> str:
        candidates = list(dict.fromkeys([c.strip() for c in candidates if c.strip()]))
        
        if len(candidates) <= 1:
            return candidates[0] if candidates else ""
        
        # Cap pool
        if len(candidates) > config.mbr_pool_cap:
            candidates = candidates[:config.mbr_pool_cap]
        
        best_idx, best_score = 0, -1e9
        for i in range(len(candidates)):
            scores = []
            for j in range(len(candidates)):
                if i != j:
                    sim = self.similarity(candidates[i], candidates[j])
                    if sim > 95:
                        sim *= (1 - config.mbr_diversity_penalty)
                    scores.append(sim)
            avg = sum(scores) / len(scores) if scores else 0
            if avg > best_score:
                best_score = avg
                best_idx = i
        
        return candidates[best_idx]

mbr = MBRSelector()

# ============================================================
# Ensemble Inference
# ============================================================

class EnsembleEngine:
    def __init__(self):
        self.model = AutoModelForSeq2SeqLM.from_pretrained(config.model_path)
        self.model = self.model.to(config.device).eval()
        self.tokenizer = AutoTokenizer.from_pretrained(config.model_path)
        
        if config.use_better_transformer and torch.cuda.is_available():
            try:
                from optimum.bettertransformer import BetterTransformer
                self.model = BetterTransformer.transform(self.model)
                logger.info("BetterTransformer applied")
            except:
                pass
        
        # Count parameters
        params = sum(p.numel() for p in self.model.parameters())
        logger.info(f"Model loaded: {params:,} parameters")
    
    def bf16_ctx(self):
        if config.use_bf16 and torch.cuda.is_bf16_supported():
            return torch.autocast(device_type="cuda", dtype=torch.bfloat16)
        return nullcontext()
    
    def generate_ensemble(self, input_text: str) -> str:
        # Get TTA variants
        variants = preprocessor.tta_variants(preprocessor.preprocess(input_text))
        
        all_candidates = []
        
        for variant in variants:
            inputs = self.tokenizer(variant, max_length=config.max_length, 
                                   truncation=True, return_tensors="pt")
            inputs = {k: v.to(config.device) for k, v in inputs.items()}
            
            with torch.inference_mode(), self.bf16_ctx():
                # Self-ensemble with different temperatures
                for temp in config.ensemble_temperatures[:config.self_ensemble_size]:
                    outputs = self.model.generate(
                        **inputs,
                        num_beams=config.num_beams,
                        max_new_tokens=config.max_new_tokens,
                        length_penalty=config.length_penalty,
                        repetition_penalty=config.repetition_penalty,
                        do_sample=temp > 0,
                        temperature=temp if temp > 0 else None,
                        top_p=0.9 if temp > 0 else None,
                        num_return_sequences=1,
                        use_cache=True
                    )
                    decoded = self.tokenizer.decode(outputs[0], skip_special_tokens=True)
                    all_candidates.append(decoded)
                
                # MBR candidates
                if config.use_mbr:
                    # Beam candidates
                    beam_out = self.model.generate(
                        **inputs,
                        num_beams=config.num_beams,
                        num_return_sequences=config.mbr_num_beam_cands,
                        max_new_tokens=config.max_new_tokens,
                        length_penalty=config.length_penalty,
                        use_cache=True
                    )
                    beam_texts = [self.tokenizer.decode(o, skip_special_tokens=True) 
                                 for o in beam_out]
                    all_candidates.extend(beam_texts)
                    
                    # Sampled candidates
                    sample_out = self.model.generate(
                        **inputs,
                        do_sample=True,
                        num_beams=1,
                        top_p=0.9,
                        temperature=0.7,
                        num_return_sequences=config.mbr_num_sample_cands,
                        max_new_tokens=config.max_new_tokens,
                        use_cache=True
                    )
                    sample_texts = [self.tokenizer.decode(o, skip_special_tokens=True) 
                                   for o in sample_out]
                    all_candidates.extend(sample_texts)
        
        # Postprocess all candidates
        all_candidates = postprocessor.postprocess(all_candidates)
        
        # MBR selection
        if config.use_mbr:
            return mbr.select(all_candidates)
        return all_candidates[0] if all_candidates else ""
    
    def run(self, test_df: pd.DataFrame) -> pd.DataFrame:
        results = []
        
        for idx, row in tqdm(test_df.iterrows(), total=len(test_df), desc="Translating"):
            sample_id = row["id"]
            text = row["transliteration"]
            
            try:
                translation = self.generate_ensemble(text)
                results.append({"id": sample_id, "translation": translation})
            except Exception as e:
                logger.error(f"Error on sample {sample_id}: {e}")
                results.append({"id": sample_id, "translation": ""})
            
            if (idx + 1) % 50 == 0:
                torch.cuda.empty_cache()
        
        return pd.DataFrame(results)

# ============================================================
# Main
# ============================================================

if __name__ == "__main__":
    logger.info("Loading test data...")
    test_df = pd.read_csv(config.test_data_path)
    logger.info(f"Loaded {len(test_df)} samples")
    
    engine = EnsembleEngine()
    results = engine.run(test_df)
    
    # Save
    output_path = Path(config.output_dir) / "submission.csv"
    results.to_csv(output_path, index=False)
    logger.info(f"Saved submission to {output_path}")
    
    # Validation report
    empty = (results["translation"] == "").sum()
    logger.info(f"Empty translations: {empty}/{len(results)}")
    logger.info(f"Sample translations:")
    print(results.head(10).to_string())