#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""baseline_labse_blocks.py — LaBSE على نفس كتل e5 (بذرة 41، 4 كتل×300 لكل مجموعة)
نفس عناصر e5 بالضبط (تحقق صارم من تطابق القناع) + McNemar مقترن على النظيف فقط.
LaBSE متناظر ⇒ ترميز واحد لكل نص (بلا بادئات)."""
import os, sys, json, time
import numpy as np
sys.path.insert(0,"/home/user/lhc/code"); sys.path.insert(0,"/home/user/lhc/recon")
from phase_w1_data import download_corpus
from scipy.stats import binomtest
import torch
from transformers import AutoTokenizer, AutoModel
torch.set_num_threads(2)
T0=time.time(); log=lambda *a: print(f"[{time.time()-T0:6.1f}s]", *a, flush=True)
R="/home/user/lhc/results"; BS=300; KBLOCKS=4
M=json.load(open(f"{R}/masks_v2.json")); masks={cn:set(M["masks_strong"][cn]) for cn in M["masks_strong"]}
DD="/home/user/lhc/data/opus"
OPUS=[("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)]
pairs={n:download_corpus(n,u,l,DD) for n,u,l in OPUS}
DEV="/home/user/lhc/data/flores101_dataset/devtest"
ld=lambda p:[l.strip() for l in open(p,encoding="utf-8") if l.strip()]
pairs["FLORES"]=[{"ar":a,"en":e} for a,e in zip(ld(f"{DEV}/ara.devtest"),ld(f"{DEV}/eng.devtest"))]
def blocks(n,seed=41):
    rng=np.random.RandomState(seed); idx=rng.permutation(n)
    return [idx[i:i+BS] for i in range(0,n-BS+1,BS)]
tok=AutoTokenizer.from_pretrained("/var/tmp/labse"); model=AutoModel.from_pretrained("/var/tmp/labse").eval()
def enc(texts, bs=16):
    out=[]
    with torch.no_grad():
        for i in range(0,len(texts),bs):
            t=tok(texts[i:i+bs],padding=True,truncation=True,max_length=128,return_tensors="pt")
            h=model(**t).last_hidden_state; m=t["attention_mask"].unsqueeze(-1).float()
            out.append(((h*m).sum(1)/m.sum(1)).numpy())
    return np.vstack(out)
norm=lambda X: X/np.maximum(np.linalg.norm(X,axis=1,keepdims=True),1e-9)
zh=np.load(f"{R}/block_hits.npz")
out={"baseline":"LaBSE (sentence-transformers/LaBSE)","blocks_per_set":KBLOCKS,"block_size":BS,"seed":41,
     "protocol":"نفس عناصر e5 في baseline_clean_blocks.json (تحقق تطابق القناع)","subsets":{}}
store={}
for cn in ("Tanzil","Bible","NeuLab-TED","FLORES"):
    P=pairs[cn]; bl=blocks(len(P))[:KBLOCKS]; items=[int(i) for b in bl for i in b]
    ar=[P[i]["ar"] for i in items]; en=[P[i]["en"] for i in items]
    inmask=np.array([i in masks[cn] for i in items],bool)
    log(f"{cn}: {len(items)} عنصرًا — ترميز LaBSE ...")
    A=norm(enc(ar)); E=norm(enc(en))
    s1=A@E.T; pred_aen=s1.argmax(1)   # ar→en: لكل العربي، أقرب إنجليزي
    s2=E@A.T; pred_ear=s2.argmax(1)   # en→ar
    hit_ae=np.array([en[pred_aen[i]]==en[i] for i in range(len(items))],bool)
    hit_ea=np.array([ar[pred_ear[j]]==ar[j] for j in range(len(items))],bool)
    n_used=KBLOCKS*BS
    lhc_ae=zh[f"{cn}|W4|ar_en"][:n_used]; lhc_ea=zh[f"{cn}|W4|en_ar"][:n_used]; lhc_m=zh[f"{cn}|W4|mask"][:n_used]
    assert np.array_equal(lhc_m,inmask), "عدم تطابق قناع مع e5/LHC!"
    c=~inmask
    res={"n":len(items),"n_clean":int(c.sum()),"n_masked":int(inmask.sum()),
      "labse":{"r1_ar_en":round(float(hit_ae.mean()),4),"r1_en_ar":round(float(hit_ea.mean()),4),
        "clean_ar_en":round(float(hit_ae[c].mean()),4),"clean_en_ar":round(float(hit_ea[c].mean()),4),
        "masked_ar_en":round(float(hit_ae[inmask].mean()),4) if inmask.any() else None,
        "masked_en_ar":round(float(hit_ea[inmask].mean()),4) if inmask.any() else None},
      "lhc_w4":{"r1_ar_en":round(float(lhc_ae.mean()),4),"r1_en_ar":round(float(lhc_ea.mean()),4),
        "clean_ar_en":round(float(lhc_ae[c].mean()),4),"clean_en_ar":round(float(lhc_ea[c].mean()),4),
        "masked_ar_en":round(float(lhc_ae[inmask].mean()),4) if inmask.any() else None,
        "masked_en_ar":round(float(lhc_ea[inmask].mean()),4) if inmask.any() else None}}
    for dname,a,b in (("ar_en",lhc_ae[c],hit_ae[c]),("en_ar",lhc_ea[c],hit_ea[c])):
        b01=int(np.sum(~a&b)); b10=int(np.sum(a&~b)); n=b01+b10
        p=float(binomtest(b01,n,0.5).pvalue) if n>0 else 1.0
        res[f"mcnemar_clean_{dname}"]={"delta_lhc_minus_labse":round(float(a.mean()-b.mean()),4),
          "b01_labse_only":b01,"b10_lhc_only":b10,"p":round(p,6) if n>0 else None,"n_discordant":n}
    # فجوة التحكم لـLaBSE (مموّه−نظيف) — على نفس عناصر e5
    res["labse_gap_masked_minus_clean"]={"ar_en":round(float(hit_ae[inmask].mean()-hit_ae[c].mean()),4) if inmask.any() else None,
      "en_ar":round(float(hit_ea[inmask].mean()-hit_ea[c].mean()),4) if inmask.any() else None}
    out["subsets"][cn]=res
    store[f"{cn}_labse_ar_en"]=hit_ae; store[f"{cn}_labse_en_ar"]=hit_ea; store[f"{cn}_items"]=np.array(items)
    log(f"  LaBSE: كامل {res['labse']['r1_ar_en']}/{res['labse']['r1_en_ar']} | نظيف {res['labse']['clean_ar_en']}/{res['labse']['clean_en_ar']} | LHC نظيف {res['lhc_w4']['clean_ar_en']}/{res['lhc_w4']['clean_en_ar']}")
    for dname in ("ar_en","en_ar"):
        m=res[f"mcnemar_clean_{dname}"]; log(f"  McNemar-نظيف [{dname}]: Δ={m['delta_lhc_minus_labse']:+.4f} p={m['p']} (b01={m['b01_labse_only']},b10={m['b10_lhc_only']})")
np.savez_compressed(f"{R}/baseline_blocks_labse.npz", **store)
json.dump(out, open(f"{R}/baseline_labse_blocks.json","w"), ensure_ascii=False, indent=1)
log("SAVED results/baseline_labse_blocks.json + baseline_blocks_labse.npz")