"""
Deep Past - ByT5 with Minimum Bayes Risk (MBR) Decoding

MBR Algorithm:
1. Generate N diverse candidate translations (beam, sampling, nucleus)
2. Score each candidate by average similarity to all others (BLEU/chrF)
3. Select candidate with highest consensus score

Research shows +1-2 BLEU improvement over standard beam search.

Expected BLEU: ~38-39 (vs 37.2 for standard beam search)
"""

import os
import re
import pandas as pd
from pathlib import Path
from typing import List, Tuple
import warnings
warnings.filterwarnings('ignore')

# Install dependencies (NO sacrebleu - not available offline)
os.system("pip install -q transformers sentencepiece accelerate")

import torch
from transformers import AutoModelForSeq2SeqLM, AutoTokenizer

# Simple character n-gram similarity (replacement for sacrebleu - works offline)
def char_ngram_similarity(hyp: str, ref: str, n: int = 6) -> float:
    """Compute character n-gram F1 score (chrF-like metric)."""
    if not hyp or not ref:
        return 0.0

    def get_ngrams(text, n):
        text = text.lower()
        ngrams = set()
        for i in range(len(text) - n + 1):
            ngrams.add(text[i:i+n])
        return ngrams

    hyp_ngrams = get_ngrams(hyp, n)
    ref_ngrams = get_ngrams(ref, n)

    if not hyp_ngrams or not ref_ngrams:
        return 0.0

    overlap = len(hyp_ngrams & ref_ngrams)
    precision = overlap / len(hyp_ngrams) if hyp_ngrams else 0
    recall = overlap / len(ref_ngrams) if ref_ngrams else 0

    if precision + recall == 0:
        return 0.0

    f1 = 2 * precision * recall / (precision + recall)
    return f1

print("=" * 60)
print("Deep Past - ByT5 with MBR Decoding")
print("=" * 60)

# ============================================================
# CONFIGURATION
# ============================================================

DATASET_NAME = "byt5-combined-finetuned"  # pascalledesma/byt5-combined-finetuned
TASK_PREFIX = "translate Akkadian to English: "

# MBR Configuration
MBR_NUM_CANDIDATES = 8  # Generate 8 diverse candidates
MBR_USE_CHRF = True     # Use chrF for consensus (more robust than BLEU)
MBR_FALLBACK_TO_BEAM = True  # Fall back to beam search if MBR fails

# ============================================================
# 1. LOCATE MODEL (KAGGLE DATASET - NO INTERNET)
# ============================================================

print(f"\nLoading model from dataset: {DATASET_NAME}")
print("(Using Kaggle dataset input - offline mode)")

def find_model_path():
    """Find model directory in /kaggle/input datasets"""
    base = Path("/kaggle/input")

    # Primary path - dataset uploaded via kaggle_model_to_dataset.py
    primary = base / DATASET_NAME
    if primary.exists():
        # Find config.json recursively
        for cfg in primary.rglob("config.json"):
            return str(cfg.parent)
        # If no config.json, return base path (model files might be flat)
        return str(primary)

    # Search all inputs for model files
    try:
        for item in base.iterdir():
            if DATASET_NAME.replace("-", "") in item.name.replace("-", ""):
                for cfg in item.rglob("config.json"):
                    return str(cfg.parent)
    except Exception:
        pass

    return None

model_path = find_model_path()

if model_path:
    print(f"Found model at: {model_path}")
    try:
        tokenizer = AutoTokenizer.from_pretrained(model_path)
        model = AutoModelForSeq2SeqLM.from_pretrained(model_path)
        print("Model loaded successfully from Kaggle dataset")
    except Exception as e:
        print(f"Error loading model: {e}")
        raise RuntimeError(f"Failed to load model from {model_path}")
else:
    print(f"CRITICAL ERROR: Dataset not found in /kaggle/input")
    print(f"\nExpected dataset: pascalledesma/{DATASET_NAME}")
    print(f"This should be auto-attached via kernel-metadata.json")
    raise RuntimeError("Model dataset not found.")

device = "cuda" if torch.cuda.is_available() else "cpu"
model = model.to(device)
model.eval()
print(f"Device: {device}")

# MBR uses char_ngram_similarity (offline-compatible)
print(f"\nMBR Configuration (using char n-gram F1 - offline compatible):")
print(f"  Candidates per input: {MBR_NUM_CANDIDATES}")
print(f"  Consensus metric: {'chrF++' if MBR_USE_CHRF else 'BLEU'}")
print(f"  Fallback enabled: {MBR_FALLBACK_TO_BEAM}")

# ============================================================
# 2. LOAD TEST DATA
# ============================================================

print("\nLoading test data...")

