#!/usr/bin/env python3 -u
"""eval_and_compare.py — Evaluate all ours (28) + externals, generate plot, save results."""

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

import numpy as np
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")

import regex
_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>"}

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 RawShared:
    def __init__(self, j):
        from tokenizers import Tokenizer
        self.tok = Tokenizer.from_file(j)

    def encode(self, text):
        enc = self.tok.encode(text)
        return enc.tokens, enc.ids, detect_script(text)

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


class HFTok:
    def __init__(self, repo):
        from transformers import AutoTokenizer
        self.tok = AutoTokenizer.from_pretrained(repo, trust_remote_code=True)

    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():
    # Load test texts
    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)} texts", flush=True)

    results = []

    # --- Our tokenizers ---
    for vsz in (8000, 16000, 32000):
        for algo in ("bpe", "unigram", "wordpiece", "bbpe"):
            # Shared
            jp = os.path.join(TOK_DIR, f"shared_{algo}_{vsz}.json")
            if os.path.exists(jp):
                name = f"shared_{algo}_{vsz}"
                print(f"\n{name}", flush=True)
                tok = RawShared(jp)
                r = evaluate(tok, name, "ours", algo, "shared", 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()

            # Concat
            ar_j = os.path.join(TOK_DIR, f"concat_ar_{algo}_{vsz//2}.json")
            az_j = os.path.join(TOK_DIR, f"concat_az_{algo}_{vsz//2}.json")
            if os.path.exists(ar_j) and os.path.exists(az_j):
                name = f"concat_{algo}_{vsz}"
                print(f"\n{name}", flush=True)
                tok = RawConcat(ar_j, az_j)
                r = evaluate(tok, name, "ours", algo, "concatenated", 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()

    # --- External ---
    externals = [
        ("CaMeLBERT-MSA", "external_msa", "WordPiece", "shared", 30000, "CAMeL-Lab/bert-base-arabic-camelbert-msa"),
        ("Asafaya-BERT", "external_msa", "WordPiece", "shared", 32000, "asafaya/bert-base-arabic"),
        ("Aranizer-SP-86k", "external_msa", "SentencePiece", "shared", 86000, "riotu-lab/Aranizer-SP-86k"),
        ("DarijaBERT-ar", "external_darija", "WordPiece", "shared", 80000, "SI2M-Lab/DarijaBERT"),
        ("DarijaBERT-az", "external_darija", "WordPiece", "shared", 110000, "SI2M-Lab/DarijaBERT-arabizi"),
    ]
    for name, src, algo, arch, vsz, repo in externals:
        print(f"\n{name} ({repo})", flush=True)
        try:
            tok = HFTok(repo)
            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" + "=" * 130, flush=True)
    hdr = f"{'Name':<30} {'Source':<16} {'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("-" * 130, flush=True)
    for r in sorted(results, key=lambda x: (0 if x.source == "ours" else 1, x.vocab_size)):
        print(f"{r.name:<30} {r.source:<16} {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("=" * 130, flush=True)

    # Generate comparison plot
    ours_best = []
    for vsz in (8000, 16000, 32000):
        cands = [r for r in results if r.vocab_size == vsz and r.architecture == "concatenated" and r.source == "ours"]
        if cands:
            best = min(cands, key=lambda x: x.fertility_overall)
            ours_best.append(best)
    ext = [r for r in results if r.source != "ours"]
    plot_data = ours_best + ext

    fig, axes = plt.subplots(2, 2, figsize=(16, 11))
    colors = {"8000": "#E69F00", "16000": "#009E73", "32000": "#0072B2",
             "external_msa": "#CC79A7", "external_darija": "#D55E00"}
    labels = [f"Ours\n{r.name}\n({r.vocab_size:,})" for r in ours_best] + \
            [f"{r.name}\n({r.vocab_size:,})" for r in ext]
    bar_c = [colors[str(r.vocab_size)] for r in ours_best] + \
            [colors.get(r.source, "#999") for r in ext]
    n = len(plot_data)

    for idx, (key, vals_fn, title, ylabel) in enumerate([
        ("fert", lambda r: r.fertility_overall, "Overall Fertility (Lower = Better)", "Fertility"),
        ("disp", 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=6, ha="center")
        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=6)

    # 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")
    ax.bar(x+w/2, [r.exact_match_az*100 for r in plot_data], w, label="Arabizi", color="#333")
    ax.set_xticks(x); ax.set_xticklabels(labels, fontsize=6, ha="center")
    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")
    ax.bar(x+w/2, [r.cpt_az for r in plot_data], w, label="Arabizi", color="#333")
    ax.set_xticks(x); ax.set_xticklabels(labels, fontsize=6, ha="center")
    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=colors["8000"], label="Ours (8K)"),
        Patch(fc=colors["16000"], label="Ours (16K)"),
        Patch(fc=colors["32000"], label="Ours (32K)"),
        Patch(fc=colors["external_msa"], label="External (MSA)"),
        Patch(fc=colors["external_darija"], label="External (Darija)"),
    ], loc="upper center", ncol=5, fontsize=8, 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=150, bbox_inches="tight")
    plt.close(fig)
    print(f"\nPlot: {os.path.join(PLOTS_DIR, 'external_comparison.png')}", flush=True)
    print("DONE!", flush=True)

if __name__ == "__main__":
    main()