#!/usr/bin/env python3 -u
"""
eval_codeswitch_and_new_baselines.py
1) Evaluate mixed-script texts SEPARATELY (not forced into ar/az binary)
2) Add atlasia/darija_bpe_tokenizer baseline
3) Evaluate on independent DODa dataset (Arabic-only)
"""

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

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

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]")
_LAT_PAT = regex.compile(r"[a-zA-Z]")
_SPECIAL = {"<unk>", "<s>", "</s>", "[CLS]", "[SEP]", "[PAD]", "[UNK]", "<pad>", "",
            "<|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 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

def classify_script_detailed(t):
    """Classify as 'ar', 'az', or 'mi' (mixed)."""
    ar_chars = len(_AR_PAT.findall(t))
    lat_chars = len(_LAT_PAT.findall(t))
    total_alpha = ar_chars + lat_chars
    if total_alpha == 0:
        return "ar"
    ar_ratio = ar_chars / total_alpha
    lat_ratio = lat_chars / total_alpha
    # Pure = >90% one script, mixed = both scripts present with >10% each
    if ar_ratio > 0.9 and lat_ratio < 0.1:
        return "ar"
    elif lat_ratio > 0.9 and ar_ratio < 0.1:
        return "az"
    else:
        return "mi"

def detect_script(t): return "ar" if len(_AR_PAT.findall(t)) > len(t) * 0.3 else "az"


@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_mi: float = 0.0
    fertility_overall: float = 0.0
    disparity: float = 0.0
    cpt_ar: float = 0.0
    cpt_az: float = 0.0
    cpt_mi: float = 0.0
    exact_match_ar: float = 0.0
    exact_match_az: float = 0.0
    exact_match_mi: 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_with_mixed(tok, name, source, algo, arch, vsz, texts):
    """Evaluate with separate mi bucket for code-switched texts."""
    m = M(name=name, source=source, algorithm=algo, architecture=arch, vocab_size=vsz)
    buckets = {"ar": [], "az": [], "mi": []}
    cpt_buckets = {"ar": [], "az": [], "mi": []}
    em_buckets = {"ar": {"ok": 0, "n": 0}, "az": {"ok": 0, "n": 0}, "mi": {"ok": 0, "n": 0}}
    all_f = []

    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)

            # Use DETAILED classification for separate mi bucket
            sc = classify_script_detailed(text)

            buckets[sc].append(fert)
            cpt_buckets[sc].append(cpt)
            em_buckets[sc]["n"] += 1

            try:
                dec = tok.decode(ids, script)
                if normalize_decode(dec) == normalize_decode(text):
                    em_buckets[sc]["ok"] += 1
            except:
                pass
        except:
            pass
    for sc in ("ar", "az", "mi"):
        setattr(m, f"fertility_{sc}", float(np.mean(buckets[sc])) if buckets[sc] else 0)
        setattr(m, f"cpt_{sc}", float(np.mean(cpt_buckets[sc])) if cpt_buckets[sc] else 0)
        b = em_buckets[sc]
        setattr(m, f"exact_match_{sc}", b["ok"] / max(b["n"], 1))

    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
    return m


