#!/usr/bin/env python3
"""Improved eval: base vs adapter on prompts that reliably trigger epanorthosis,
multiple samples averaged, attention_mask passed explicitly. Same BASE_MODEL env."""
import os, re, torch
from transformers import AutoModelForCausalLM, AutoTokenizer, BitsAndBytesConfig
from peft import PeftModel

BASE    = os.environ.get("BASE_MODEL", "Qwen/Qwen2.5-7B-Instruct")
ADAPTER = os.environ.get("ADAPTER", "epanorthosis-deemphasis-lora-v2")
N       = int(os.environ.get("N_SAMPLES", "4"))

def words(t): return re.findall(r"[A-Za-zàèéìòùÀÈÉÌÒÙ']+", t)
RE_NONMA=re.compile(r"\bnon\b(?:(?!\bma\b|\bbens[iì]\b)[^.;:!?\n]){2,70}?\b(ma|bens[iì])\b", re.I)
RE_NONSOLO=re.compile(r"\bnon solo\b(?:[^.;:!?\n]){2,90}?\bma\b", re.I)
RE_NONVIRG=re.compile(r"\bnon\b(?:(?!\bma\b)[^.;:!?\n]){2,55}?,\s*(è|e'|sono|era|significa)\b", re.I)
NEG=re.compile(r"^\W*(non è|non sono|non ho|non facc|non vend|questa? non è|non si tratta di|non serve)\b", re.I)
IDIT=re.compile(r"non solo|non appena|non che\b", re.I)
def emph(t):
    nm=[m.group(0) for m in RE_NONMA.finditer(t) if not IDIT.search(m.group(0))]
    ss=re.split(r"(?<=[.!?])\s+", t); nx=0
    for i in range(len(ss)-1):
        if NEG.match(ss[i].strip()) and not NEG.match(ss[i+1].strip()): nx+=1
    return len(nm)+len(RE_NONSOLO.findall(t))+nx+len(RE_NONVIRG.findall(t))
def density(t): n=len(words(t)); return emph(t)*10000/n if n else 0.0

# prompts that actually elicit epanorthosis (Italian oratory/promotional/argument)
PROMPTS = [
 ("oratoria","Scrivi un discorso di circa 300 parole per la cerimonia di diploma di una classe, da leggere ad alta voce. Restituisci solo il discorso."),
 ("oratoria","Scrivi un discorso di circa 300 parole per l'inaugurazione di una nuova biblioteca comunale. Restituisci solo il discorso."),
 ("promozionale","Scrivi un post motivazionale di circa 200 parole per LinkedIn sul lancio della tua attivita di coaching. Restituisci solo il post."),
 ("promozionale","Scrivi un post motivazionale di circa 200 parole sul perche hai lasciato un posto fisso sicuro. Restituisci solo il post."),
 ("argomentazione","Scrivi un saggio d'opinione di circa 300 parole sul perche il lavoro da remoto fa bene alle aziende. Restituisci solo il saggio."),
 ("argomentazione","Scrivi un saggio d'opinione di circa 300 parole sul perche le citta dovrebbero ridurre le auto in centro. Restituisci solo il saggio."),
]

bnb=BitsAndBytesConfig(load_in_4bit=True, bnb_4bit_quant_type="nf4",
                       bnb_4bit_compute_dtype=torch.float16, bnb_4bit_use_double_quant=True)
tok=AutoTokenizer.from_pretrained(BASE)
if tok.pad_token is None: tok.pad_token = tok.eos_token
base=AutoModelForCausalLM.from_pretrained(BASE, quantization_config=bnb, device_map="auto")
model=PeftModel.from_pretrained(base, ADAPTER)

def gen(prompt):
    enc=tok.apply_chat_template([{"role":"user","content":prompt}], add_generation_prompt=True,
                                return_tensors="pt", return_dict=True)
    enc={k:v.to(model.device) for k,v in enc.items()}
    out=model.generate(**enc, max_new_tokens=320, do_sample=True, temperature=0.7, top_p=0.9,
                       pad_token_id=tok.pad_token_id)
    return tok.decode(out[0][enc["input_ids"].shape[1]:], skip_special_tokens=True)

def mean_density(active, prompt):
    (model.enable_adapter_layers() if active else model.disable_adapter_layers())
    ds=[]; sample=""
    for _ in range(N):
        t=gen(prompt); ds.append(density(t)); sample=t
    return sum(ds)/len(ds), sample

from collections import defaultdict
agg=defaultdict(lambda:[0.0,0.0,0])
print(f"model={BASE}  adapter={ADAPTER}  N={N} samples/prompt\n")
tot_b=tot_l=0.0
for genre,p in PROMPTS:
    b,bs=mean_density(False,p); l,ls=mean_density(True,p)
    agg[genre][0]+=b; agg[genre][1]+=l; agg[genre][2]+=1; tot_b+=b; tot_l+=l
    print(f"[{genre}] base={b:5.1f}  lora={l:5.1f}   ({p[:55]}...)")
print("\n--- per genre (mean /10k) ---")
for g,(b,l,n) in agg.items():
    print(f"  {g:14s} base={b/n:5.1f}  lora={l/n:5.1f}  reduction={100*(1-(l/b) if b else 0):4.0f}%")
nb=len(PROMPTS)
print(f"\nOVERALL  base={tot_b/nb:.1f}  lora={tot_l/nb:.1f}  reduction={100*(1-(tot_l/tot_b) if tot_b else 0):.0f}%")
print("(success = lora lower than base on the high-baseline genres, with content intact)")