#!/usr/bin/env python3
"""Calibration via the adapter 'dial' (paper §7.2/§7.7).
Measures the Epanorthosis Index at several LoRA strengths (alpha) to find the
setting that lands NEAR the human genre rate instead of at zero.
Run after the v2 adapter is trained. Same BASE_MODEL / ADAPTER env vars."""
import os, re, torch
from transformers import AutoModelForCausalLM, AutoTokenizer, BitsAndBytesConfig
from peft import PeftModel
from peft.tuners.lora import LoraLayer

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", "6"))
ALPHAS  = [0.0, 0.25, 0.5, 0.75, 1.0]
HUMAN   = {"oratoria": 14.0, "argomentazione": 12.0}   # promotional: no human baseline

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 = [
 ("oratoria","Scrivi un discorso di circa 250 parole per la cerimonia di diploma di una classe."),
 ("oratoria","Scrivi un discorso di circa 250 parole per l'inaugurazione di una biblioteca comunale."),
 ("promozionale","Scrivi un post motivazionale di circa 180 parole per LinkedIn sul lancio della tua attivita di coaching."),
 ("promozionale","Scrivi un post motivazionale di circa 180 parole sul perche hai lasciato un posto fisso."),
 ("argomentazione","Scrivi un saggio d'opinione di circa 250 parole sul perche il lavoro da remoto fa bene alle aziende."),
 ("argomentazione","Scrivi un saggio d'opinione di circa 250 parole sul perche le citta dovrebbero ridurre le auto in centro."),
]

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)
model.enable_adapter_layers()

def set_scale(mult):
    for m in model.modules():
        if isinstance(m, LoraLayer):
            for a in list(m.scaling.keys()):
                m.scaling[a] = (m.lora_alpha[a] / m.r[a]) * mult

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=300, 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)

print(f"model={BASE} adapter={ADAPTER} N={N}\n")
print(f"{'alpha':>6} | {'overall':>7} | {'oratoria':>8} | {'promoz.':>8} | {'argom.':>8}")
print("-"*52)
for al in ALPHAS:
    set_scale(al)
    from collections import defaultdict
    g=defaultdict(lambda:[0.0,0]); tot=0.0; ntot=0
    for genre,p in PROMPTS:
        ds=sum(density(gen(p)) for _ in range(N))/N
        g[genre][0]+=ds; g[genre][1]+=1; tot+=ds; ntot+=1
    ov=tot/ntot
    o=g["oratoria"][0]/g["oratoria"][1]; pr=g["promozionale"][0]/g["promozionale"][1]; ar=g["argomentazione"][0]/g["argomentazione"][1]
    print(f"{al:>6.2f} | {ov:7.1f} | {o:8.1f} | {pr:8.1f} | {ar:8.1f}")
print(f"\nhuman baseline (target): oratoria~{HUMAN['oratoria']}, argomentazione~{HUMAN['argomentazione']} /10k")
print("Pick the alpha whose EI lands nearest the human rate (calibration, not elimination).")