#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
run_baseline_eval.py — المرجع الخارجي الصادق (مراجعة 2026-10-01)
====================================================================
LaBSE + multilingual-e5-small على نفس بروتوكول twin-retrieval بالضبط:
    n=300 عينة (seed=41) لكل مجموعة، المرشحون = 300 (التوأم + 299 مشتّتًا)،
    نفس التنظيف (download_corpus)، نفس الاتجاهين ar→en / en→ar.
النجاح = نص التوأم المُسترجَع يطابق التوأم الحقيقي (نفس دلالة cross_lingual_stats).
"""
import os, sys, json, random
import numpy as np

HERE = os.path.dirname(os.path.abspath(__file__))
sys.path.insert(0, HERE)
from phase_w1_data import download_corpus
from phase_c_encoder import TrainableEncoder
from phase_p2_slot import hashed_features_slot
from run_tests_phase_w1 import cross_lingual_stats   # نفس الكود = تطابق تام

HELDOUT = [
    ("Tanzil", "https://object.pouta.csc.fi/OPUS-Tanzil/v1/moses/ar-en.txt.zip", 20000),
    ("Bible", "https://object.pouta.csc.fi/OPUS-bible-uedin/v1/moses/ar-en.txt.zip", 20000),
    ("NeuLab-TED", "https://object.pouta.csc.fi/OPUS-NeuLab-TedTalks/v1/moses/ar-en.txt.zip", 20000),
]
N, SEED = 300, 41


def sample_pairs(pairs, n=N, seed=SEED):
    rng = random.Random(seed)
    return rng.sample(pairs, min(n, len(pairs)))


def retrieve_acc(emb_q, texts_q, emb_t, texts_t):
    """Twin retrieval — نص التوأم المُرجَّع يطابق الحقيقي (كما في cross_lingual_stats)."""
    q = emb_q / np.maximum(np.linalg.norm(emb_q, axis=1, keepdims=True), 1e-9)
    t = emb_t / np.maximum(np.linalg.norm(emb_t, axis=1, keepdims=True), 1e-9)
    pred = (q @ t.T).argmax(axis=1)
    ok = sum(1 for i, j in enumerate(pred) if texts_t[j] == texts_t[i])
    return round(ok / len(q), 4), ok


def wilson(k, n, z=1.96):
    p = k / n
    d = 1 + z * z / n
    c = (p + z * z / (2 * n)) / d
    h = z * ((p * (1 - p) / n + z * z / (4 * n * n)) ** 0.5) / d
    return [round(c - h, 3), round(c + h, 3)]


def main():
    corpora = {}
    for name, url, lim in HELDOUT:
        ps = download_corpus(name, url, lim, os.path.join(HERE, "w1_data"))
        corpora[name] = ps
        print(f"corpus {name}: {len(ps)} pairs", flush=True)
    corpora["ALL"] = sum((corpora[n] for n, _, _ in HELDOUT), [])

    # ---- (1) LHC-0 W4 — نفس مسار الكود الذي نُشرت أرقامه ----
    z = np.load(os.path.join(HERE, "phase_w1_W.npz"))
    enc = TrainableEncoder(F=int(z["F"]), d=int(z["d"]), tau=float(z["tau"]),
                           W=z["W"], features_fn=hashed_features_slot)
    results = {}
    for name, ps in corpora.items():
        r = cross_lingual_stats(enc, ps, n=N, seed=SEED)
        results[f"ours_W4[{name}]"] = r
        print(json.dumps({"name": f"ours_W4[{name}]", **r}, ensure_ascii=False), flush=True)

    # ---- (2) مراجع خارجية ----
    import torch
    from transformers import AutoTokenizer, AutoModel

    def encode_batch(model, tok, texts, prefix=""):
        embs = []
        with torch.no_grad():
            for i in range(0, len(texts), 32):
                batch = [prefix + t for t in texts[i:i + 32]]
                toks = tok(batch, padding=True, truncation=True,
                           max_length=128, return_tensors="pt")
                h = model(**toks).last_hidden_state
                m = toks["attention_mask"].unsqueeze(-1).float()
                embs.append(((h * m).sum(1) / m.sum(1)).cpu().numpy())
        return np.vstack(embs)

    for mname, qp, pp in [("sentence-transformers/LaBSE", "", ""),
                          ("intfloat/multilingual-e5-small", "query: ", "passage: ")]:
        print(f"\n=== {mname} ===", flush=True)
        tok = AutoTokenizer.from_pretrained(mname)
        model = AutoModel.from_pretrained(mname)
        model.eval()
        for name, ps in corpora.items():
            s = sample_pairs(ps)
            ar = [p["ar"] for p in s]
            en = [p["en"] for p in s]
            e_ar_q = encode_batch(model, tok, ar, qp)   # ar كاستعلام
            e_en_p = encode_batch(model, tok, en, pp)   # en كمرور
            e_en_q = encode_batch(model, tok, en, qp)   # en كاستعلام
            e_ar_p = encode_batch(model, tok, ar, pp)   # ar كمرور
            a1, k1 = retrieve_acc(e_ar_q, ar, e_en_p, en)     # ar→en
            a2, k2 = retrieve_acc(e_en_q, en, e_ar_p, ar)     # en→ar
            results[f"{mname}[{name}]"] = {
                "ar_to_en": a1, "en_to_ar": a2, "n": N,
                "ci95_ar_en": wilson(k1, N), "ci95_en_ar": wilson(k2, N)}
            print(json.dumps({"name": f"{mname}[{name}]", **results[f"{mname}[{name}]"]},
                             ensure_ascii=False), flush=True)
        del model

    with open(os.path.join(HERE, "baseline_results.json"), "w", encoding="utf-8") as f:
        json.dump(results, f, ensure_ascii=False, indent=1)
    print("\nSAVED baseline_results.json ✓", flush=True)


if __name__ == "__main__":
    main()