def evaluate_on_doda(tok, name, source, algo, arch, vsz, texts):
    """Evaluate on independent DODa data (Arabic only, no ar/az split)."""
    all_f, all_c = [], []
    em_ok, em_n = 0, 0

    for i, text in enumerate(texts):
        if (i + 1) % 5000 == 0:
            print(f"    [{i+1}/{len(texts)}] {name} (doda)", 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)
            all_c.append(cpt)
            try:
                dec = tok.decode(ids, script)
                if normalize_decode(dec) == normalize_decode(text):
                    em_ok += 1
            except:
                pass
            em_n += 1
        except:
            pass

    return {
        "name": name, "source": source, "algorithm": algo,
        "architecture": arch, "vocab_size": vsz,
        "n_texts": em_n,
        "fertility": float(np.mean(all_f)) if all_f else 0,
        "cpt": float(np.mean(all_c)) if all_c else 0,
        "exact_match": em_ok / max(em_n, 1),
    }


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

    # Check script distribution with DETAILED classification
    print("=== Script distribution (detailed) ===", flush=True)
    from collections import Counter
    dist = Counter()
    for f in [test_ar, test_az, test_mi]:
        for t in f:
            dist[classify_script_detailed(t)] += 1
    total = sum(dist.values())
    for sc in ("ar", "az", "mi"):
        print(f"  {sc}: {dist[sc]} ({dist[sc]/total*100:.1f}%)", flush=True)
    print(f"  Total: {total}", flush=True)

    all_texts = test_ar + test_az + test_mi
    mi_texts = test_mi  # Only mixed-script texts for dedicated eval

    print(f"\n=== 1. Code-switching evaluation (3 best ours + key externals) ===", flush=True)
    cs_results = []

    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_with_mixed(tok, name, "ours", algo, arch, vsz, all_texts)
            cs_results.append(r)
            print(f"  F_ar={r.fertility_ar:.3f} F_az={r.fertility_az:.3f} F_mi={r.fertility_mi:.3f} EM_mi={r.exact_match_mi:.2%}", flush=True)
            del tok; gc.collect()

    # Key externals for code-switching
    externals_cs = [
        ("DarijaBERT-ar", "external_darija", "WordPiece", "shared", 80000,
         "SI2M-Lab/DarijaBERT", False),
        ("DarijaBERT-mix", "external_darija", "WordPiece", "shared", 160000,
         "SI2M-Lab/DarijaBERT-mix", False),
        ("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_cs:
        print(f"\n{name} ({repo})", flush=True)
        try:
            tok = HFTok(repo, use_token=gated)
            r = evaluate_with_mixed(tok, name, src, algo, arch, vsz, all_texts)
            cs_results.append(r)
            print(f"  F_ar={r.fertility_ar:.3f} F_az={r.fertility_az:.3f} F_mi={r.fertility_mi:.3f} EM_mi={r.exact_match_mi:.2%}", flush=True)
            del tok; gc.collect()
        except Exception as e:
            print(f"  FAILED: {e}", flush=True)

    # Save code-switching results
    cs_csv = os.path.join(BASE, "codeswitch_results.csv")
    cs_json = os.path.join(BASE, "codeswitch_results.json")
    with open(cs_csv, "w", newline="") as f:
        w = csv.DictWriter(f, fieldnames=list(asdict(cs_results[0]).keys()))
        w.writeheader()
        for r in cs_results:
            w.writerow(asdict(r))
    with open(cs_json, "w") as f:
        json.dump([asdict(r) for r in cs_results], f, indent=2)
    print(f"\nCode-switching results saved to {cs_csv}", flush=True)

    # Print code-switching table
    print("\n" + "=" * 130, flush=True)
    hdr = f"{'Name':<30} {'F_ar':>7} {'F_az':>7} {'F_mi':>7} {'CPT_ar':>7} {'CPT_az':>7} {'CPT_mi':>7} {'EM_ar':>7} {'EM_az':>7} {'EM_mi':>7}"
    print(hdr, flush=True)
    print("-" * 130, flush=True)
    for r in cs_results:
        print(f"{r.name:<30} {r.fertility_ar:>7.3f} {r.fertility_az:>7.3f} {r.fertility_mi:>7.3f} {r.cpt_ar:>7.3f} {r.cpt_az:>7.3f} {r.cpt_mi:>7.3f} {r.exact_match_ar:>7.2%} {r.exact_match_az:>7.2%} {r.exact_match_mi:>7.2%}", flush=True)
    print("=" * 130, flush=True)

    # ========== 2. atlasia/darija_bpe_tokenizer ==========
    print("\n=== 2. Evaluating atlasia/darija_bpe_tokenizer ===", flush=True)
    try:
        tok = HFTok("atlasia/darija_bpe_tokenizer", use_token=True)
        r = evaluate_with_mixed(tok, "atlasia_darija_bpe", "external_darija", "BPE", "shared", 0, all_texts)
        # Also need vocab size
        r.vocab_size = tok.tok.vocab_size
        print(f"  Vocab size: {r.vocab_size}", flush=True)
        print(f"  F={r.fertility_overall:.3f} F_ar={r.fertility_ar:.3f} F_az={r.fertility_az:.3f} ΔF={r.disparity:.3f}", flush=True)
        cs_results.append(r)
        del tok; gc.collect()
        print("  atlasia/darija_bpe_tokenizer evaluated successfully", flush=True)
    except Exception as e:
        print(f"  atlasia/darija_bpe_tokenizer FAILED: {e}", flush=True)
        import traceback; traceback.print_exc()

    # ========== 3. Independent DODa evaluation ==========
    print("\n=== 3. Independent dataset evaluation (DODa) ===", flush=True)

    # Try to load DODa from HF
    doda_texts = []
    try:
        from datasets import load_dataset
        print("  Loading DODa from HuggingFace...", flush=True)
        ds = load_dataset("OussamaElbaz/DODa", split="train", token=HF_TOKEN, trust_remote_code=True)
        if ds is not None:
            # Extract Arabic text
            for row in ds:
                t = row.get("text", "") or row.get("arabic", "") or row.get("word", "") or row.get("sentence", "")
                if t and len(t.strip()) > 5:
                    doda_texts.append(t.strip())
            print(f"  Loaded {len(doda_texts)} DODa entries", flush=True)
    except Exception as e:
        print(f"  DODa load failed: {e}", flush=True)
        import traceback; traceback.print_exc()

    # Try alternative DODa repos
    if not doda_texts:
        for repo in ["OussamaElbaz/DODa", "DODa"]:
            try:
                from datasets import load_dataset
                ds = load_dataset(repo, split="train", token=HF_TOKEN, trust_remote_code=True)
                for row in ds:
                    for k, v in row.items():
                        if isinstance(v, str) and len(v.strip()) > 5 and any(c in v for c in "ابتثج"):
                            doda_texts.append(v.strip())
                            break
                if doda_texts:
                    print(f"  Loaded {len(doda_texts)} from {repo}", flush=True)
                    break
            except:
                continue

    if not doda_texts:
        print("  No DODa data available locally. Skipping independent evaluation.", flush=True)
        print("  (Would need to download DODa separately)", flush=True)
    else:
        print(f"  Evaluating {len(doda_texts)} DODa texts...", flush=True)
        doda_results = []
        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):
                tok = RawConcat(ar_j, az_j)
                r = evaluate_on_doda(tok, name, "ours", algo, arch, vsz, doda_texts)
                doda_results.append(r)
                print(f"  {name}: F={r['fertility']:.3f} CPT={r['cpt']:.3f} EM={r['exact_match']:.2%}", flush=True)
                del tok; gc.collect()

        # Key externals
        for name, repo in [("CaMeLBERT-MSA", "CAMeL-Lab/bert-base-arabic-camelbert-msa"),
                           ("Qwen2.5-Darija", "GemMaroc/Qwen2.5-7B-Instruct-darija")]:
            try:
                tok = HFTok(repo, use_token=False)
                r = evaluate_on_doda(tok, name, "external", "WordPiece", "shared", 0, doda_texts)
                doda_results.append(r)
                print(f"  {name}: F={r['fertility']:.3f} CPT={r['cpt']:.3f} EM={r['exact_match']:.2%}", flush=True)
                del tok; gc.collect()
            except Exception as e:
                print(f"  {name} FAILED: {e}", flush=True)

        doda_csv = os.path.join(BASE, "doda_independent_results.csv")
        with open(doda_csv, "w", newline="") as f:
            w = csv.DictWriter(f, fieldnames=list(doda_results[0].keys()))
            w.writeheader()
            for r in doda_results:
                w.writerow(r)
        print(f"\nDODa results saved to {doda_csv}", flush=True)

    print("\n=== ALL DONE ===", flush=True)


if __name__ == "__main__":
    main()