#!/usr/bin/env python3 -u
"""eval_all_externals.py — Evaluate ALL 9 external tokenizers + our best 3 on test set."""

import json, os, sys, time, csv, gc, warnings
from collections import Counter
from dataclasses import dataclass, asdict
from typing import List, Dict

import numpy as np
import regex
warnings.filterwarnings("ignore")

import matplotlib
matplotlib.use("Agg")
import matplotlib.pyplot as plt

BASE = "/root/oiq_cc_tokenizer/results"
CORPORA = os.path.join(BASE, "corpora")
TOK_DIR = os.path.join(BASE, "tokenizers")
PLOTS_DIR = os.path.join(BASE, "plots")

HF_TOKEN = os.environ.get("HF_TOKEN", "")

_WORD_PAT = regex.compile(r"[\p{L}\p{M}\p{N}]+", regex.UNICODE)
_AR_PAT = regex.compile(r"[\u0600-\u06FF\u0750-\u077F]")
_SPECIAL = {"<unk>", "<s>", "</s>", "[CLS]", "[SEP]", "[PAD]", "[UNK]", "<pad>", "<|endoftext|>", "<|im_start|>", "<|im_end|>"}

def segment_words(t): return _WORD_PAT.findall(t)
def count_graphemes(t): return len(regex.findall(r"\X", t))
def detect_script(t): return "ar" if len(_AR_PAT.findall(t)) > len(t) * 0.3 else "az"
def filter_sp(tokens): return [t for t in tokens if t not in _SPECIAL]


@dataclass
class M:
    name: str = ""
    source: str = ""
    algorithm: str = ""
    architecture: str = ""
    vocab_size: int = 0
    fertility_ar: float = 0.0
    fertility_az: float = 0.0
    fertility_overall: float = 0.0
    disparity: float = 0.0
    cpt_ar: float = 0.0
    cpt_az: float = 0.0
    exact_match_ar: float = 0.0
    exact_match_az: float = 0.0


class RawConcat:
    def __init__(self, ar_j, az_j):
        from tokenizers import Tokenizer
        self.ar = Tokenizer.from_file(ar_j)
        self.az = Tokenizer.from_file(az_j)

    def encode(self, text):
        s = detect_script(text)
        t = self.ar if s == "ar" else self.az
        enc = t.encode(text)
        return enc.tokens, enc.ids, s

    def decode(self, ids, script):
        t = self.ar if script == "ar" else self.az
        return t.decode(ids, skip_special_tokens=True)


class HFTok:
    def __init__(self, repo, use_token=False):
        from transformers import AutoTokenizer
        kwargs = {"trust_remote_code": True}
        if use_token:
            kwargs["token"] = HF_TOKEN
        self.tok = AutoTokenizer.from_pretrained(repo, **kwargs)

    def encode(self, text):
        ids = self.tok.encode(text, add_special_tokens=False)
        return self.tok.convert_ids_to_tokens(ids), ids, detect_script(text)

    def decode(self, ids, script):
        return self.tok.decode(ids, skip_special_tokens=True)


def evaluate(tok, name, source, algo, arch, vsz, texts):
    m = M(name=name, source=source, algorithm=algo, architecture=arch, vocab_size=vsz)
    ar_f, az_f, all_f = [], [], []
    ar_c, az_c = [], []
    ar_ok, az_ok, ar_n, az_n = 0, 0, 0, 0

    for i, text in enumerate(texts):
        if (i + 1) % 5000 == 0:
            print(f"    [{i+1}/{len(texts)}] {name}", flush=True)
        try:
            tokens, ids, script = tok.encode(text)
            content = filter_sp(tokens)
            words = segment_words(text)
            if not words:
                continue
            fert = len(content) / len(words)
            all_f.append(fert)
            cpt = count_graphemes(text) / max(len(content), 1)
            try:
                dec = tok.decode(ids, script)
                exact = dec.strip() == text.strip()
            except:
                exact = False
            if script == "ar":
                ar_f.append(fert); ar_c.append(cpt); ar_n += 1
                if exact: ar_ok += 1
            else:
                az_f.append(fert); az_c.append(cpt); az_n += 1
                if exact: az_ok += 1
        except:
            pass

    m.fertility_ar = float(np.mean(ar_f)) if ar_f else 0
    m.fertility_az = float(np.mean(az_f)) if az_f else 0
    m.fertility_overall = float(np.mean(all_f)) if all_f else 0
    mx = max(m.fertility_ar, m.fertility_az, 1e-9)
    m.disparity = abs(m.fertility_ar - m.fertility_az) / mx
    m.cpt_ar = float(np.mean(ar_c)) if ar_c else 0
    m.cpt_az = float(np.mean(az_c)) if az_c else 0
    m.exact_match_ar = ar_ok / max(ar_n, 1)
    m.exact_match_az = az_ok / max(az_n, 1)
    return m