test_path = "/kaggle/input/deep-past-initiative-machine-translation/test.csv"
if not Path(test_path).exists():
    # Search for test.csv
    for p in Path("/kaggle/input").rglob("test.csv"):
        test_path = str(p)
        break

test_df = pd.read_csv(test_path)
print(f"Loaded {len(test_df)} test samples")

# Find source column
source_col = "transliteration"
if source_col not in test_df.columns:
    for col in test_df.columns:
        if col.lower() not in ["id", "target", "translation"]:
            source_col = col
            break
print(f"Source column: {source_col}")

# ============================================================
# 3. PREPROCESSING
# ============================================================

def preprocess(text):
    """Normalize Akkadian input."""
    if pd.isna(text) or not str(text).strip():
        return ""

    text = str(text)

    # Normalize gap markers
    text = re.sub(r'\.\.\.+', '<gap>', text)
    text = re.sub(r'\bxx?\b', '<gap>', text)
    text = re.sub(r'\[\.+\]', '<gap>', text)

    return text.strip()

# ============================================================
# 4. POSTPROCESSING
# ============================================================

def postprocess(text):
    """Clean and format output."""
    if not text:
        return ""

    # Remove non-ASCII
    text = re.sub(r'[^\x20-\x7E\s]', '', text)

    # Fix spacing around punctuation
    text = re.sub(r'\s+([.,;:!?])', r'\1', text)
    text = re.sub(r'([.,;:!?])([A-Za-z])', r'\1 \2', text)

    # Remove word repetitions
    text = re.sub(r'\b(\w+)\s+\1\b', r'\1', text, flags=re.IGNORECASE)

    # Fix fractions
    text = re.sub(r'(\d)\s*/\s*(\d)', r'\1/\2', text)

    # Clean whitespace
    text = re.sub(r'\s+', ' ', text).strip()

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

    return text

# ============================================================
# 5. MBR DECODING IMPLEMENTATION
# ============================================================

def generate_diverse_candidates(
    inputs,
    num_candidates: int = 8,
    max_length: int = 256
) -> List[str]:
    """
    Generate diverse candidate translations using multiple decoding strategies.

    Strategies:
    1. Beam search with varying beam sizes
    2. Sampling with temperature
    3. Nucleus sampling (top-p)
    4. Top-k sampling
    """
    candidates = []

    with torch.no_grad():
        # Strategy 1: Beam search with different beam sizes (4 candidates)
        for num_beams in [10, 7, 5, 12]:
            try:
                outputs = model.generate(
                    **inputs,
                    max_length=max_length,
                    num_beams=num_beams,
                    early_stopping=True,
                    no_repeat_ngram_size=3,
                    repetition_penalty=1.5,
                    length_penalty=1.2,
                    do_sample=False,
                )
                decoded = tokenizer.decode(outputs[0], skip_special_tokens=True)
                candidates.append(postprocess(decoded))
            except Exception:
                pass

        # Strategy 2: Sampling with temperature (2 candidates)
        for temp in [0.7, 0.9]:
            try:
                outputs = model.generate(
                    **inputs,
                    max_length=max_length,
                    do_sample=True,
                    temperature=temp,
                    top_k=50,
                    top_p=0.95,
                    no_repeat_ngram_size=3,
                    repetition_penalty=1.5,
                )
                decoded = tokenizer.decode(outputs[0], skip_special_tokens=True)
                candidates.append(postprocess(decoded))
            except Exception:
                pass

        # Strategy 3: Nucleus sampling (top-p) (2 candidates)
        for top_p in [0.92, 0.85]:
            try:
                outputs = model.generate(
                    **inputs,
                    max_length=max_length,
                    do_sample=True,
                    temperature=0.8,
                    top_p=top_p,
                    no_repeat_ngram_size=3,
                    repetition_penalty=1.5,
                )
                decoded = tokenizer.decode(outputs[0], skip_special_tokens=True)
                candidates.append(postprocess(decoded))
            except Exception:
                pass

    # Remove duplicates while preserving order
    seen = set()
    unique_candidates = []
    for cand in candidates:
        if cand and cand not in seen:
            seen.add(cand)
            unique_candidates.append(cand)

    return unique_candidates[:num_candidates]


def compute_consensus_score(candidate: str, all_candidates: List[str], use_chrf: bool = True) -> float:
    """
    Compute consensus score for a candidate.
    Score = average similarity to all other candidates.

    Higher score = more similar to other candidates = consensus translation.
    Uses char_ngram_similarity (offline-compatible replacement for sacrebleu).
    """
    if not candidate or not all_candidates:
        return 0.0

    scores = []

    for other in all_candidates:
        if not other or other == candidate:
            continue

        try:
            # Use char n-gram similarity (works offline, no sacrebleu needed)
            score = char_ngram_similarity(candidate, other, n=6) * 100  # Scale to 0-100
            scores.append(score)
        except Exception:
            # If scoring fails, assign neutral score
            scores.append(0.0)

    return sum(scores) / len(scores) if scores else 0.0


