#!/usr/bin/env python3 -u
"""eval_darijabert_mix.py — Evaluate DarijaBERT-mix on the same test set and append to external_comparison.csv."""

import json, os, csv, gc, warnings
from dataclasses import dataclass, asdict

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

BASE = "/root/oiq_cc_tokenizer/results"
CORPORA = os.path.join(BASE, "corpora")
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]
def normalize_decode(s):
    s = s.replace("##", "")
    s = " ".join(s.split())
    return s


@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


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:
            ids = tok.encode(text, add_special_tokens=False)
            tokens = tok.convert_ids_to_tokens(ids)
            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, skip_special_tokens=True)
                exact = normalize_decode(dec) == normalize_decode(text)
            except:
                exact = False
            script = detect_script(text)
            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():
    from transformers import AutoTokenizer

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

    # Load DarijaBERT-mix
    repo = "SI2M-Lab/DarijaBERT-mix"
    print(f"\nLoading {repo} ...", flush=True)
    tok = AutoTokenizer.from_pretrained(repo, trust_remote_code=True)
    vsz = tok.vocab_size
    print(f"  vocab_size = {vsz}", flush=True)

    # Evaluate
    print(f"\nEvaluating DarijaBERT-mix ...", flush=True)
    r = evaluate(tok, "DarijaBERT-mix", "external_darija", "WordPiece", "shared", vsz, texts)
    print(f"  F={r.fertility_overall:.3f}  F_ar={r.fertility_ar:.3f}  F_az={r.fertility_az:.3f}", flush=True)
    print(f"  D={r.disparity:.3f}  CPT_ar={r.cpt_ar:.3f}  CPT_az={r.cpt_az:.3f}", flush=True)
    print(f"  EM_ar={r.exact_match_ar:.2%}  EM_az={r.exact_match_az:.2%}", flush=True)

    # Append to external_comparison.csv
    csv_path = os.path.join(BASE, "external_comparison.csv")
    with open(csv_path, "a", newline="") as f:
        w = csv.DictWriter(f, fieldnames=list(asdict(r).keys()))
        w.writerow(asdict(r))
    print(f"\nAppended to {csv_path}", flush=True)

    print("DONE!", flush=True)


if __name__ == "__main__":
    main()