#!/usr/bin/env python
"""
Two things:
  1. Tokenizer fertility analysis across all 4 languages + English baseline
  2. Add new lang tokens for Tulu (tcy_Knda) and Gondi (gon_Deva)
     initialized from their proxy token embeddings (kan_Knda, hin_Deva)

Run this BEFORE nllb_pretrain.py.
Saves extended tokenizer to ./nllb-extended-tokenizer/
"""

import os
import random
import torch
import numpy as np
import pandas as pd
from datasets import load_dataset
from transformers import AutoTokenizer
from transformers.models.m2m_100 import M2M100ForConditionalGeneration

MODEL_ID     = "facebook/nllb-200-1.3B"
MONO_DIR     = "../../Dataset/Monolingual"
OUTPUT_DIR   = "../nllb-extended-tokenizer"
SAMPLE_SIZE  = 1_000   # sentences per language for fertility analysis
RANDOM_SEED  = 42

random.seed(RANDOM_SEED)
os.makedirs(OUTPUT_DIR, exist_ok=True)

# ── Language config ────────────────────────────────────────────────────────────
# existing_token: the NLLB token used during CPT (proxy)
# new_token:      new token to add (None if language already in NLLB)
LANGUAGES = [
    {
        "name":           "Mizo",
        "file":           "mizo_mono_cpt_train.tsv",
        "existing_token": "lus_Latn",
        "new_token":      None,           # Mizo is in NLLB as lus_Latn
    },
    {
        "name":           "Khasi",
        "file":           "khasi_mono_cpt_train.tsv",
        "existing_token": "kha_Latn",
        "new_token":      None,           # Khasi is in NLLB as kha_Latn
    },
    {
        "name":           "Kannada (proxy for Tulu)",
        "file":           "kannada_mono_cpt_train.tsv",
        "existing_token": "kan_Knda",
        "new_token":      None,
    },
    {
        "name":           "Tulu",
        "file":           None,           # no monolingual data yet
        "existing_token": "kan_Knda",     # proxy used during CPT
        "new_token":      "tcy_Knda",     # new token, init from kan_Knda
    },
    {
        "name":           "Hindi (proxy for Gondi)",
        "file":           "hindi_mono_cpt_train.tsv",
        "existing_token": "hin_Deva",
        "new_token":      None,
    },
    {
        "name":           "Gondi",
        "file":           None,           # no monolingual data yet
        "existing_token": "hin_Deva",     # proxy used during CPT
        "new_token":      "gon_Deva",     # new token, init from hin_Deva
    },
]

# English baseline for comparison
ENGLISH_SAMPLE = [
    "The government announced new policies for education reform.",
    "She walked to the market every morning to buy fresh vegetables.",
    "The river flooded several villages in the northern district.",
    "Children played in the courtyard while their parents worked.",
    "The doctor advised him to rest for at least a week.",
]

# ── Load tokenizer ─────────────────────────────────────────────────────────────
print(f"Loading tokenizer: {MODEL_ID}")
tokenizer = AutoTokenizer.from_pretrained(MODEL_ID)
print(f"Vocab size before extension: {tokenizer.vocab_size:,}")

# ── Fertility analysis ─────────────────────────────────────────────────────────
def compute_fertility(texts, lang_token, tokenizer, sample_size=SAMPLE_SIZE):
    """
    Fertility = avg tokens per word.
    Lower is better — 1.0 = perfect (one token per word).
    Also reports % unknown tokens (shown as <unk>).
    """
    if not texts:
        return {"fertility": None, "unk_rate": None, "n": 0}

    sample = random.sample(texts, min(sample_size, len(texts)))
    tokenizer.src_lang = lang_token

    fertilities = []
    unk_counts  = []
    total_toks  = []

    for text in sample:
        words      = text.split()
        if not words:
            continue
        token_ids  = tokenizer.encode(text, add_special_tokens=False)
        n_toks     = len(token_ids)
        n_words    = len(words)
        n_unk      = token_ids.count(tokenizer.unk_token_id)

        fertilities.append(n_toks / n_words)
        unk_counts.append(n_unk)
        total_toks.append(n_toks)

    return {
        "fertility":  round(np.mean(fertilities), 3),
        "fertility_std": round(np.std(fertilities), 3),
        "unk_rate":   round(sum(unk_counts) / max(sum(total_toks), 1) * 100, 2),
        "n":          len(fertilities),
    }


def load_sample(filepath):
    if filepath is None or not os.path.exists(filepath):
        return []
    df = pd.read_csv(filepath, sep='\t', header=None, names=['text'],
                     on_bad_lines='skip', quoting=3)
    return df['text'].dropna().str.strip().tolist()


print("\n" + "="*65)
print("  TOKENIZER FERTILITY ANALYSIS")
print("="*65)
print(f"  Model   : {MODEL_ID}")
print(f"  Sample  : {SAMPLE_SIZE} sentences per language")
print("="*65)

fertility_rows = []

# English baseline
tokenizer.src_lang = "eng_Latn"
eng_fertilities = []
for text in ENGLISH_SAMPLE * 20:   # repeat to get ~100 samples
    ids = tokenizer.encode(text, add_special_tokens=False)
    eng_fertilities.append(len(ids) / len(text.split()))

eng_fertility = round(np.mean(eng_fertilities), 3)
print(f"\n  {'Language':<35} {'Token':>10} {'Fertility':>10} {'Std':>6} {'UNK%':>7} {'N':>6}")
print(f"  {'-'*35} {'-'*10} {'-'*10} {'-'*6} {'-'*7} {'-'*6}")
print(f"  {'English (baseline)':<35} {'eng_Latn':>10} {eng_fertility:>10.3f} {'—':>6} {'—':>7} {'~100':>6}")