def main():
    texts = []
    for s in ("test_ar", "test_az", "test_mi"):
        p = os.path.join(CORPORA, f"{s}.txt")
        if os.path.exists(p):
            with open(p) as f:
                texts.extend(l.strip() for l in f if l.strip())
    print(f"{len(texts)} test texts", flush=True)

    results = []

    # --- Our best per size (from raw tokenizer JSONs) ---
    ours_cfg = [
        ("concat_bpe_8000", "concat_ar_bpe_4000", "concat_az_bpe_4000", "bpe", "concatenated", 8000),
        ("concat_wordpiece_16000", "concat_ar_wordpiece_8000", "concat_az_wordpiece_8000", "wordpiece", "concatenated", 16000),
        ("concat_bpe_32000", "concat_ar_bpe_16000", "concat_az_bpe_16000", "bpe", "concatenated", 32000),
    ]
    for name, ar_sub, az_sub, algo, arch, vsz in ours_cfg:
        ar_j = os.path.join(TOK_DIR, f"{ar_sub}.json")
        az_j = os.path.join(TOK_DIR, f"{az_sub}.json")
        if os.path.exists(ar_j) and os.path.exists(az_j):
            print(f"\n{name}", flush=True)
            tok = RawConcat(ar_j, az_j)
            r = evaluate(tok, name, "ours", algo, arch, vsz, texts)
            print(f"  F={r.fertility_overall:.3f} D={r.disparity:.3f} EM_ar={r.exact_match_ar:.2%}", flush=True)
            results.append(r)
            del tok; gc.collect()

    # --- ALL External tokenizers ---
    externals = [
        # MSA tokenizers
        ("CaMeLBERT-MSA", "external_msa", "WordPiece", "shared", 30000,
         "CAMeL-Lab/bert-base-arabic-camelbert-msa", False),
        ("Asafaya-BERT", "external_msa", "WordPiece", "shared", 32000,
         "asafaya/bert-base-arabic", False),
        ("Aranizer-SP-86k", "external_msa", "SentencePiece", "shared", 86000,
         "riotu-lab/Aranizer-SP-86k", False),
        ("B2BERT", "external_msa", "WordPiece", "shared", 30000,
         "AHAAM/B2BERT", False),
        # Darija tokenizers
        ("DarijaBERT-ar", "external_darija", "WordPiece", "shared", 80000,
         "SI2M-Lab/DarijaBERT", False),
        ("DarijaBERT-az", "external_darija", "WordPiece", "shared", 110000,
         "SI2M-Lab/DarijaBERT-arabizi", False),
        ("Moroccan-Darija-Tokenizer", "external_darija", "BPE", "shared", 30000,
         "BounharAbdelaziz/Moroccan-Darija-Tokenizer", True),
        ("Translit-Darija", "external_darija", "BPE", "shared", 30000,
         "atlasia/Transliteration-Moroccan-Darija", True),
        ("Qwen2.5-Darija", "external_darija", "SentencePiece", "shared", 151643,
         "GemMaroc/Qwen2.5-7B-Instruct-darija", False),
    ]

    for name, src, algo, arch, vsz, repo, gated in externals:
        print(f"\n{name} ({repo})", flush=True)
        try:
            tok = HFTok(repo, use_token=gated)
            r = evaluate(tok, name, src, algo, arch, vsz, texts)
            print(f"  F={r.fertility_overall:.3f} D={r.disparity:.3f} EM_ar={r.exact_match_ar:.2%}", flush=True)
            results.append(r)
            del tok; gc.collect()
        except Exception as e:
            print(f"  FAILED: {e}", flush=True)

    # Save
    out_csv = os.path.join(BASE, "external_comparison.csv")
    out_json = os.path.join(BASE, "external_comparison.json")
    with open(out_csv, "w", newline="") as f:
        w = csv.DictWriter(f, fieldnames=list(asdict(results[0]).keys()))
        w.writeheader()
        for r in results: w.writerow(asdict(r))
    with open(out_json, "w") as f:
        json.dump([asdict(r) for r in results], f, indent=2)

    # Print table
    print("\n" + "=" * 140, flush=True)
    hdr = f"{'Name':<35} {'Source':<16} {'Algo':<14} {'V':>7} {'Fert':>7} {'F_ar':>7} {'F_az':>7} {'Disp':>7} {'CPT_ar':>7} {'CPT_az':>7} {'EM_ar':>7} {'EM_az':>7}"
    print(hdr, flush=True)
    print("-" * 140, flush=True)
    for r in sorted(results, key=lambda x: (0 if x.source == "ours" else 1, 0 if x.source == "external_darija" else 2, x.vocab_size)):
        print(f"{r.name:<35} {r.source:<16} {r.algorithm:<14} {r.vocab_size:>7,} {r.fertility_overall:>7.3f} {r.fertility_ar:>7.3f} {r.fertility_az:>7.3f} {r.disparity:>7.3f} {r.cpt_ar:>7.3f} {r.cpt_az:>7.3f} {r.exact_match_ar:>7.2%} {r.exact_match_az:>7.2%}", flush=True)
    print("=" * 140, flush=True)

    # Generate plot
    ours_best = [r for r in results if r.source == "ours"]
    ext_msa = [r for r in results if r.source == "external_msa"]
    ext_dar = [r for r in results if r.source == "external_darija"]
    plot_data = ours_best + ext_msa + ext_dar

    colors = {"ours": "#0072B2", "external_msa": "#CC79A7", "external_darija": "#D55E00"}
    color_vals = {"8000": "#E69F00", "16000": "#009E73", "32000": "#0072B2"}
    labels = [f"Ours\n{r.name}" for r in ours_best] + \
             [f"{r.name}" for r in ext_msa] + \
             [f"{r.name}" for r in ext_dar]
    bar_c = [color_vals.get(str(r.vocab_size), "#0072B2") for r in ours_best] + \
            [colors.get(r.source, "#999") for r in ext_msa] + \
            [colors.get(r.source, "#999") for r in ext_dar]
    n = len(plot_data)

    fig, axes = plt.subplots(2, 2, figsize=(20, 11))

    for idx, (vals_fn, title, ylabel) in enumerate([
        (lambda r: r.fertility_overall, "Overall Fertility (Lower = Better)", "Fertility"),
        (lambda r: r.disparity, "Cross-Script Disparity (Lower = Better)", "Disparity"),
    ]):
        ax = axes[0, idx]
        vals = [vals_fn(r) for r in plot_data]
        bars = ax.bar(range(n), vals, color=bar_c, edgecolor="gray", linewidth=0.5)
        ax.set_xticks(range(n))
        ax.set_xticklabels(labels, fontsize=5.5, ha="center", rotation=30)
        ax.set_ylabel(ylabel, fontsize=9)
        ax.set_title(title, fontsize=10, fontweight="bold")
        for b, v in zip(bars, vals):
            ax.text(b.get_x()+b.get_width()/2, b.get_height()+0.005, f"{v:.3f}", ha="center", va="bottom", fontsize=5.5, rotation=45)

    # Exact match grouped
    ax = axes[1, 0]
    x = np.arange(n)
    w = 0.35
    ax.bar(x-w/2, [r.exact_match_ar*100 for r in plot_data], w, label="Arabic", color="#56B4E9", edgecolor="gray", linewidth=0.5)
    ax.bar(x+w/2, [r.exact_match_az*100 for r in plot_data], w, label="Arabizi", color="#333", alpha=0.6, edgecolor="gray", linewidth=0.5)
    ax.set_xticks(x); ax.set_xticklabels(labels, fontsize=5.5, ha="center", rotation=30)
    ax.set_ylabel("Exact Match (%)"); ax.set_title("Exact Reconstruction", fontsize=10, fontweight="bold")
    ax.legend(fontsize=7); ax.set_ylim(0, 108)

    # CPT grouped
    ax = axes[1, 1]
    ax.bar(x-w/2, [r.cpt_ar for r in plot_data], w, label="Arabic", color="#56B4E9", edgecolor="gray", linewidth=0.5)
    ax.bar(x+w/2, [r.cpt_az for r in plot_data], w, label="Arabizi", color="#333", alpha=0.6, edgecolor="gray", linewidth=0.5)
    ax.set_xticks(x); ax.set_xticklabels(labels, fontsize=5.5, ha="center", rotation=30)
    ax.set_ylabel("CPT"); ax.set_title("Characters Per Token (Higher = Better)", fontsize=10, fontweight="bold")
    ax.legend(fontsize=7)

    from matplotlib.patches import Patch
    fig.legend(handles=[
        Patch(fc="#E69F00", label="Ours (8K)"),
        Patch(fc="#009E73", label="Ours (16K)"),
        Patch(fc="#0072B2", label="Ours (32K)"),
        Patch(fc="#CC79A7", label="External (MSA)"),
        Patch(fc="#D55E00", label="External (Darija)"),
    ], loc="upper center", ncol=5, fontsize=9, bbox_to_anchor=(0.5, 0.98), frameon=True)
    plt.tight_layout(rect=[0, 0, 1, 0.95])
    fig.savefig(os.path.join(PLOTS_DIR, "external_comparison.png"), dpi=200, bbox_inches="tight")
    plt.close(fig)
    print(f"\nPlot saved: {os.path.join(PLOTS_DIR, 'external_comparison.png')}", flush=True)
    print("DONE!", flush=True)


if __name__ == "__main__":
    main()