# ============================================================
# Akkadian MT: Ensemble + Data Augmentation Pipeline
# ============================================================
# OFFLINE MODE - No external dependencies beyond Kaggle defaults
# ============================================================

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, 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

warnings.filterwarnings("ignore")

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)

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

# ============================================================
# Custom chrF-like similarity (no sacrebleu dependency)
# ============================================================

def chrf_similarity(a: str, b: str, ngram_order: int = 6) -> float:
    """Compute character n-gram F-score similarity between two strings."""
    if not a or not b:
        return 0.0
    
    def get_ngrams(s, n):
        s = s.lower()
        return [s[i:i+n] for i in range(len(s) - n + 1)]
    
    scores = []
    for n in range(1, ngram_order + 1):
        a_ngrams = Counter(get_ngrams(a, n))
        b_ngrams = Counter(get_ngrams(b, n))
        
        if not a_ngrams or not b_ngrams:
            continue
        
        overlap = sum((a_ngrams & b_ngrams).values())
        precision = overlap / sum(a_ngrams.values()) if sum(a_ngrams.values()) > 0 else 0
        recall = overlap / sum(b_ngrams.values()) if sum(b_ngrams.values()) > 0 else 0
        
        if precision + recall > 0:
            f_score = 2 * precision * recall / (precision + recall)
            scores.append(f_score)
    
    return sum(scores) / len(scores) * 100 if scores else 0.0

# ============================================================
# 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/"
    
    # Self-Ensemble
    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 Helper
# ============================================================

class LexiconHelper:
    def __init__(self, lexicon_path: str, dict_path: str):
        self.akka_to_eng: Dict[str, List[str]] = {}
        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 with TTA
# ============================================================

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()
            text = re.sub(r'\b(\w+)(?:\s+\1\b)+', r'\1', text)
            results.append(text)
        return results

postprocessor = Postprocessor()

# ============================================================
# MBR Selector (using custom similarity)
# ============================================================

class MBRSelector:
    def similarity(self, a: str, b: str) -> float:
        return chrf_similarity(a, b)
    
    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 ""
        
        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):
        logger.info("Loading model...")
        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
        
        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:
        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_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)
                    
                    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)
        
        all_candidates = postprocessor.postprocess(all_candidates)
        
        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)
    
    output_path = Path(config.output_dir) / "submission.csv"
    results.to_csv(output_path, index=False)
    logger.info(f"Saved submission to {output_path}")
    
    empty = (results["translation"] == "").sum()
    logger.info(f"Empty translations: {empty}/{len(results)}")
    print(results.head(10).to_string())