fertility_rows.append({
    "Language": "English (baseline)",
    "NLLB Token": "eng_Latn",
    "Fertility": eng_fertility,
    "Fertility Std": "—",
    "UNK%": "—",
})

for lang in LANGUAGES:
    if lang["file"] is None:
        # Tulu / Gondi — no data yet, skip fertility
        print(f"  {lang['name']:<35} {lang['existing_token']:>10} {'N/A (no data)':>10}")
        fertility_rows.append({
            "Language": lang["name"],
            "NLLB Token": lang["existing_token"],
            "Fertility": "N/A",
            "Fertility Std": "N/A",
            "UNK%": "N/A",
        })
        continue

    filepath = os.path.join(MONO_DIR, lang["file"])
    texts    = load_sample(filepath)

    if not texts:
        print(f"  {lang['name']:<35} FILE NOT FOUND: {filepath}")
        continue

    stats = compute_fertility(texts, lang["existing_token"], tokenizer)
    flag  = "  ⚠ HIGH" if stats["fertility"] and stats["fertility"] > 3.5 else ""

    print(f"  {lang['name']:<35} {lang['existing_token']:>10} "
          f"{stats['fertility']:>10.3f} {stats['fertility_std']:>6.3f} "
          f"{stats['unk_rate']:>7.2f}% {stats['n']:>6}{flag}")

    fertility_rows.append({
        "Language":       lang["name"],
        "NLLB Token":     lang["existing_token"],
        "Fertility":      stats["fertility"],
        "Fertility Std":  stats["fertility_std"],
        "UNK%":           stats["unk_rate"],
    })

print("\n  Fertility interpretation:")
print("  < 1.5 → Excellent  |  1.5-2.5 → Good  |  2.5-3.5 → Poor  |  > 3.5 → Very poor")
print("  UNK% > 5% → tokenizer has no subwords for this language\n")

# Save fertility table
fertility_df = pd.DataFrame(fertility_rows)
fertility_path = os.path.join(OUTPUT_DIR, "fertility_report.tsv")
fertility_df.to_csv(fertility_path, sep='\t', index=False)
print(f"  Fertility report saved → {fertility_path}")

# ── Add new lang tokens for Tulu and Gondi ─────────────────────────────────────
print("\n" + "="*65)
print("  ADDING NEW LANG TOKENS: tcy_Knda (Tulu), gon_Deva (Gondi)")
print("="*65)

NEW_TOKENS = [
    {"new": "tcy_Knda", "proxy": "kan_Knda"},  # Tulu from Kannada
    {"new": "gon_Deva", "proxy": "hin_Deva"},  # Gondi from Hindi
]

# Check which tokens already exist
tokens_to_add = []
for t in NEW_TOKENS:
    if t["new"] in tokenizer.all_special_tokens:
        print(f"  {t['new']} already exists — skipping")
    else:
        tokens_to_add.append(t["new"])
        print(f"  Will add: {t['new']}  (init from {t['proxy']})")

if tokens_to_add:
    # Add to tokenizer
    tokenizer.add_special_tokens({"additional_special_tokens": tokens_to_add})
    print(f"\n  Vocab size after extension: {len(tokenizer):,}  "
          f"(added {len(tokens_to_add)} tokens)")

    # Load model to extend embeddings
    print("\n  Loading model to extend embedding layer...")
    model = M2M100ForConditionalGeneration.from_pretrained(MODEL_ID, dtype=torch.float32)
    model.resize_token_embeddings(len(tokenizer))

    # Initialize new token embeddings from proxy token embeddings
    with torch.no_grad():
        embed = model.model.shared.weight   # shared encoder+decoder embeddings

        for t in NEW_TOKENS:
            if t["new"] not in tokens_to_add:
                continue
            new_id   = tokenizer.convert_tokens_to_ids(t["new"])
            proxy_id = tokenizer.convert_tokens_to_ids(t["proxy"])

            # Copy proxy embedding + small random noise for differentiation
            proxy_emb = embed[proxy_id].clone()
            noise     = torch.randn_like(proxy_emb) * 0.01
            embed[new_id] = proxy_emb + noise

            print(f"  Initialized {t['new']} (id={new_id}) "
                  f"from {t['proxy']} (id={proxy_id})")

    # Save extended tokenizer (model weights don't need saving — 
    # they'll be loaded fresh in pretrain script and resized there)
    tokenizer.save_pretrained(OUTPUT_DIR)
    print(f"\n  Extended tokenizer saved → {OUTPUT_DIR}/")
    print(f"  NOTE: Update nllb_pretrain.py to load from '{OUTPUT_DIR}'")
    print(f"        and call model.resize_token_embeddings(len(tokenizer))")

else:
    tokenizer.save_pretrained(OUTPUT_DIR)
    print("  No new tokens needed. Tokenizer saved as-is.")

# ── Verify new tokens ──────────────────────────────────────────────────────────
print("\n" + "="*65)
print("  VERIFICATION")
print("="*65)
for t in NEW_TOKENS:
    tid = tokenizer.convert_tokens_to_ids(t["new"])
    pid = tokenizer.convert_tokens_to_ids(t["proxy"])
    print(f"  {t['new']:15s} → id {tid:6d}  |  "
          f"proxy {t['proxy']:15s} → id {pid:6d}")

print("\nDone. Run nllb_pretrain.py next.")