#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""robust_clean.py — المتانة بعشر بذور على ALL-نظيف (Tanzil-نظيف + Bible-نظيف فقط؛
TED مستثناة لأن النظيف منها 154 عنصرًا فقط — تُذكر كقيد)."""
import sys, re, random, json
sys.path.insert(0,"/home/user/lhc/code"); sys.path.insert(0,"/home/user/lhc/recon")
import numpy as np
from phase_w1_data import download_corpus
from phase_c_encoder import TrainableEncoder
from phase_p2_slot import hashed_features_slot
from lhc_core import LHC0

DD="/home/user/lhc/data/opus"
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)]
_AR_LETTER=re.compile(r"[\u0621-\u064A]")
_HARAKAT=["\u064E","\u064F","\u0650","\u0652","\u0651","\u064B","\u064C","\u064D"]
_DIGIT_MAP=str.maketrans("0123456789","\u0660\u0661\u0662\u0663\u0664\u0665\u0666\u0667\u0668\u0669")
def distort_kind(t,nrs,kind):
    if kind=="tashkeel":
        out=[]
        for ch in t:
            out.append(ch)
            if _AR_LETTER.match(ch) and nrs.rand()<0.28: out.append(_HARAKAT[int(nrs.randint(len(_HARAKAT)))])
        return "".join(out)
    if kind=="ortho":
        for a,b in (("\u0623","\u0627"),("\u0625","\u0627"),("\u0622","\u0627"),("\u0671","\u0627"),("\u0629","\u0647"),("\u0649","\u064A")): t=t.replace(a,b)
        return t.translate(_DIGIT_MAP)
    if kind=="tatweel":
        ws=t.split(" ")
        for i in range(len(ws)):
            w=ws[i]
            if len(w)>3 and nrs.rand()<0.25:
                j=int(nrs.randint(1,len(w)-2)); ws[i]=w[:j]+"\u0640"+w[j:]
        return " ".join(ws)
    if kind=="combined": return distort_kind(distort_kind(t,nrs,"tashkeel"),nrs,"ortho")
    return t

M=json.load(open("/home/user/lhc/results/masks_v2.json"))
pool=[]
for name,url,lim in HELDOUT:
    ps=download_corpus(name,url,lim,DD)
    mask=set(M["masks_strong"][name])
    pool += [p for i,p in enumerate(ps) if i not in mask]
print(f"ALL-نظيف (Tanzil+Bible): {len(pool):,} عنصرًا")

def acc(enc,queries,targets):
    lhc=LHC0(kb=[{"prompt":t,"response":str(i),"domain":"x"} for i,t in enumerate(targets)],encoder=enc)
    ok=0
    for i,q in enumerate(queries):
        e,_,_=lhc.memory.read(q); ok+=int(e["prompt"]==targets[i])
    return ok/max(1,len(queries))

encs={}
for tag,wp in [("W4","weights/phase_w1_W.npz"),("W4.1","weights/phase_w1_W_robust.npz")]:
    z=np.load(wp); encs[tag]=TrainableEncoder(F=int(z["F"]),d=int(z["d"]),tau=float(z["tau"]),W=z["W"],features_fn=hashed_features_slot)

rng=random.Random(41); samp=rng.sample(pool,300)
ar=[p["ar"] for p in samp]; en=[p["en"] for p in samp]
out={"pool":"ALL-clean (Tanzil+Bible clean)","n_pool":len(pool),"n_sample":300,"sample_seed":41,"seeds":list(range(10)),
     "note":"TED مستثناة (النظيف 154 عنصرًا) — قيد معلن"}
for kind in ("clean","tashkeel","ortho","tatweel","combined"):
    for tag in ("W4","W4.1"):
        a1s,a2s=[],[]
        for cs in range(10):
            if kind=="clean": ar_d=ar
            else:
                nrs=np.random.RandomState(cs); ar_d=[distort_kind(a,nrs,kind) for a in ar]
            a1s.append(acc(encs[tag],ar_d,en)); a2s.append(acc(encs[tag],en,ar_d))
            if kind=="clean": break
        m1,s1=float(np.mean(a1s)),float(np.std(a1s)); m2,s2=float(np.mean(a2s)),float(np.std(a2s))
        out[f"{kind}|{tag}"]={"mean_ar_en":round(m1,4),"std_ar_en":round(s1,4),"mean_en_ar":round(m2,4),"std_en_ar":round(s2,4)}
        print(f"[{kind:9s} {tag:4s}] {m1:.3f}±{s1:.3f} / {m2:.3f}±{s2:.3f}", flush=True)
json.dump(out, open("/home/user/lhc/results/robust_clean.json","w"), ensure_ascii=False, indent=1)
print("SAVED results/robust_clean.json")