def mbr_decode(
    inputs,
    num_candidates: int = 8,
    max_length: int = 256,
    use_chrf: bool = True
) -> Tuple[str, List[str], List[float]]:
    """
    Minimum Bayes Risk decoding.

    Returns:
        - Best translation (consensus)
        - All candidates
        - Consensus scores for each candidate
    """
    # Step 1: Generate diverse candidates
    candidates = generate_diverse_candidates(inputs, num_candidates, max_length)

    if not candidates:
        return "", [], []

    # If only one candidate, return it
    if len(candidates) == 1:
        return candidates[0], candidates, [1.0]

    # Step 2: Compute consensus scores
    consensus_scores = []
    for candidate in candidates:
        score = compute_consensus_score(candidate, candidates, use_chrf)
        consensus_scores.append(score)

    # Step 3: Select candidate with highest consensus
    best_idx = consensus_scores.index(max(consensus_scores))
    best_translation = candidates[best_idx]

    return best_translation, candidates, consensus_scores

# ============================================================
# 6. TRANSLATE WITH MBR
# ============================================================

def translate_mbr(text, max_length=256):
    """Translate Akkadian to English using MBR decoding."""
    text = preprocess(text)
    if not text:
        return ""

    input_text = TASK_PREFIX + text

    inputs = tokenizer(
        input_text,
        return_tensors="pt",
        max_length=512,
        truncation=True,
        padding=True
    ).to(device)

    try:
        # MBR decoding
        best_translation, candidates, scores = mbr_decode(
            inputs,
            num_candidates=MBR_NUM_CANDIDATES,
            max_length=max_length,
            use_chrf=MBR_USE_CHRF
        )

        if best_translation:
            return best_translation

        # Fallback to standard beam search if MBR fails
        if MBR_FALLBACK_TO_BEAM:
            with torch.no_grad():
                outputs = model.generate(
                    **inputs,
                    max_length=max_length,
                    num_beams=10,
                    early_stopping=True,
                    no_repeat_ngram_size=3,
                    repetition_penalty=1.5,
                    length_penalty=1.2,
                )
            result = tokenizer.decode(outputs[0], skip_special_tokens=True)
            return postprocess(result)

        return ""

    except Exception as e:
        # Emergency fallback to beam search
        print(f"MBR failed: {e}, using beam search fallback")
        try:
            with torch.no_grad():
                outputs = model.generate(
                    **inputs,
                    max_length=max_length,
                    num_beams=10,
                    early_stopping=True,
                    no_repeat_ngram_size=3,
                    repetition_penalty=1.5,
                    length_penalty=1.2,
                )
            result = tokenizer.decode(outputs[0], skip_special_tokens=True)
            return postprocess(result)
        except Exception:
            return ""

# ============================================================
# 7. RUN INFERENCE
# ============================================================

print("\nTranslating test data with MBR decoding...")
print("(This will be slower than standard beam search but should produce better translations)")

translations = []
for i, row in test_df.iterrows():
    source = row[source_col]
    trans = translate_mbr(source)
    translations.append(trans)

    # Progress indicator with examples
    if i < 3:
        print(f"\n[{i}] Source: {str(source)[:70]}...")
        print(f"    MBR Output: {trans[:70]}..." if trans else "    MBR Output: (empty)")
    elif i % 50 == 0:
        print(f"  [{i}/{len(test_df)}] Progress...")

# ============================================================
# 8. CREATE SUBMISSION
# ============================================================

print("\nCreating submission...")

submission = pd.DataFrame({
    "id": test_df["id"],
    "translation": translations  # CRITICAL: must be "translation" not "target"
})

# Validate
empty_count = (submission["translation"] == "").sum()
if empty_count > 0:
    print(f"Warning: {empty_count} empty translations")

total_chars = submission["translation"].str.len().sum()
print(f"Total characters: {total_chars}")

# Quality statistics
avg_length = submission["translation"].str.len().mean()
print(f"Average translation length: {avg_length:.1f} characters")

submission.to_csv("submission.csv", index=False)
print(f"\nSaved: submission.csv")
print(f"  Rows: {len(submission)}")
print(f"  Column names: {list(submission.columns)}")
print(f"\nFirst 5 rows:")
print(submission.head())

print("\n" + "=" * 60)
print("ByT5 MBR Inference Complete!")
print(f"  Dataset: pascalledesma/{DATASET_NAME}")
print(f"  Decoding: Minimum Bayes Risk (MBR)")
print(f"  Candidates per input: {MBR_NUM_CANDIDATES}")
print(f"  Expected BLEU: ~38-39 (+1-2 vs beam search)")
print("=" * 60)