"""A/B test corrigé : GLM-OCR sait-il distinguer accord vs désaccord sur un crop ?"""
from pathlib import Path
from PIL import Image
from pipeline.ocr_glm import GLMOCR
from pipeline.ingest import pdf_to_images

ZONE = (0.55, 0.82, 0.92, 0.90)

# Ground truth vérifié visuellement + confirmé par legacy
CASES = [
    ("2018 CARC/OGC 7.pdf",  1, "accord"),
    ("2018 CARC/OGC 55.pdf", 1, "accord"),
    ("2018 CARC/OGC 27.pdf", 1, "accord"),
    ("2018 CARC/OGC 86.pdf", 1, "désaccord"),
    ("2018 CARC/OGC 9.pdf",  1, None),  # ground truth à vérifier
    ("2018 CARC/OGC 43.pdf", 1, None),
    ("2018 CARC/OGC 68.pdf", 1, None),
]

PROMPTS = {
    "P1_simple":      """Quelle case est cochée (X) : "Accord" ou "Désaccord" ? Un seul mot.""",
    "P2_json":        """Renvoie UNIQUEMENT {"cochee":"accord"} ou {"cochee":"désaccord"} selon la case marquée d'une croix.""",
    "P3_negatif":     """Sur cette image, il y a DEUX cases. L'UNE est cochée par une croix noire, l'AUTRE est vide. Dis-moi laquelle est cochée. Réponds par un seul mot : accord ou désaccord.""",
    "P4_explicite":   """Regarde ATTENTIVEMENT les deux cases à cocher. Une case VIDE ressemble à [ ]. Une case COCHÉE ressemble à [X]. Quelle case est cochée ? Réponds UNIQUEMENT : accord ou désaccord.""",
}


def crop_rel(img, z):
    w, h = img.size
    return img.crop((int(z[0]*w), int(z[1]*h), int(z[2]*w), int(z[3]*h)))


def classify(txt):
    low = txt.lower()
    # Retirer les séquences "accord_desaccord" et "accord ou désaccord"
    low = low.replace("accord_desaccord", "").replace("accord ou désaccord", "").replace("accord ou desaccord", "")
    has_des = "désaccord" in low or "desaccord" in low
    has_acc = "accord" in low
    if has_des and not has_acc: return "désaccord"
    if has_acc and not has_des: return "accord"
    if has_acc and has_des:     return "both"
    return "?"


def main():
    ocr = GLMOCR()
    print(f"VRAM = {ocr.vram_gb:.2f} Go\n")

    Path("/tmp/ogc_crops").mkdir(exist_ok=True)
    scores = {name: [0, 0] for name in PROMPTS}  # [hits, evaluated]
    for pdf, page, expected in CASES:
        images = pdf_to_images(pdf)
        img = Image.open(images[page - 1])
        crop = crop_rel(img, ZONE)
        crop_path = f"/tmp/ogc_crops/{Path(pdf).stem.replace(' ', '_')}_cb.png"
        crop.save(crop_path)
        label = f"OGC {Path(pdf).stem.split()[-1]}"
        print(f"=== {label} (attendu={expected}) ===")
        for name, prompt in PROMPTS.items():
            res = ocr.run(crop_path, prompt, max_new_tokens=48)
            got = classify(res["text"])
            verdict = ""
            if expected:
                if got == expected: scores[name][0] += 1; verdict = " ✓"
                else: verdict = " ✗"
                scores[name][1] += 1
            print(f"  [{name}] → {got:10s} (raw={res['text'].strip()[:80]!r}){verdict}")
        print()

    print("=== RÉCAPITULATIF (sur cas à ground truth vérifié) ===")
    for name, (h, n) in scores.items():
        print(f"  {name:18s}: {h}/{n}")


if __name__ == "__main__":
    main()