pipeline_multipage.py

# -*- coding: utf-8 -*-
"""pipeline_multipage.ipynb

Automatically generated by Colab.

Original file is located at
    https://colab.research.google.com/drive/1kVRCAPJDDpeWj2ILhDiEPEGZu0-QA3-d

## SETUP

Monto Drive e definisco i percorsi per MP-DocVQA. In più rispetto alla
versione single-page, definisco WINDOW_SIZE: il numero di pagine del
documento che passo al modello, centrate sulla pagina che contiene la
risposta (W=1 equivale al caso single-page).
"""

# ══ SETUP GLOBALE ══ CELLA DA RUNNARE SEMPRE


from google.colab import drive
drive.mount('/content/drive')

import json, os, re, random, tarfile, gc
from collections import Counter, defaultdict

SEED = 42
random.seed(SEED)

# ── Percorsi ──
DRIVE_DIR  = '/content/drive/MyDrive/MP-docvqa'
QAS_DIR    = f'{DRIVE_DIR}/qas'
IMG_DIR    = f'{DRIVE_DIR}/qas/images'          # ← immagini già su Drive, no tar
BENCH_PATH = f'{DRIVE_DIR}/benchmark_final.json'
PRED_DIR   = f'{DRIVE_DIR}/predictions'
os.makedirs(PRED_DIR, exist_ok=True)

# Bilanciamento ragionevole
N_ENTITY  = 34   # ~33% delle corrotte
N_ELEMENT = 33   # ~33% delle corrotte
N_LAYOUT  = 33   # ~33% delle corrotte
N_CLEAN   = 100  # 50/50 con le corrotte

# Totale: 100 corrotte + 100 clean = 200 campioni


# ── Parametro finestra MP-DocVQA ──
# 1 = solo pagina della risposta (equivalente a single-page)
# 2 = pagina risposta + una pagina adiacente
# 3 = finestra più ampia
WINDOW_SIZE = 2

"""## CARICAMENTO DATI

Carico le annotazioni di MP-DocVQA (train + validation). A differenza di
SP-DocVQA qui non c'è il campo question_types: aggiungo un placeholder per
riusare lo stesso campionamento stratificato della versione single-page.
"""

# ──  Caricamento annotazioni MP-DocVQA (train + val) ──

with open(f'{QAS_DIR}/train.json') as f:
    data_train = json.load(f)
with open(f'{QAS_DIR}/val.json') as f:
    data_val = json.load(f)

all_data = data_train['data'] + data_val['data']

# MP-DocVQA non ha question_types — aggiunge placeholder
for s in all_data:
    if 'question_types' not in s:
        s['question_types'] = ['unknown']

print(f"Train: {len(data_train['data'])} | Val: {len(data_val['data'])} | Totale: {len(all_data)}")

print("\nEsempio di campione:")
campi = ('questionId', 'question', 'page_ids', 'answers', 'answer_page_idx')
esempio = {k: all_data[0][k] for k in campi}
print(json.dumps(esempio, indent=2))

print("\nCampi disponibili:", list(all_data[0].keys()))

"""## PARTE 1 — CORRUPTION ENGINE

Stessa logica della versione single-page (le stesse 3 funzioni: entity,
element, layout). Cambia solo come ricavo l'immagine da mostrare al
modello: prendo la pagina che contiene la risposta da
page_ids[answer_page_idx].
"""

# ──  Corruption Engine ──


# Regola fondamentale: la corruzione rende l'informazione richiesta ASSENTE dal documento.
# La domanda rimane grammaticalmente corretta,  l'unanswerability si scopre SOLO guardando l'immagine.

import spacy

# Carica il modello NLP di spaCy per riconoscere le entità nelle domande (date, luoghi, ecc.)
# Se non è installato, lo scarica automaticamente.
try:
    nlp = spacy.load("en_core_web_sm")
except OSError:
    import subprocess
    subprocess.run(["python", "-m", "spacy", "download", "en_core_web_sm"], check=True)
    nlp = spacy.load("en_core_web_sm")

########################################   ENTITY    ####################################
ENTITY_POOLS = {
    "DATE":     ["March 2027", "1965", "April 12, 2031", "the year 2044",
                 "September 1958", "February 2003", "July 1999", "October 2015"],
    "GPE":      ["Reykjavik", "Montevideo", "Tashkent", "Wellington",
                 "Nairobi", "Bratislava", "Ulaanbaatar"],
    "LOC":      ["the Andes", "the Baltic Sea", "the Sahara", "the Pyrenees", "the Gobi Desert"],
    "ORG":      ["Zentech Holdings", "Marwick & Vale", "the Halverson Institute",
                 "Penrose Corp", "Valtrade Group", "Nordex Solutions"],
    "PERSON":   ["Eleanor Whitfield", "Marcus Delgado", "Priya Ramachandran",
                 "Tobias Lindqvist", "Fatima Al-Rashid", "Chen Wei"],
    "CARDINAL": ["847", "13", "2,915", "66", "1,204", "342", "7,800", "29"],
    "MONEY":    ["$3,720", "$48.50", "$192,000", "$7.25", "$5,430", "$0.99"],
    "PERCENT":  ["73%", "12.4%", "5%", "88%", "31%", "6.7%"],
    "ORDINAL":  ["seventh", "twenty-third", "ninth", "fifteenth", "fourth"],
    "TIME":     ["3:45 PM", "11:20 AM", "midnight", "6:00 AM", "9:30 PM"],
    "QUANTITY": ["12 kilometers", "3.5 liters", "80 kg", "450 grams", "2.3 miles"],
}


# Parole interrogative: evita di sostituire "when" o "what" scambiandoli per entità
INTERROGATIVE = {"which", "what", "when", "how", "where", "who", "many", "much"}

# Numeri generici che spaCy classifica come CARDINAL ma non sono entità significative
# Es: "one table", "few rows" — sostituirli non renderebbe la domanda unanswerable
CARDINAL_BLACKLIST = {
    "one", "two", "three", "four", "five", "six", "seven", "eight",
    "nine", "ten", "few", "many", "several", "some", "any", "no"
}

# Pattern da escludere:
# - "year-old" / "years-old": età, non date (es. "a 5 year-old company")
# - r"^[A-Z]{2,5}$": acronimi puri come "USA", "CEO" — spaCy li vede come ORG ma non sono sostituibili
BLACKLIST_PATTERNS = ["year-old", "years-old", r"^[A-Z]{2,5}$"]



def blacklisted(text):
    """
    Restituisce True se l'entità NON deve essere sostituita.
    Controlla due casi:
      - pattern regex (iniziano con ^): es. acronimi come "USA", "CEO"
      - sottostringa semplice: es. "year-old" dentro "5 year-old company"
    """
    for pattern in BLACKLIST_PATTERNS:
        if pattern.startswith("^"):
            # È un pattern regex: controlla se il testo corrisponde esattamente
            if re.match(pattern, text):
                return True
        else:
            # È una parola semplice: controlla se compare nel testo
            if pattern in text.lower():
                return True
    return False


def entity_candidates(doc):
    """
    Restituisce la lista delle entità della domanda che possono essere sostituite.
    Un'entità è candidata solo se supera tutti e 4 i filtri:
    """
    candidates = []

    for ent in doc.ents:

        # 1) Il tipo di entità deve essere tra quelli che sappiamo sostituire
        if ent.label_ not in ENTITY_POOLS:
            continue

        # 2) La prima parola non deve essere interrogativa (es. "what", "when")
        #    spaCy a volte classifica queste parole come entità per errore
        if ent.text.split()[0].lower() in INTERROGATIVE:
            continue

        # 3) Il testo non deve essere in blacklist (acronimi, età, ecc.)
        if blacklisted(ent.text):
            continue

        # 4) I CARDINAL generici ("one", "few", "many"...) non sono entità significative
        if ent.label_ == "CARDINAL" and ent.text.lower() in CARDINAL_BLACKLIST:
            continue

        candidates.append(ent)

    return candidates

def corrupt_entity(question, doc=None):
    """
    Sostituisce una entità nella domanda con un valore diverso dello stesso tipo.
    Restituisce (domanda_corrotta, descrizione_della_modifica) oppure None se non è possibile.
    """

    # Analizza la domanda con spaCy (se non è già stata analizzata)
    if doc is None:
        doc = nlp(question)

    # Trova le entità sostituibili
    candidates = entity_candidates(doc)
    if not candidates:
        return None

    # Sceglie una entità a caso tra i candidati
    ent = random.choice(candidates)

# Prende i valori del pool escludendo il valore originale
    # (per evitare di "sostituire" con la stessa cosa)
    pool = []
    for v in ENTITY_POOLS[ent.label_]:
        if v.lower() != ent.text.lower():
            pool.append(v)

    # Se il pool è vuoto non si può corrompere
    if not pool:
        return None

# Controllo di sicurezza: il carattere subito dopo l'entità deve essere
    # uno spazio o punteggiatura,  altrimenti la sostituzione produrrebbe
    # una parola "incollata" al testo successivo (es. "2019report")

    # Se l'entità è l'ultima parola della domanda, non c'è nessun carattere dopo
    if ent.end_char < len(question):
        char_after_entity = question[ent.end_char]
    else:
        char_after_entity = " "  # spazio fittizio: nessun problema di incollamento

    if char_after_entity not in " .,;:?!\n)'\"":
        return None


    # Sostituisce l'entità originale con il nuovo valore
    new_val = random.choice(pool)
    corrupted_question = question[:ent.start_char] + new_val + question[ent.end_char:]

    # Descrizione della modifica (utile per il debug e la relazione)
    detail = f"{ent.text} ({ent.label_}) -> {new_val}"

    return corrupted_question, detail


    ########################################   ELEMENTS    ####################################



# ELEMENT: gruppi semanticamente lontani. Sostituire tra gruppi DIVERSI riduce il rischio
# che l'elemento sostitutivo sia comunque presente nella pagina (label noise).
ELEMENT_GROUPS = [
    ["table", "row", "column", "sheet", "field"],  # elementi tabulari
    ["figure", "chart", "graph", "diagram"],       # elementi visivi
    ["footnote", "caption", "header", "heading", "section"],  # elementi testuali
]
# Lista piatta di tutti gli elementi di pagina (tutti i gruppi uniti in una lista sola)
ELEMENTS = []
for group in ELEMENT_GROUPS:
    for element in group:
        ELEMENTS.append(element)





def corrupt_element(question):
    # Trova tutti gli elementi di pagina presenti nella domanda
    found = []
    for element in ELEMENTS:
        if re.search(rf"\b{element}\b", question, re.I):
            found.append(element)

    if not found:
        return None

    # Sceglie a caso quale elemento sostituire
    target = random.choice(found)

    # Raccoglie tutti gli elementi di gruppi DIVERSI da quello del target
    other_group_elements = []
    for group in ELEMENT_GROUPS:
        if target not in group:
            other_group_elements.extend(group)

    replacement = random.choice(other_group_elements)

    corrupted_question = re.sub(rf"\b{target}\b", replacement, question, count=1, flags=re.I)
    detail = f"{target} -> {replacement}"
    return corrupted_question, detail

    ########################################   LAYOUT    ####################################


# LAYOUT: opposti spaziali + posizioni ordinali legate a nomi strutturali
LAYOUT_OPPOSITES = {
    "top left": ["bottom right"], "top right": ["bottom left"],
    "bottom left": ["top right"], "bottom right": ["top left"],
    "upper left": ["lower right"], "upper right": ["lower left"],
    "lower left": ["upper right"], "lower right": ["upper left"],
    "top": ["bottom"], "bottom": ["top"], "left": ["right"], "right": ["left"],
    "upper": ["lower"], "lower": ["upper"], "above": ["below"], "below": ["above"],
    "corner": ["center"],
}
POSITIONAL  = {"first": ["last","third"], "last": ["first","third"],
               "second": ["last","first"], "third": ["first","last"]}
STRUCT_NOUN = r"(row|column|page|section|paragraph|line|item|entry|table|figure)"

def corrupt_layout(question):
    """
    Sostituisce un riferimento spaziale o ordinale nella domanda con il suo opposto.
    Restituisce (domanda_corrotta, descrizione) oppure None se non trova nulla da corrompere.
    """
    q = question.lower()

    # ── Parte 1: termini spaziali (top, bottom, left, right, ecc.) ──
    # Scorre dal più lungo al più corto per trovare prima "top left" rispetto a "top"
    for term in sorted(LAYOUT_OPPOSITES, key=len, reverse=True):
        term_escaped = re.escape(term)

        # Il termine deve essere seguito da una parola strutturale del documento
        # per escludere usi generici come "top priority" o "left alone"
        is_position_reference = re.search(
            rf"\b{term_escaped}\s+(the\s+)?(of|corner|part|side|margin|page|section|table|figure)\b", q
        )

        if is_position_reference:
            replacement = random.choice(LAYOUT_OPPOSITES[term])
            corrupted_question = re.sub(rf"\b{term_escaped}\b", replacement, question, count=1, flags=re.I)
            detail = f"{term} -> {replacement}"
            return corrupted_question, detail

    # ── Parte 2: posizioni ordinali + nome strutturale (es. "first row", "last column") ──
    for word, alternatives in POSITIONAL.items():
        if re.search(rf"\b{word}\b\s+{STRUCT_NOUN}", q):
            replacement = random.choice(alternatives)
            corrupted_question = re.sub(rf"\b{word}\b", replacement, question, count=1, flags=re.I)
            detail = f"{word} -> {replacement}"
            return corrupted_question, detail

    return None

# Dizionario che mappa ogni tipo di corruzione alla sua funzione
CORRUPTORS = {
    "entity":  corrupt_entity,
    "element": corrupt_element,
    "layout":  corrupt_layout,
}

"""## CORRUZIONE + CAMPIONAMENTO

Applico le tre funzioni di corruzione con campionamento stratificato e
campiono in parallelo le domande clean, per arrivare a un benchmark
bilanciato 50/50 (stessa logica del single-page).
"""

# ──  Pool di candidati, campionamento stratificato e applicazione delle corruzioni ──
random.seed(SEED)

questions = [s["question"] for s in all_data]

pools = {
    "entity":  [],
    "element": [],
    "layout":  []
}

for doc, s in zip(nlp.pipe(questions, batch_size=512), all_data):
    if entity_candidates(doc):
        pools["entity"].append(s)
    if corrupt_element(s["question"]):
        pools["element"].append(s)
    if corrupt_layout(s["question"]):
        pools["layout"].append(s)

print("Domande corrompibili per tipo:", {k: len(v) for k, v in pools.items()})

used_qids      = set()
seen_questions = set()


def sample_and_corrupt(pool, ctype, n):
    corruption_fn = CORRUPTORS[ctype]

    by_type = defaultdict(list)
    for s in pool:
        by_type[s['question_types'][0]].append(s)

    quota   = max(1, n // len(by_type))
    records = []

    def try_add(s):
        question_normalized = re.sub(r'\s+', ' ', s['question'].lower().strip())

        if s['questionId'] in used_qids or question_normalized in seen_questions:
            return False

        result = corruption_fn(s["question"])
        if result is None:
            return False

        corrupted_question, detail = result

        used_qids.add(s['questionId'])
        seen_questions.add(question_normalized)

        # ── MP-DocVQA: ricava la pagina della risposta da page_ids ──
        answer_page = s["page_ids"][s["answer_page_idx"]]

        records.append({
            **s,
            "image":             answer_page,        # pagina che contiene la risposta
            "question":          corrupted_question,
            "original_question": s["question"],
            "corruption_type":   ctype,
            "corruption_detail": detail,
            "label":             "unanswerable"
        })
        return True

    for question_type, samples in by_type.items():
        added = 0
        for s in random.sample(samples, len(samples)):
            if added == quota or len(records) == n:
                break
            if try_add(s):
                added += 1

    for s in random.sample(pool, len(pool)):
        if len(records) == n:
            break
        try_add(s)

    return records


entity_data  = sample_and_corrupt(pools["entity"],  "entity",  N_ENTITY)
element_data = sample_and_corrupt(pools["element"], "element", N_ELEMENT)
layout_data  = sample_and_corrupt(pools["layout"],  "layout",  N_LAYOUT)

corrupted_data = entity_data + element_data + layout_data

n_entity  = len([d for d in corrupted_data if d['corruption_type'] == 'entity'])
n_element = len([d for d in corrupted_data if d['corruption_type'] == 'element'])
n_layout  = len([d for d in corrupted_data if d['corruption_type'] == 'layout'])

print(f"Entity:  {n_entity}/{N_ENTITY}")
print(f"Element: {n_element}/{N_ELEMENT}")
print(f"Layout:  {n_layout}/{N_LAYOUT}")


# ── Campionamento clean ──
clean_pool = [s for s in all_data if s['questionId'] not in used_qids]

by_type = defaultdict(list)
for s in clean_pool:
    by_type[s['question_types'][0]].append(s)

clean_data = []
quota = N_CLEAN // len(by_type)
for question_type, samples in by_type.items():
    n_to_take = min(quota, len(samples))
    clean_data += random.sample(samples, n_to_take)

already_used = {c['questionId'] for c in clean_data}
extra_pool   = [s for s in clean_pool if s['questionId'] not in already_used]
remaining    = N_CLEAN - len(clean_data)
clean_data  += random.sample(extra_pool, min(remaining, len(extra_pool)))

# ── MP-DocVQA: aggiunge campo image anche per i clean ──
clean_data = [
    {
        **s,
        "image":           s["page_ids"][s["answer_page_idx"]],
        "label":           "answerable",
        "corruption_type": "clean"
    }
    for s in clean_data
]

print(f"\nCorrotte: {len(corrupted_data)} | Clean: {len(clean_data)}")
print("Distribuzione corruzioni:", dict(Counter(d['corruption_type'] for d in corrupted_data)))

for ctype in ["entity", "element", "layout"]:
    esempio = next(d for d in corrupted_data if d["corruption_type"] == ctype)
    print(f"\n[{ctype}]")
    print(f"  Originale : {esempio['original_question']}")
    print(f"  Corrotta  : {esempio['question']}")
    print(f"  Modifica  : {esempio['corruption_detail']}")
    print(f"  Pagina    : {esempio['image']}")

"""## LLM AS A JUDGE

Verifica qualità con LLM-as-a-judge. Il judge (LLaMA-3.3-70B via Groq) è
volutamente text-only, non vede l'immagine: una corruzione è valida solo se

- well_formed : la domanda è grammaticalmente corretta e naturale;
- needs_document : dal solo testo è impossibile capire che è unanswerable,
  serve ispezionare il documento.

Se un criterio fallisce, la domanda viene riscritta automaticamente (max 1
tentativo) e ri-giudicata.
"""

# ──  LLM-as-a-judge (Groq, llama-3.3-70b) ──
!pip install groq -q

from groq import Groq
from google.colab import userdata
import time

# Inizializza il client Groq con la chiave API
groq_client = Groq(api_key=userdata.get("GROQ_API_KEY"))

# ── Prompt di sistema: ruolo del judge ──
JUDGE_SYSTEM = "You are an expert annotator for document question answering benchmarks."

# ── Prompt di valutazione ──
# Il judge valuta la domanda corrotta su due criteri:
#   1. well_formed     -> la domanda è grammaticalmente corretta e naturale
#   2. needs_document  -> dal solo testo non si capisce che è unanswerable,
#                        bisogna guardare l'immagine del documento
JUDGE_PROMPT = """
Original question:  {orig}
Corrupted question: {corr}
Corruption type:    {ctype}

The corrupted question will be shown to a vision model together with a document image.
Evaluate it on two criteria:

1. well_formed: the question is grammatically correct and reads naturally.
2. needs_document: from the text alone it is impossible to tell the question is
   unanswerable — one must look at the document image to find out.

Reply ONLY with this JSON:
{{"well_formed": true, "needs_document": true, "reason": "..."}}
"""


# ── Prompt di riscrittura ──
# Se la domanda fallisce uno dei criteri, il judge prova a riscriverla
# mantenendo la corruzione ma migliorando forma e naturalezza
REWRITE_PROMPT = """
Original question:  {orig}
Corrupted question: {corr}
Corruption type:    {ctype}
Problem:            {reason}

Rewrite the corrupted question so that:
- it is grammatically correct and natural
- it PRESERVES the corruption (do not restore the original value)
- its unanswerability can only be verified by looking at the document image

Reply ONLY with this JSON:
{{"rewritten_question": "..."}}
"""


def call_api(messages, max_tokens=200, max_retries=3):
    """
    Chiama l'API di Groq e restituisce il testo della risposta.
    Se riceve un errore di rate limit (429), aspetta e riprova.
    """
    for attempt in range(max_retries):
        try:
            response = groq_client.chat.completions.create(
                model="llama-3.3-70b-versatile",
                messages=messages,
                max_tokens=max_tokens,
                temperature=0.0
            )
            return response.choices[0].message.content

        except Exception as e:
            if "429" in str(e):
                wait_seconds = 10 * (attempt + 1)
                print(f"Rate limit raggiunto, aspetto {wait_seconds}s...")
                time.sleep(wait_seconds)
            else:
                raise

    return None


def extract_json(text):
    """
    Estrae il primo oggetto JSON dalla risposta del judge.
    Restituisce None se non trova JSON valido.
    """
    match = re.search(r"\{.*\}", text or "", re.S)
    if match is None:
        return None
    try:
        return json.loads(match.group())
    except json.JSONDecodeError:
        return None


def evaluate(sample, question):
    """
    Chiede al judge di valutare la domanda corrotta.
    Restituisce un dizionario con well_formed, needs_document e reason.
    """
    messages = [
        {"role": "system", "content": JUDGE_SYSTEM},
        {"role": "user",   "content": JUDGE_PROMPT.format(
            orig  = sample["original_question"],
            corr  = question,
            ctype = sample["corruption_type"]
        )}
    ]

    raw_response = call_api(messages)
    result = extract_json(raw_response)

    if result is None:
        return {"well_formed": None, "needs_document": None, "reason": "PARSE_FAIL"}

    return result


def rewrite(sample, question, reason):
    """
    Chiede al judge di riscrivere la domanda che non ha superato la valutazione.
    Restituisce la domanda riscritta oppure None se la riscrittura fallisce.
    """
    messages = [
        {"role": "system", "content": JUDGE_SYSTEM},
        {"role": "user",   "content": REWRITE_PROMPT.format(
            orig   = sample["original_question"],
            corr   = question,
            ctype  = sample["corruption_type"],
            reason = reason
        )}
    ]

    raw_response = call_api(messages)
    result = extract_json(raw_response)

    if result is None:
        return None

    return result.get("rewritten_question")


def is_valid(verdict):
    """
    Restituisce True se la domanda ha superato entrambi i criteri del judge.
    """
    return verdict.get("well_formed") is True and verdict.get("needs_document") is True


# ── Valutazione di ogni domanda corrotta ──
MAX_REWRITES = 1  # al massimo 1 tentativo di riscrittura per domanda

for i, sample in enumerate(corrupted_data):

    current_question = sample["question"]
    n_rewrites = 0

    # Valutazione iniziale
    verdict = evaluate(sample, current_question)

    # Se non valida, prova a riscrivere una volta
    if not is_valid(verdict) and n_rewrites < MAX_REWRITES:
        rewritten = rewrite(sample, current_question, verdict.get("reason", ""))
        n_rewrites += 1

        if rewritten is not None:
            current_question = rewritten
            verdict = evaluate(sample, current_question)

    # Salva la domanda finale nel record
    if current_question != sample["question"]:
        sample["question_before_rewrite"] = sample["question"]
        sample["question"] = current_question

    # Salva il verdetto nel record
    sample["n_rewrites"]    = n_rewrites
    sample["judge_verdict"] = "KEEP" if is_valid(verdict) else "REJECT"
    sample["judge_reason"]  = verdict.get("reason", "")

    # Stampa il progresso
    tag = f"R{n_rewrites}" if n_rewrites > 0 else "  "
    print(f"[{i+1:02d}/{len(corrupted_data)}] {tag} {sample['judge_verdict']:6s} "
          f"({sample['corruption_type']}) {current_question[:60]}")

# Conta quanti KEEP e quanti REJECT
conteggio_verdetti = Counter(s['judge_verdict'] for s in corrupted_data)

print("\nVerdetti finali:")
print(f"  KEEP   : {conteggio_verdetti['KEEP']}")
print(f"  REJECT : {conteggio_verdetti['REJECT']}")

"""## BILANCIAMENTO E SALVATAGGIO SU DRIVE

Stessa procedura della versione single-page: tengo solo i KEEP del judge,
bilancio 50/50 con le clean, salvo benchmark_final.json. Le immagini di
MP-DocVQA (22GB) le ho già processate ed estratte in locale: qui verifico
solo che quelle necessarie al benchmark siano presenti su Drive.
"""

# ──  Bilanciamento e salvataggio del benchmark su Drive ──
random.seed(SEED)

# Tiene solo le domande approvate dal judge
corrupted_keep = [d for d in corrupted_data if d["judge_verdict"] == "KEEP"]
n_reject       = len(corrupted_data) - len(corrupted_keep)

# Rimuove altrettante clean per mantenere il bilanciamento 50/50
# Esempio: 5 REJECT -> rimuove 5 clean -> rimane sempre metà corrotte, metà clean
if n_reject > 0:
    clean_data = random.sample(clean_data, len(clean_data) - n_reject)

# Unisce corrotte e clean, mescola per non averle in ordine
benchmark_final = corrupted_keep + clean_data
random.shuffle(benchmark_final)

# Salva il benchmark su Drive (sovrascrive se esiste già)
with open(BENCH_PATH, "w") as f:
    json.dump(benchmark_final, f, indent=2)

print(f"Benchmark salvato in: {BENCH_PATH}")
print(f"  Unanswerable : {len(corrupted_keep)}")
print(f"  Answerable   : {len(clean_data)}")
print(f"  Totale       : {len(benchmark_final)}")

# ── Immagini del benchmark ──
# Il file tar completo (22GB) è stato processato in locale.
# Le immagini necessarie al benchmark sono state estratte
# e caricate su Drive in IMG_DIR.

IMG_DIR = f'{DRIVE_DIR}/qas/images'

# Verifica che le immagini necessarie siano presenti
needed = set()
for s in benchmark_final:
    page_ids   = s.get("page_ids", [s["image"]])
    answer_idx = s.get("answer_page_idx", 0)
    start = max(0, answer_idx - WINDOW_SIZE // 2)
    end   = min(len(page_ids), start + WINDOW_SIZE)
    for pid in page_ids[start:end]:
        needed.add(f"{pid}.jpg")

missing = {f for f in needed if not os.path.isfile(os.path.join(IMG_DIR, f))}
print(f"Immagini necessarie : {len(needed)}")
print(f"  presenti su Drive : {len(needed) - len(missing)}")
print(f"  mancanti          : {len(missing)}")
if missing:
    for m in sorted(missing)[:5]:
        print(f"  [!] {m}")

def wrap(text, width=45):
    return "\n".join(textwrap.wrap(text, width))

def show_grid(corruption_type, n=6):
    samples = [d for d in corrupted_keep if d["corruption_type"] == corruption_type][:n]

    fig, axes = plt.subplots(2, 3, figsize=(18, 16))
    fig.suptitle(f"Spot-check: {corruption_type.upper()}", fontweight='bold', fontsize=14)

    for ax, sample in zip(axes.flat, samples):

        # pagina della risposta (page_id già salvato nel campo "image")
        image_path = os.path.join(IMG_DIR, f"{sample['image']}.jpg")  # ← .jpg

        try:
            image = Image.open(image_path).convert('RGB')
            ax.imshow(image)
        except FileNotFoundError:
            ax.text(0.5, 0.5, "Immagine\nnon trovata",
                    ha='center', va='center', transform=ax.transAxes)

        ax.axis('off')
        ax.set_title(
            f"ORI: {wrap(sample['original_question'])}\n"
            f"\nCOR: {wrap(sample['question'])}\n"
            f"\nMOD: {wrap(sample['corruption_detail'])}\n"
            f"\nPAG: {sample['image']}\n"
            f"TOT pagine doc: {len(sample.get('page_ids', []))}",  # ← info multi-page
            fontsize=7,
            loc='left'
        )

    plt.tight_layout()
    plt.show()

for corruption_type in ["entity", "element", "layout"]:
    show_grid(corruption_type)

from PIL import Image
import matplotlib.pyplot as plt
import textwrap

def show_window(sample, window_size=2):
    page_ids   = sample["page_ids"]
    answer_idx = sample["answer_page_idx"]
    start = max(0, answer_idx - window_size // 2)
    end   = min(len(page_ids), start + window_size)
    window = page_ids[start:end]

    fig, axes = plt.subplots(1, len(window), figsize=(8 * len(window), 10))
    if len(window) == 1:
        axes = [axes]

    for ax, pid in zip(axes, window):
        img_path = os.path.join(IMG_DIR, f"{pid}.jpg")
        try:
            ax.imshow(Image.open(img_path).convert('RGB'))
        except FileNotFoundError:
            ax.text(0.5, 0.5, "Non trovata", ha='center', va='center', transform=ax.transAxes)
        ax.set_title(f"{pid}" + (" ← RISPOSTA" if pid == sample['image'] else ""), fontsize=9)
        ax.axis('off')

    plt.suptitle(
        f"TIPO:      {sample.get('corruption_type','').upper()}\n"
        f"MODIFICA:  {sample.get('corruption_detail','')}\n"
        f"ORIGINALE: {sample.get('original_question','')}\n"
        f"CORROTTA:  {sample['question']}",
        fontsize=10, ha='left', x=0.01
    )
    plt.tight_layout()
    plt.show()

# Esempio su un campione per tipo
for ctype in ["entity", "element", "layout"]:
    sample = next(s for s in corrupted_keep if s["corruption_type"] == ctype)
    show_window(sample)

# ── Caricamento benchmark da Drive ──
# Esegui questa cella dopo ogni restart del runtime prima di procedere con la Parte 2.

with open(BENCH_PATH) as f:
    benchmark_final = json.load(f)

# Ricostruisce le variabili necessarie per le celle successive
corrupted_keep = [s for s in benchmark_final if s["label"] == "unanswerable"]
clean_data     = [s for s in benchmark_final if s["label"] == "answerable"]

print(f"Benchmark caricato: {len(benchmark_final)} campioni")
print(f"  Unanswerable : {len(corrupted_keep)}")
print(f"  Answerable   : {len(clean_data)}")
print()

"""## PARTE 2 : Benchmark

## INFERENZA BASELINE

Come nella versione single-page, ma qui il modello riceve una FINESTRA di
pagine (definita da WINDOW_SIZE) invece di una sola immagine: tutte le
pagine della finestra vengono passate insieme al prompt testuale.
"""

# ══ SETUP GLOBALE MP-DocVQA ══ (rieseguire dopo OGNI restart)
from google.colab import drive
drive.mount('/content/drive')

import json, os, re, random, shutil, gc
from collections import Counter, defaultdict
import torch
from PIL import Image

SEED = 42
random.seed(SEED)

# ── Percorsi ──
DRIVE_DIR  = '/content/drive/MyDrive/MP-docvqa'
QAS_DIR    = f'{DRIVE_DIR}/qas'
IMG_SRC    = f'{DRIVE_DIR}/qas/images'       # immagini su Drive
IMG_DIR    = '/content/mpdocvqa_images'      # destinazione locale (RAM/SSD Colab)
BENCH_PATH = f'{DRIVE_DIR}/benchmark_final.json'
PRED_DIR   = f'{DRIVE_DIR}/predictions'
os.makedirs(PRED_DIR, exist_ok=True)
os.makedirs(IMG_DIR,  exist_ok=True)

# ── Parametri benchmark ──
N_ENTITY    = 34
N_ELEMENT   = 33
N_LAYOUT    = 33
N_CLEAN     = 100
WINDOW_SIZE = 2

# ── Carica benchmark ──
with open(BENCH_PATH) as f:
    benchmark_final = json.load(f)
print(f"Benchmark: {len(benchmark_final)} campioni")
print(f"Window size: {WINDOW_SIZE}")

# ── Copia immagini necessarie da Drive a /content/ ──
# Drive è lento per lettura random durante l'inferenza;
# copiare tutto in locale all'inizio è molto più veloce.
needed = set()
for s in benchmark_final:
    page_ids   = s.get("page_ids", [s["image"]])
    answer_idx = s.get("answer_page_idx", 0)
    start = max(0, answer_idx - WINDOW_SIZE // 2)
    end   = min(len(page_ids), start + WINDOW_SIZE)
    for pid in page_ids[start:end]:
        needed.add(f"{pid}.jpg")

missing = {f for f in needed if not os.path.isfile(os.path.join(IMG_DIR, f))}
print(f"Immagini necessarie : {len(needed)}")
print(f"  già in /content/  : {len(needed) - len(missing)}")
print(f"  da copiare        : {len(missing)}")

for fname in missing:
    src = os.path.join(IMG_SRC, fname)
    dst = os.path.join(IMG_DIR, fname)
    if os.path.isfile(src):
        shutil.copy2(src, dst)
    else:
        print(f"  [!] Non trovata su Drive: {fname}")

print("Immagini pronte in /content/")

import torch
from PIL import Image

with open(BENCH_PATH) as f:
    benchmark_final = json.load(f)
print(f"Benchmark caricato da Drive: {len(benchmark_final)} campioni")


# ── Funzione finestra MP-DocVQA ──
def get_window_pages(sample, window_size=None):
    global WINDOW_SIZE
    if window_size is None:
        window_size = WINDOW_SIZE
    page_ids   = sample["page_ids"]
    answer_idx = sample["answer_page_idx"]
    start = max(0, answer_idx - window_size // 2)
    end   = min(len(page_ids), start + window_size)
    return page_ids[start:end]

def resolve_image_path(page_id):
    filename  = f"{page_id}.jpg"   # ← .jpg non .png
    full_path = os.path.join(IMG_DIR, filename)
    if not os.path.isfile(full_path):
        raise FileNotFoundError(f"Immagine non trovata: {full_path}")
    return full_path

BASE_INSTRUCTION = (
    "You are a binary classifier. "
    "Look at the document pages and decide whether "
    "the question can be answered using ONLY the information visible in the provided pages."
)

FEW_SHOT_EXAMPLES = """
Example 1:
Question: "What is the value shown in the bottom-left table?"
Reasoning: The provided pages contain a table at the top-right, not bottom-left.
The referenced element does not exist in this position.
Answer: UNANSWERABLE

Example 2:
Question: "What is the total amount due?"
Reasoning: The provided pages contain a row labeled 'Total Amount Due' with a value.
The information requested is present in the document.
Answer: ANSWERABLE
"""

KNOWLEDGE_NOTE = """Before answering, verify the following in the provided document pages:
- Entities: check that any name, date, number, or organization mentioned
  in the question actually appears in the pages.
- Elements: check that any referenced element (table, figure, chart, footnote)
  exists in the pages.
- Layout: check that any referenced position (top, bottom, left, right)
  matches the actual position of elements in the pages.

Only after this verification, answer with ANSWERABLE or UNANSWERABLE."""

COT_INSTRUCTION = """Think step by step:
1) What does the question ask for?
2) Is it visible in any of the provided document pages?
   - YES → ANSWERABLE
   - NO  → UNANSWERABLE

Your last line must be exactly: ANSWERABLE or UNANSWERABLE."""

SYSTEM_MESSAGE = (
    "You are a binary classifier. "
    "You must reply with ONLY one word: ANSWERABLE or UNANSWERABLE. "
    "Never extract text from the document. Never answer the question itself."
)


def build_prompt(question, strategy="baseline"):
    if strategy == "baseline":
        return f"{BASE_INSTRUCTION}\n\nQuestion: {question}"
    if strategy == "few_shot":
        return f"{BASE_INSTRUCTION}\n\n{FEW_SHOT_EXAMPLES}\n\nQuestion: {question}"
    if strategy == "cot":
        return f"Question: {question}\n\n{COT_INSTRUCTION}"
    if strategy == "knowledge":
        return f"{BASE_INSTRUCTION}\n\n{KNOWLEDGE_NOTE}\n\nQuestion: {question}"
    raise ValueError(f"Strategia sconosciuta: '{strategy}'")


def normalize_pred(raw):
    text = (raw or "").lower()
    pos_unanswerable = text.rfind("unanswerable")
    pos_answerable   = text.rfind("answerable")

    if pos_unanswerable == -1 and pos_answerable == -1:
        return "answerable", True
    if pos_unanswerable != -1 and pos_answerable == pos_unanswerable + 2:
        return "unanswerable", False
    if pos_unanswerable > pos_answerable:
        return "unanswerable", False
    else:
        return "answerable", False


def run_inference(infer_fn, model_name, strategy="baseline"):
    global WINDOW_SIZE
    max_new_tokens = 64 if strategy == "cot" else 10
    results    = []
    fail_count = 0

    for i, sample in enumerate(benchmark_final):
        prompt          = build_prompt(sample["question"], strategy)
        # ── Recupera le pagine della finestra ──
        window_page_ids = get_window_pages(sample)
        image_paths     = []
        for pid in window_page_ids:
            try:
                image_paths.append(resolve_image_path(pid))
            except FileNotFoundError as e:
                print(f"[{model_name}] Pagina mancante: {e}", flush=True)

        if not image_paths:
            print(f"[{model_name}] Nessuna immagine per campione {i}", flush=True)
            raw_output = ""
        else:
            try:
                # infer_fn riceve lista di path invece di path singolo
                raw_output = infer_fn(image_paths, prompt, max_new_tokens)
            except Exception as e:
                print(f"[{model_name}] ERRORE campione {i}: {e}", flush=True)
                raw_output = ""

        prediction, parse_fail = normalize_pred(raw_output)
        fail_count += parse_fail

        results.append({
            "id"             : sample["questionId"],
            "label"          : sample["label"],
            "prediction"     : prediction,
            "corruption_type": sample.get("corruption_type", "clean"),
            "question"       : sample["question"],
            "window_pages"   : window_page_ids,   # pagine usate per questa inferenza
            "raw"            : raw_output.strip(),
            "parse_fail"     : parse_fail
        })

        if (i + 1) % 10 == 0:
            print(f"[{model_name} | {strategy} | W={WINDOW_SIZE}] {i+1}/{len(benchmark_final)}", flush=True)

    filename = f"predictions_{model_name}_w{WINDOW_SIZE}"
    if strategy != "baseline":
        filename += f"_{strategy}"

    output_path = os.path.join(PRED_DIR, f"{filename}.json")
    with open(output_path, "w") as f:
        json.dump(results, f, indent=2)

    print(f"\n[{model_name} | {strategy}] completato")
    print(f"  window_size : {WINDOW_SIZE}")
    print(f"  parse_fail  : {fail_count}/{len(results)}")
    print(f"  salvato in  : {output_path}")

    return results


def free_vram(*names):
    g = globals()
    for n in names:
        if n in g: del g[n]
    gc.collect(); torch.cuda.empty_cache()
    print(f"VRAM allocata: {torch.cuda.memory_allocated()/1e9:.2f} GB")

"""## GEMMA

Carico Gemma-3-4B-it e giro l'inferenza baseline con la finestra di pagine
di default (W=2).
"""

pip install "transformers>=4.50" -q

# ── Gemma-3-4B-it ──
from huggingface_hub import login
from google.colab import userdata
from transformers import pipeline

login(token=userdata.get('HF_TOKEN'))

pipe_gemma = pipeline(
    "image-text-to-text",
    model="google/gemma-3-4b-it",
    dtype=torch.bfloat16,
    device_map="cuda"            # ← forza GPU, non auto
)

def infer_gemma(image_paths, prompt, max_new_tokens):
    # Carica tutte le pagine della finestra
    images = []
    for path in image_paths:
        img = Image.open(path).convert("RGB")
        img.thumbnail((512, 512))
        images.append(img)

    # Costruisce il contenuto: una image per ogni pagina + il testo
    content = [{"type": "system", "content": [{"type": "text", "text": SYSTEM_MESSAGE}]}]
    user_content = []
    for img in images:
        user_content.append({"type": "image", "image": img})
    user_content.append({"type": "text", "text": prompt})

    messages = [
        {"role": "system", "content": [{"type": "text", "text": SYSTEM_MESSAGE}]},
        {"role": "user",   "content": user_content}
    ]

    out = pipe_gemma(messages, max_new_tokens=max_new_tokens)
    return out[0]["generated_text"][-1]["content"]


# Baseline su entrambe le finestre (stessa logica di Qwen e Phi)
WINDOW_SIZE = 1
predictions_gemma_w1 = run_inference(infer_gemma, "gemma3_4b")

WINDOW_SIZE = 2
predictions_gemma_w2 = run_inference(infer_gemma, "gemma3_4b")

free_vram("pipe_gemma")

"""## QWEN

Eseguo la baseline sia con finestra W=1 (equivalente a single-page) sia
W=2, per avere già un primo confronto sull'effetto della window size.
"""

pip install "transformers>=4.52" qwen-vl-utils -q

# ── Sessione B: Qwen2-VL-2B-Instruct ──
# !pip install "transformers>=4.52" qwen-vl-utils -q
# → Runtime > Restart session → Setup globale → Utilities → questa cella

from transformers import Qwen2VLForConditionalGeneration, AutoProcessor

# Caricamento modello
qwen_model = Qwen2VLForConditionalGeneration.from_pretrained(
    "Qwen/Qwen2-VL-2B-Instruct",
    torch_dtype=torch.float16,
    device_map="auto"
)
qwen_proc = AutoProcessor.from_pretrained("Qwen/Qwen2-VL-2B-Instruct")




def infer_qwen(image_paths, prompt, max_new_tokens):
    # Carica tutte le pagine della finestra
    images = []
    for path in image_paths:
        img = Image.open(path).convert("RGB")
        img.thumbnail((512, 512))
        images.append(img)

    # Costruisce il contenuto: una image per ogni pagina + il testo
    content = []
    for img in images:
        content.append({"type": "image", "image": img})
    content.append({"type": "text", "text": prompt})

    messages = [
        {"role": "system", "content": SYSTEM_MESSAGE},
        {"role": "user",   "content": content}
    ]

    text   = qwen_proc.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
    inputs = qwen_proc(text=[text], images=images, return_tensors="pt").to(qwen_model.device)

    with torch.no_grad():
        output_ids = qwen_model.generate(**inputs, max_new_tokens=max_new_tokens)

    new_tokens = output_ids[0][inputs["input_ids"].shape[-1]:]
    return qwen_proc.decode(new_tokens, skip_special_tokens=True)




# Inferenza su tutti i 200 campioni → salva predictions_qwen2vl_2b.json su Drive


# Run 1 — window 1 (single-page equivalent)
WINDOW_SIZE = 1
predictions_qwen_w1 = run_inference(infer_qwen, "qwen2vl_2b")
# salva predictions_qwen2vl_2b_w1.json

# Run 2 — window 2
WINDOW_SIZE = 2
predictions_qwen_w2 = run_inference(infer_qwen, "qwen2vl_2b")
# salva predictions_qwen2vl_2b_w2.json
# Libera la VRAM prima della sessione successiva
free_vram("qwen_model", "qwen_proc")

"""## PHI 3.5

Sessione separata (transformers==4.48.0). Eseguo la baseline sia con finestra W=1 (equivalente a single-page) sia
W=2, per avere già un primo confronto sull'effetto della window size.
"""

pip install "transformers==4.48.0" -q

# ── Phi-3.5-vision-instruct ──
# SESSIONE SEPARATA: pip install "transformers==4.48.0" -q → restart → Setup → Utilities

from transformers import AutoModelForCausalLM, AutoProcessor

phi_model = AutoModelForCausalLM.from_pretrained(
    "microsoft/Phi-3.5-vision-instruct",
    torch_dtype=torch.float16,
    device_map="auto",
    trust_remote_code=True,
    _attn_implementation="eager"
)
phi_proc = AutoProcessor.from_pretrained(
    "microsoft/Phi-3.5-vision-instruct",
    trust_remote_code=True
)

def infer_phi(image_paths, prompt, max_new_tokens):
    images = []
    for path in image_paths:
        img = Image.open(path).convert("RGB")
        img.thumbnail((512, 512))
        images.append(img)

    placeholders = "\n".join(f"<|image_{i+1}|>" for i in range(len(images)))

    messages = [
        {"role": "system", "content": SYSTEM_MESSAGE},
        {"role": "user",   "content": f"{placeholders}\n{prompt}"}
    ]

    text = phi_proc.tokenizer.apply_chat_template(
        messages, tokenize=False, add_generation_prompt=True
    )
    inputs = phi_proc(text, images, return_tensors="pt").to(phi_model.device)

    with torch.no_grad():
        output_ids = phi_model.generate(**inputs, max_new_tokens=max_new_tokens)

    new_tokens = output_ids[0][inputs["input_ids"].shape[-1]:]
    return phi_proc.decode(new_tokens, skip_special_tokens=True)


# Baseline su entrambe le finestre (stessa logica della cella Qwen)
WINDOW_SIZE = 1
predictions_phi_w1 = run_inference(infer_phi, "phi35_vision")

WINDOW_SIZE = 2
predictions_phi_w2 = run_inference(infer_phi, "phi35_vision")

free_vram("phi_model", "phi_proc")

"""## VALUTAZIONE E ANALISI (BASELINE)

Calcolo le metriche per ciascun modello/finestra e visualizzo le confusion
matrix, per un primo confronto baseline tra i tre VLM.
"""

import numpy as np
import matplotlib.pyplot as plt
from sklearn.metrics import (
    accuracy_score, precision_score, recall_score,
    f1_score, confusion_matrix
)

# Classe positiva = unanswerable (quella che vogliamo rilevare)
POS = "unanswerable"
NEG = "answerable"


def evaluate(predictions, model_name):
    """
    Calcola le metriche di classificazione per un modello.
    Restituisce un dizionario con accuracy, precision, recall, F1
    globali e F1 per tipo di corruzione.
    """
    y_true = [p["label"]      for p in predictions]
    y_pred = [p["prediction"] for p in predictions]

    # ── Metriche per tipo di corruzione ──
    # Per le corrotte (entity/element/layout): F1 sulla classe unanswerable
    # Per le clean: accuracy (non ha senso F1 su una classe assente)
    per_type = {}
    corruption_types = sorted({p["corruption_type"] for p in predictions})

    for ctype in corruption_types:
        subset  = [p for p in predictions if p["corruption_type"] == ctype]
        yt      = [p["label"]      for p in subset]
        yp      = [p["prediction"] for p in subset]

        if ctype == "clean":
            per_type[ctype] = accuracy_score(yt, yp)
        else:
            per_type[ctype] = f1_score(
                yt, yp,
                pos_label=POS,
                labels=[NEG, POS],
                zero_division=0
            )

    return {
        "model"      : model_name,
        "acc"        : accuracy_score(y_true, y_pred),
        "prec"       : precision_score(y_true, y_pred, pos_label=POS, zero_division=0),
        "rec"        : recall_score(y_true, y_pred,    pos_label=POS, zero_division=0),
        "f1"         : f1_score(y_true, y_pred,        pos_label=POS, zero_division=0),
        "f1_macro"   : f1_score(y_true, y_pred, average="macro",      zero_division=0),
        "cm"         : confusion_matrix(y_true, y_pred, labels=[NEG, POS]),
        "parse_fail" : sum(p.get("parse_fail", False) for p in predictions),
        "per_type"   : per_type,
    }


# ── Carica tutti i JSON di predizione e calcola le metriche ──
all_metrics = {}

for fname in sorted(os.listdir(PRED_DIR)):
    if not (fname.startswith("predictions_") and fname.endswith(".json")):
        continue

    model_name = fname[len("predictions_"):-len(".json")]

    with open(os.path.join(PRED_DIR, fname)) as f:
        predictions = json.load(f)

    all_metrics[model_name] = evaluate(predictions, model_name)


# ── Tabella 1: metriche globali ──
print(f"\n{'Modello':<26} {'Acc':>6} {'Prec':>6} {'Rec':>6} {'F1':>6} {'F1-mac':>7} {'fail':>5}")
print("─" * 68)
for name, m in all_metrics.items():
    print(
        f"{name:<26} "
        f"{m['acc']:>6.3f} "
        f"{m['prec']:>6.3f} "
        f"{m['rec']:>6.3f} "
        f"{m['f1']:>6.3f} "
        f"{m['f1_macro']:>7.3f} "
        f"{m['parse_fail']:>5d}"
    )

# ── Tabella 2: F1 per tipo di corruzione ──
CTYPES = ["entity", "element", "layout", "clean"]

print(f"\n{'F1 per tipo (clean = accuracy)'}")
print(f"{'Modello':<26}" + "".join(f"{c:>10}" for c in CTYPES))
print("─" * 66)
for name, m in all_metrics.items():
    row = "".join(
        f"{m['per_type'].get(c, float('nan')):>10.3f}"
        for c in CTYPES
    )
    print(f"{name:<26}{row}")

# ── Confusion matrix per tutti i modelli ──
models = list(all_metrics.keys())
n = len(models)

fig, axes = plt.subplots(1, n, figsize=(3.5 * n, 3.5))
if n == 1:
    axes = [axes]

for ax, name in zip(axes, models):
    m  = all_metrics[name]
    cm = m["cm"]

    im = ax.imshow(cm, cmap="Blues")

    # valori nelle celle
    for i in range(2):
        for j in range(2):
            ax.text(j, i, cm[i, j],
                    ha="center", va="center", fontsize=12,
                    color="white" if cm[i, j] > cm.max() / 2 else "black")

    ax.set_xticks([0, 1])
    ax.set_xticklabels([NEG, POS], rotation=45, ha="right", fontsize=8)
    ax.set_yticks([0, 1])
    ax.set_yticklabels([NEG, POS], fontsize=8)
    ax.set_xlabel("Predicted", fontsize=8)
    ax.set_ylabel("True", fontsize=8)
    ax.set_title(f"{name}\nAcc={m['acc']:.2f}  F1={m['f1']:.2f}  Rec={m['rec']:.2f}", fontsize=8)

plt.suptitle("Confusion matrices — tutti i modelli", fontsize=11, y=1.02)
plt.tight_layout()
plt.show()

"""## PARTE 3  MITIGAZIONE IN-CONTEXT

## QWEN

Applico le tre strategie di mitigazione (few-shot, CoT, knowledge) a Qwen,
su entrambe le finestre W=1 e W=2.
"""

pip install "transformers>=4.52" qwen-vl-utils -q

# ══════════════════════════════════════════════════════════════
# PARTE 3 — Mitigazione Qwen2-VL-2B (W=1 e W=2)
# ══════════════════════════════════════════════════════════════

from transformers import Qwen2VLForConditionalGeneration, AutoProcessor

qwen_model = Qwen2VLForConditionalGeneration.from_pretrained(
    "Qwen/Qwen2-VL-2B-Instruct",
    torch_dtype=torch.float16,
    device_map="auto"
)
qwen_proc = AutoProcessor.from_pretrained("Qwen/Qwen2-VL-2B-Instruct")

def infer_qwen(image_paths, prompt, max_new_tokens):
    images = []
    for path in image_paths:
        img = Image.open(path).convert("RGB")
        img.thumbnail((512, 512))
        images.append(img)

    content = []
    for img in images:
        content.append({"type": "image", "image": img})
    content.append({"type": "text", "text": prompt})

    messages = [
        {"role": "system", "content": SYSTEM_MESSAGE},
        {"role": "user",   "content": content}
    ]

    text = qwen_proc.apply_chat_template(
        messages, tokenize=False, add_generation_prompt=True
    )
    inputs = qwen_proc(
        text=[text], images=images, return_tensors="pt"
    ).to(qwen_model.device)

    with torch.no_grad():
        output_ids = qwen_model.generate(**inputs, max_new_tokens=max_new_tokens)

    new_tokens = output_ids[0][inputs["input_ids"].shape[-1]:]
    return qwen_proc.decode(new_tokens, skip_special_tokens=True)


# ── Inferenza W=1 e W=2 per ogni strategia ──
global WINDOW_SIZE

for ws in [1, 2]:
    WINDOW_SIZE = ws
    for strategy in ["few_shot", "cot", "knowledge"]:
        print(f"\nW={ws} | {strategy.upper()}")
        run_inference(infer_qwen, "qwen2vl_2b", strategy=strategy)

free_vram("qwen_model", "qwen_proc")

!pip install "transformers==4.48.0" -q

"""## PHI MITIGAZIONE

Applico le tre strategie di mitigazione (few-shot, CoT, knowledge) a
Phi-3.5-vision, su entrambe le finestre W=1 e W=2.  
"""

# ══════════════════════════════════════════════════════════════
# PARTE 3 — Mitigazione Phi3.5-Vision (W=1 e W=2)
# pip install "transformers==4.48.0" -q → restart → Setup → Utilities
# ══════════════════════════════════════════════════════════════

from transformers import AutoModelForCausalLM, AutoProcessor

phi_model = AutoModelForCausalLM.from_pretrained(
    "microsoft/Phi-3.5-vision-instruct",
    torch_dtype=torch.float16,
    device_map="auto",
    trust_remote_code=True,
    _attn_implementation="eager"
)
phi_proc = AutoProcessor.from_pretrained(
    "microsoft/Phi-3.5-vision-instruct",
    trust_remote_code=True
)

def infer_phi(image_paths, prompt, max_new_tokens):
    images = []
    for path in image_paths:
        img = Image.open(path).convert("RGB")
        img.thumbnail((512, 512))
        images.append(img)

    # Placeholder numerati per ogni pagina della finestra
    placeholders = "\n".join(f"<|image_{i+1}|>" for i in range(len(images)))

    messages = [
        {"role": "system", "content": SYSTEM_MESSAGE},
        {"role": "user",   "content": f"{placeholders}\n{prompt}"}
    ]

    text = phi_proc.tokenizer.apply_chat_template(
        messages, tokenize=False, add_generation_prompt=True
    )
    inputs = phi_proc(text, images, return_tensors="pt").to(phi_model.device)

    with torch.no_grad():
        output_ids = phi_model.generate(**inputs, max_new_tokens=max_new_tokens)

    new_tokens = output_ids[0][inputs["input_ids"].shape[-1]:]
    return phi_proc.decode(new_tokens, skip_special_tokens=True)


# ── Inferenza W=1 e W=2 per ogni strategia ──
for ws in [1, 2]:
    WINDOW_SIZE = ws
    for strategy in ["few_shot", "cot", "knowledge"]:
        print(f"\n{'═'*50}")
        print(f"  W={ws} | Strategia: {strategy.upper()}")
        print(f"{'═'*50}")
        run_inference(infer_phi, "phi35_vision", strategy=strategy)

free_vram("phi_model", "phi_proc")

"""## GEMMA MITIGAZIONE

Stessa procedura della Parte 3, applicata a Gemma-3-4B su entrambe le
finestre.
"""

# ══════════════════════════════════════════════════════════════
# PARTE 3 — Mitigazione Gemma3-4B (W=1 e W=2)
# ══════════════════════════════════════════════════════════════

from huggingface_hub import login
from google.colab import userdata
from transformers import pipeline

login(token=userdata.get('HF_TOKEN'))

pipe_gemma = pipeline(
    "image-text-to-text",
    model="google/gemma-3-4b-it",
    dtype=torch.bfloat16,
    device_map="cuda"
)

def infer_gemma(image_paths, prompt, max_new_tokens):
    images = []
    for path in image_paths:
        img = Image.open(path).convert("RGB")
        img.thumbnail((512, 512))
        images.append(img)

    user_content = []
    for img in images:
        user_content.append({"type": "image", "image": img})
    user_content.append({"type": "text", "text": prompt})

    messages = [
        {"role": "system", "content": [{"type": "text", "text": SYSTEM_MESSAGE}]},
        {"role": "user",   "content": user_content}
    ]

    out = pipe_gemma(messages, max_new_tokens=max_new_tokens)
    return out[0]["generated_text"][-1]["content"]


for ws in [1, 2]:
    WINDOW_SIZE = ws  # gira sia W=1 sia W=2

    for strategy in [ "few_shot", "cot", "knowledge"]:
        print(f"\nW={ws} | {strategy.upper()}")
        run_inference(infer_gemma, "gemma3_4b", strategy=strategy)

free_vram("pipe_gemma")

"""## CONFRONTO FINALE E GRAFICI

Carico tutte le predizioni (modello × strategia × finestra) e genero
tabelle e grafici di confronto finali: F1 per strategia e window size,
delta F1 tra W=1 e W=2, heatmap F1 per tipo di corruzione.
"""

# ── Disconnetti GPU quando non serve ──
from google.colab import runtime
runtime.unassign()

import numpy as np
import matplotlib.pyplot as plt
from sklearn.metrics import (
    accuracy_score, precision_score, recall_score,
    f1_score, confusion_matrix
)
import re

POS = "unanswerable"
NEG = "answerable"

STRATEGY_ORDER = {"baseline": 0, "few_shot": 1, "cot": 2, "knowledge": 3}
MODEL_ORDER    = {"gemma3_4b": 0, "qwen2vl_2b": 1, "phi35_vision": 2}

def sort_key(item):
    name = item[0] if isinstance(item, tuple) else item
    m    = item[1] if isinstance(item, tuple) else all_metrics[item]
    base = next((k for k in MODEL_ORDER if k in name), name)
    return (MODEL_ORDER.get(base, 99), m["window"] or 0, STRATEGY_ORDER.get(m["strategy"], 99))

def evaluate(predictions, model_name):
    y_true = [p["label"]      for p in predictions]
    y_pred = [p["prediction"] for p in predictions]

    per_type = {}
    corruption_types = sorted({p["corruption_type"] for p in predictions})
    for ctype in corruption_types:
        subset = [p for p in predictions if p["corruption_type"] == ctype]
        yt     = [p["label"]      for p in subset]
        yp     = [p["prediction"] for p in subset]
        if ctype == "clean":
            per_type[ctype] = accuracy_score(yt, yp)
        else:
            per_type[ctype] = f1_score(yt, yp, pos_label=POS, labels=[NEG, POS], zero_division=0)

    w_match  = re.search(r"_w(\d+)", model_name)
    window   = int(w_match.group(1)) if w_match else None
    strategy = "baseline"
    for s in ["few_shot", "cot", "knowledge"]:
        if model_name.endswith(s):
            strategy = s
            break

    return {
        "model"      : model_name,
        "window"     : window,
        "strategy"   : strategy,
        "acc"        : accuracy_score(y_true, y_pred),
        "prec"       : precision_score(y_true, y_pred, pos_label=POS, zero_division=0),
        "rec"        : recall_score(y_true, y_pred,    pos_label=POS, zero_division=0),
        "f1"         : f1_score(y_true, y_pred,        pos_label=POS, zero_division=0),
        "f1_macro"   : f1_score(y_true, y_pred, average="macro",      zero_division=0),
        "cm"         : confusion_matrix(y_true, y_pred, labels=[NEG, POS]),
        "parse_fail" : sum(p.get("parse_fail", False) for p in predictions),
        "per_type"   : per_type,
    }

# ── Carica tutti i JSON ──
all_metrics = {}
for fname in sorted(os.listdir(PRED_DIR)):
    if not (fname.startswith("predictions_") and fname.endswith(".json")):
        continue
    model_name = fname[len("predictions_"):-len(".json")]
    with open(os.path.join(PRED_DIR, fname)) as f:
        predictions = json.load(f)
    all_metrics[model_name] = evaluate(predictions, model_name)

# ── Tabella 1 ──
print(f"\n{'Modello':<36} {'W':>3} {'Strategy':<10} {'Acc':>6} {'Prec':>6} {'Rec':>6} {'F1':>6} {'F1-mac':>7} {'fail':>5}")
print("─" * 90)
for name, m in sorted(all_metrics.items(), key=sort_key):
    print(
        f"{name:<36} {str(m['window']):>3} {m['strategy']:<10} "
        f"{m['acc']:>6.3f} {m['prec']:>6.3f} {m['rec']:>6.3f} "
        f"{m['f1']:>6.3f} {m['f1_macro']:>7.3f} {m['parse_fail']:>5d}"
    )

# ── Tabella 2 ──
CTYPES = ["entity", "element", "layout", "clean"]
print(f"\n{'F1 per tipo (clean = accuracy)'}")
print(f"{'Modello':<36} {'W':>3}" + "".join(f"{c:>10}" for c in CTYPES))
print("─" * 80)
for name, m in sorted(all_metrics.items(), key=sort_key):
    row = "".join(f"{m['per_type'].get(c, float('nan')):>10.3f}" for c in CTYPES)
    print(f"{name:<36} {str(m['window']):>3}{row}")

# ── Tabella 3: delta W=1 vs W=2 ──
print(f"\n{'Confronto W=1 vs W=2 — delta F1 per modello'}")
print(f"{'Modello':<20} {'Strategia':<12} {'F1 W=1':>8} {'F1 W=2':>8} {'Δ F1':>8}")
print("─" * 60)
for model_base in ["gemma3_4b", "qwen2vl_2b", "phi35_vision"]:
    for strategy in ["baseline", "few_shot", "cot", "knowledge"]:
        w1 = next((m for m in all_metrics.values() if m["window"] == 1 and m["strategy"] == strategy and model_base in m["model"]), None)
        w2 = next((m for m in all_metrics.values() if m["window"] == 2 and m["strategy"] == strategy and model_base in m["model"]), None)
        if w1 and w2:
            delta = w2["f1"] - w1["f1"]
            print(f"{model_base:<20} {strategy:<12} {w1['f1']:>8.3f} {w2['f1']:>8.3f} {delta:>+8.3f}")

# ── Confusion matrices ──
models = list(all_metrics.keys())
n      = len(models)
n_cols = 4
n_rows = (n + n_cols - 1) // n_cols

fig, axes = plt.subplots(n_rows, n_cols, figsize=(4 * n_cols, 4 * n_rows))
axes = axes.flatten()

for idx, name in enumerate(sorted(models, key=lambda x: sort_key((x, all_metrics[x])))):
    ax = axes[idx]
    m  = all_metrics[name]
    cm = m["cm"]

    ax.imshow(cm, cmap="Blues")
    for i in range(2):
        for j in range(2):
            ax.text(j, i, cm[i, j], ha="center", va="center", fontsize=11,
                    color="white" if cm[i, j] > cm.max() / 2 else "black")

    ax.set_xticks([0, 1]); ax.set_xticklabels([NEG, POS], rotation=45, ha="right", fontsize=7)
    ax.set_yticks([0, 1]); ax.set_yticklabels([NEG, POS], fontsize=7)
    ax.set_xlabel("Predicted", fontsize=7)
    ax.set_ylabel("True", fontsize=7)
    ax.set_title(
        f"{name}\nW={m['window']} | {m['strategy']}\nAcc={m['acc']:.2f}  F1={m['f1']:.2f}",
        fontsize=7
    )

for idx in range(n, len(axes)):
    axes[idx].axis("off")

plt.suptitle("Confusion matrices", fontsize=13, y=1.01)
plt.tight_layout()
plt.show()

import numpy as np
import matplotlib.pyplot as plt
import matplotlib.patches as mpatches
import os, json, re
from sklearn.metrics import accuracy_score, precision_score, recall_score, f1_score

# ── CONFIG ──
PRED_DIR = "/content/drive/MyDrive/MP-docvqa/predictions"
POS = "unanswerable"
NEG = "answerable"

STRATEGY_ORDER = {"baseline": 0, "few_shot": 1, "cot": 2, "knowledge": 3}
MODEL_ORDER    = {"gemma3_4b": 0, "qwen2vl_2b": 1, "phi35_vision": 2}
MODEL_LABELS   = {"gemma3_4b": "Gemma-3 4B", "qwen2vl_2b": "Qwen2-VL 2B", "phi35_vision": "Phi-3.5 Vision"}
STRATEGY_LABELS = {"baseline": "Baseline", "few_shot": "Few-shot", "cot": "CoT", "knowledge": "Knowledge"}
CTYPES = ["entity", "element", "layout", "clean"]

# ── COLORI ──
MODEL_COLORS = {
    "gemma3_4b":    {"w1": "#4C72B0", "w2": "#1A3D6B"},
    "qwen2vl_2b":   {"w1": "#DD8452", "w2": "#8B4A1F"},
    "phi35_vision": {"w1": "#55A868", "w2": "#2A6B3A"},
}
CTYPE_COLORS = {"entity": "#4C72B0", "element": "#DD8452", "layout": "#55A868", "clean": "#C44E52"}

# ── EVALUATE ──
def evaluate(predictions, model_name):
    y_true = [p["label"]      for p in predictions]
    y_pred = [p["prediction"] for p in predictions]

    per_type = {}
    for ctype in CTYPES:
        subset = [p for p in predictions if p["corruption_type"] == ctype]
        if not subset:
            per_type[ctype] = float("nan")
            continue
        yt = [p["label"]      for p in subset]
        yp = [p["prediction"] for p in subset]
        if ctype == "clean":
            per_type[ctype] = accuracy_score(yt, yp)
        else:
            per_type[ctype] = f1_score(yt, yp, pos_label=POS, labels=[NEG, POS], zero_division=0)

    w_match  = re.search(r"_w(\d+)", model_name)
    window   = int(w_match.group(1)) if w_match else None
    strategy = "baseline"
    for s in ["few_shot", "cot", "knowledge"]:
        if model_name.endswith(s):
            strategy = s
            break

    base = next((k for k in MODEL_ORDER if k in model_name), model_name)

    return {
        "model"    : model_name,
        "base"     : base,
        "window"   : window,
        "strategy" : strategy,
        "f1"       : f1_score(y_true, y_pred, pos_label=POS, zero_division=0),
        "per_type" : per_type,
        "parse_fail": sum(p.get("parse_fail", False) for p in predictions),
        "n"        : len(predictions),
    }

# ── CARICA ──
all_metrics = {}
for fname in sorted(os.listdir(PRED_DIR)):
    if not (fname.startswith("predictions_") and fname.endswith(".json")):
        continue
    model_name = fname[len("predictions_"):-len(".json")]
    with open(os.path.join(PRED_DIR, fname)) as f:
        predictions = json.load(f)
    all_metrics[model_name] = evaluate(predictions, model_name)

# ════════════════════════════════════════════════════════════════
# GRAFICO 1 — F1 per strategia, W=1 vs W=2, per ogni modello
# ════════════════════════════════════════════════════════════════
fig, axes = plt.subplots(1, 3, figsize=(15, 5), sharey=True)
fig.suptitle("F1 (UNANSWERABLE) per strategia e window size", fontsize=14, fontweight="bold")

strategies = ["baseline", "few_shot", "cot", "knowledge"]
x = np.arange(len(strategies))
width = 0.35

for ax, model_base in zip(axes, ["gemma3_4b", "qwen2vl_2b", "phi35_vision"]):
    f1_w1, f1_w2 = [], []
    for s in strategies:
        m1 = next((m for m in all_metrics.values() if m["base"] == model_base and m["window"] == 1 and m["strategy"] == s), None)
        m2 = next((m for m in all_metrics.values() if m["base"] == model_base and m["window"] == 2 and m["strategy"] == s), None)
        f1_w1.append(m1["f1"] if m1 else 0)
        f1_w2.append(m2["f1"] if m2 else 0)

    bars1 = ax.bar(x - width/2, f1_w1, width, label="W=1", color=MODEL_COLORS[model_base]["w1"], alpha=0.85)
    bars2 = ax.bar(x + width/2, f1_w2, width, label="W=2", color=MODEL_COLORS[model_base]["w2"], alpha=0.85)

    ax.set_title(MODEL_LABELS[model_base], fontsize=12)
    ax.set_xticks(x)
    ax.set_xticklabels([STRATEGY_LABELS[s] for s in strategies], rotation=15, ha="right")
    ax.set_ylim(0, 1.0)
    ax.set_ylabel("F1" if ax == axes[0] else "")
    ax.legend(fontsize=9)
    ax.grid(axis="y", linestyle="--", alpha=0.4)
    ax.spines[["top", "right"]].set_visible(False)

    for bar in list(bars1) + list(bars2):
        h = bar.get_height()
        if h > 0.02:
            ax.text(bar.get_x() + bar.get_width()/2, h + 0.01, f"{h:.2f}", ha="center", va="bottom", fontsize=7.5)

plt.tight_layout()
plt.savefig("/content/plot1_f1_per_strategia.png", dpi=150, bbox_inches="tight")
plt.show()
print("Salvato: plot1_f1_per_strategia.png")


# ════════════════════════════════════════════════════════════════
# GRAFICO 2 — Delta F1 (W=2 - W=1) per modello e strategia
# ════════════════════════════════════════════════════════════════
fig, ax = plt.subplots(figsize=(10, 5))
fig.suptitle("Δ F1 (W=2 − W=1): impatto della window size", fontsize=14, fontweight="bold")

n_models = 3
n_strat  = 4
group_gap = 0.5
bar_width = 0.18

model_bases = ["gemma3_4b", "qwen2vl_2b", "phi35_vision"]
positions = []
deltas    = []
colors    = []

base_x = 0
for model_base in model_bases:
    for i, s in enumerate(strategies):
        m1 = next((m for m in all_metrics.values() if m["base"] == model_base and m["window"] == 1 and m["strategy"] == s), None)
        m2 = next((m for m in all_metrics.values() if m["base"] == model_base and m["window"] == 2 and m["strategy"] == s), None)
        delta = (m2["f1"] - m1["f1"]) if (m1 and m2) else 0
        positions.append(base_x + i * bar_width)
        deltas.append(delta)
        colors.append("#4C72B0" if delta >= 0 else "#C44E52")
    base_x += n_strat * bar_width + group_gap

bars = ax.bar(positions, deltas, width=bar_width * 0.85, color=colors, alpha=0.85)

# etichette gruppi
group_centers = [i * (n_strat * bar_width + group_gap) + (n_strat * bar_width) / 2 for i in range(n_models)]
ax.set_xticks(group_centers)
ax.set_xticklabels([MODEL_LABELS[m] for m in model_bases], fontsize=11)
ax.axhline(0, color="black", linewidth=0.8)
ax.set_ylabel("Δ F1")
ax.grid(axis="y", linestyle="--", alpha=0.4)
ax.spines[["top", "right"]].set_visible(False)

patch_pos = mpatches.Patch(color="#4C72B0", alpha=0.85, label="W=2 migliore")
patch_neg = mpatches.Patch(color="#C44E52", alpha=0.85, label="W=1 migliore")
ax.legend(handles=[patch_pos, patch_neg], fontsize=9)

for bar, delta in zip(bars, deltas):
    if abs(delta) > 0.005:
        y = bar.get_height()
        ax.text(bar.get_x() + bar.get_width()/2,
                y + 0.003 if y >= 0 else y - 0.015,
                f"{delta:+.2f}", ha="center", va="bottom" if y >= 0 else "top", fontsize=7)

# legenda strategie (mini)
for i, s in enumerate(strategies):
    ax.text(i * bar_width + bar_width/2, ax.get_ylim()[0] - 0.02,
            STRATEGY_LABELS[s][0], ha="center", fontsize=7, color="gray")

plt.tight_layout()
plt.savefig("/content/plot2_delta_window.png", dpi=150, bbox_inches="tight")
plt.show()
print("Salvato: plot2_delta_window.png")


# ════════════════════════════════════════════════════════════════
# GRAFICO 3 — F1 per tipo di corruzione (heatmap)
# ════════════════════════════════════════════════════════════════
run_labels = []
matrix     = []

sort_fn = lambda k: (MODEL_ORDER.get(all_metrics[k]["base"], 99),
                     all_metrics[k]["window"] or 0,
                     STRATEGY_ORDER.get(all_metrics[k]["strategy"], 99))

for name in sorted(all_metrics.keys(), key=sort_fn):
    m = all_metrics[name]
    row = [m["per_type"].get(c, float("nan")) for c in CTYPES]
    matrix.append(row)
    label = f"{MODEL_LABELS.get(m['base'], m['base'])} W={m['window']} {STRATEGY_LABELS.get(m['strategy'], m['strategy'])}"
    run_labels.append(label)

matrix = np.array(matrix)

fig, ax = plt.subplots(figsize=(8, 12))
fig.suptitle("F1 per tipo di corruzione (heatmap)", fontsize=14, fontweight="bold")

im = ax.imshow(matrix, aspect="auto", cmap="RdYlGn", vmin=0, vmax=1)
plt.colorbar(im, ax=ax, shrink=0.4, label="F1 / Accuracy")

ax.set_xticks(range(len(CTYPES)))
ax.set_xticklabels([c.capitalize() for c in CTYPES], fontsize=11)
ax.set_yticks(range(len(run_labels)))
ax.set_yticklabels(run_labels, fontsize=8)

for i in range(len(run_labels)):
    for j in range(len(CTYPES)):
        val = matrix[i, j]
        if not np.isnan(val):
            ax.text(j, i, f"{val:.2f}", ha="center", va="center",
                    fontsize=7.5, color="black" if 0.3 < val < 0.8 else "white")

plt.tight_layout()
plt.savefig("/content/plot3_heatmap_corruzione.png", dpi=150, bbox_inches="tight")
plt.show()
print("Salvato: plot3_heatmap_corruzione.png")


# ════════════════════════════════════════════════════════════════
# GRAFICO 4 — F1 medio per tipo corruzione, per modello (radar-style bar)
# ════════════════════════════════════════════════════════════════
fig, axes = plt.subplots(1, 3, figsize=(15, 5), sharey=True)
fig.suptitle("F1 medio per tipo di corruzione e strategia di mitigazione", fontsize=14, fontweight="bold")

for ax, model_base in zip(axes, ["gemma3_4b", "qwen2vl_2b", "phi35_vision"]):
    runs = [m for m in all_metrics.values() if m["base"] == model_base]
    x = np.arange(len(CTYPES))
    width_s = 0.18
    strat_colors = {"baseline": "#4C72B0", "few_shot": "#DD8452", "cot": "#55A868", "knowledge": "#C44E52"}

    for si, s in enumerate(strategies):
        vals = []
        for ctype in CTYPES:
            matches = [m["per_type"].get(ctype, float("nan")) for m in runs if m["strategy"] == s]
            mean_val = np.nanmean(matches) if matches else float("nan")
            vals.append(mean_val)
        offset = (si - 1.5) * width_s
        bars = ax.bar(x + offset, vals, width=width_s * 0.9,
                      label=STRATEGY_LABELS[s], color=strat_colors[s], alpha=0.85)

    ax.set_title(MODEL_LABELS[model_base], fontsize=12)
    ax.set_xticks(x)
    ax.set_xticklabels([c.capitalize() for c in CTYPES])
    ax.set_ylim(0, 1.05)
    ax.set_ylabel("F1 medio (W=1+W=2)" if ax == axes[0] else "")
    ax.legend(fontsize=8, loc="upper right")
    ax.grid(axis="y", linestyle="--", alpha=0.4)
    ax.spines[["top", "right"]].set_visible(False)

plt.tight_layout()
plt.savefig("/content/plot4_corruzione_per_strategia.png", dpi=150, bbox_inches="tight")
plt.show()
print("Salvato: plot4_corruzione_per_strategia.png")

print("\nTutti i grafici salvati in /content/")

for fname in sorted(os.listdir(PRED_DIR)):
    if fname.startswith("predictions_"):
        print(fname)

import json, os

with open(os.path.join(PRED_DIR, "predictions_qwen2vl_2b_w1.json")) as f:
    w1 = json.load(f)
with open(os.path.join(PRED_DIR, "predictions_qwen2vl_2b_w2.json")) as f:
    w2 = json.load(f)

print("W=1 — pagine per campione (primi 5):")
for p in w1[:5]:
    print(f"  {p['id']}: {p['window_pages']}")

print("\nW=2 — pagine per campione (primi 5):")
for p in w2[:5]:
    print(f"  {p['id']}: {p['window_pages']}")

"""## DIAGRAMMI RIASSUNTIVI (per la relazione)

Genero due diagrammi illustrativi:struttura dei due dataset e workflow
della pipeline, da usare nella relazione finale.
"""

import matplotlib.pyplot as plt
import matplotlib.patches as mpatches
from matplotlib.patches import FancyBboxPatch

fig, axes = plt.subplots(1, 2, figsize=(14, 9))
fig.patch.set_facecolor("white")

# ── Palette ──
C_HEADER  = "#1A5276"   # blu scuro
C_FIELD   = "#D6EAF8"   # blu chiaro
C_SPECIAL = "#1A7A4A"   # verde
C_CORRUPT = "#C0392B"   # rosso
C_WINDOW  = "#1A7A4A"   # verde
C_OUTPUT  = "#6C3483"   # viola
TEXT_DARK = "#1C1C1C"
TEXT_MID  = "#4A4A4A"
TEXT_LITE = "white"

def draw_box(ax, x, y, w, h, label, sublabel=None,
             fc="#D6EAF8", tc=TEXT_DARK, fs=10, sfs=8, radius=0.02):
    box = FancyBboxPatch((x, y), w, h,
                         boxstyle=f"round,pad=0",
                         facecolor=fc, edgecolor="#888888",
                         linewidth=0.6, zorder=2)
    ax.add_patch(box)
    if sublabel:
        ax.text(x + w/2, y + h*0.62, label,
                ha="center", va="center", fontsize=fs,
                fontweight="bold", color=tc, zorder=3)
        ax.text(x + w/2, y + h*0.28, sublabel,
                ha="center", va="center", fontsize=sfs,
                color=TEXT_MID, zorder=3)
    else:
        ax.text(x + w/2, y + h/2, label,
                ha="center", va="center", fontsize=fs,
                fontweight="bold", color=tc, zorder=3)

def draw_arrow(ax, x1, y1, x2, y2):
    ax.annotate("", xy=(x2, y2), xytext=(x1, y1),
                arrowprops=dict(arrowstyle="->", color="#555555",
                                lw=1.2, connectionstyle="arc3,rad=0.0"))

# ══════════════════════════════
# SP-DocVQA
# ══════════════════════════════
ax = axes[0]
ax.set_xlim(0, 1)
ax.set_ylim(0, 1)
ax.axis("off")
ax.set_title("SP-DocVQA  (single-page)", fontsize=13,
             fontweight="bold", color=C_HEADER, pad=10)

# Header
draw_box(ax, 0.05, 0.88, 0.90, 0.08,
         "Campione SP-DocVQA",
         "~8 GB · .png · train+val pool",
         fc=C_HEADER, tc=TEXT_LITE, fs=11, sfs=8)

# Campi base
fields_sp = [
    ("questionId",     "identificativo univoco"),
    ("question",       "testo della domanda"),
    ("answers",        "lista risposte gold"),
    ("image",          "percorso singola pagina  ← input visivo"),
    ("question_type",  "usato per campionamento stratificato"),
]
y0 = 0.84
for name, desc in fields_sp:
    y0 -= 0.085
    fc = "#D4EFDF" if name == "image" else \
         "#FAD7A0" if name == "question_type" else C_FIELD
    draw_box(ax, 0.05, y0, 0.90, 0.075, name, desc, fc=fc, fs=9, sfs=7.5)

# freccia verso corruzione
draw_arrow(ax, 0.50, y0, 0.50, y0 - 0.025)

# Box corruzione
y_c = y0 - 0.145
draw_box(ax, 0.05, y_c, 0.90, 0.13,
         "Corruzione applicata  (50 camp. × 3 tipi)",
         fc="#FADBD8", tc=TEXT_DARK, fs=9)

sub_w = 0.25
for i, (label, color) in enumerate([
        ("entity\n(spaCy NER)",    "#E74C3C"),
        ("element\n(regex)",       "#C0392B"),
        ("layout\n(regex)",        "#922B21"),
]):
    bx = 0.08 + i * (sub_w + 0.02)
    draw_box(ax, bx, y_c + 0.015, sub_w, 0.06,
             label, fc=color, tc=TEXT_LITE, fs=8)

# clean
ax.text(0.50, y_c - 0.018, "50 corrotti (UNANSWERABLE)  +  50 clean (ANSWERABLE)",
        ha="center", va="center", fontsize=7.5, color=TEXT_MID)

# freccia verso output
draw_arrow(ax, 0.50, y_c, 0.50, y_c - 0.035)

# Output
y_out = y_c - 0.105
draw_box(ax, 0.10, y_out, 0.80, 0.065,
         "VLM  →  ANSWERABLE / UNANSWERABLE",
         fc=C_OUTPUT, tc=TEXT_LITE, fs=9)

# ══════════════════════════════
# MP-DocVQA
# ══════════════════════════════
ax = axes[1]
ax.set_xlim(0, 1)
ax.set_ylim(0, 1)
ax.axis("off")
ax.set_title("MP-DocVQA  (multi-page)", fontsize=13,
             fontweight="bold", color="#0E6655", pad=10)

# Header
draw_box(ax, 0.05, 0.88, 0.90, 0.08,
         "Campione MP-DocVQA",
         "~22 GB · .jpg · train+val pool",
         fc="#0E6655", tc=TEXT_LITE, fs=11, sfs=8)

# Campi base
fields_mp = [
    ("questionId",        "identificativo univoco"),
    ("question",          "testo della domanda"),
    ("answers",           "lista risposte gold"),
    ("page_ids",          "lista ordinata di pagine del documento"),
    ("answer_page_idx",   "indice pagina contenente la risposta"),
]
y0 = 0.84
for name, desc in fields_mp:
    y0 -= 0.085
    fc = "#A9DFBF" if name in ("page_ids", "answer_page_idx") else C_FIELD
    draw_box(ax, 0.05, y0, 0.90, 0.075, name, desc, fc=fc, fs=9, sfs=7.5)

# freccia verso window
draw_arrow(ax, 0.50, y0, 0.50, y0 - 0.025)

# Box window
y_w = y0 - 0.155
draw_box(ax, 0.05, y_w, 0.90, 0.14,
         "Window W  centrata su answer_page_idx",
         fc="#D5F5E3", tc=TEXT_DARK, fs=9)

sub_w = 0.25
for i, (label, color) in enumerate([
        ("W = 1\n1 pagina",   "#1E8449"),
        ("W = 2\n2 pagine",   "#196F3D"),
        ("W = N\nN pagine",   "#145A32"),
]):
    bx = 0.08 + i * (sub_w + 0.02)
    draw_box(ax, bx, y_w + 0.015, sub_w, 0.07,
             label, fc=color, tc=TEXT_LITE, fs=8)

ax.text(0.50, y_w - 0.018,
        "98 corrotti (UNANSWERABLE)  +  98 clean (ANSWERABLE)  =  196 tot",
        ha="center", va="center", fontsize=7.5, color=TEXT_MID)

# freccia verso output
draw_arrow(ax, 0.50, y_w, 0.50, y_w - 0.035)

# Output
y_out = y_w - 0.105
draw_box(ax, 0.10, y_out, 0.80, 0.065,
         "VLM  (W immagini)  →  ANSWERABLE / UNANSWERABLE",
         fc=C_OUTPUT, tc=TEXT_LITE, fs=9)

plt.tight_layout(pad=2.0)
plt.savefig("/content/plot_dataset_structure.png", dpi=180, bbox_inches="tight",
            facecolor="white")
plt.show()
print("Salvato: /content/plot_dataset_structure.png")

import matplotlib.pyplot as plt
from matplotlib.patches import FancyBboxPatch

fig, ax = plt.subplots(figsize=(7, 10))
fig.patch.set_facecolor("white")
ax.set_xlim(0, 10)
ax.set_ylim(0, 14)
ax.axis("off")

C_GRAY     = "#4A4A4A"
C_PURPLE   = "#4A3A8A"
C_BLUE     = "#1A5276"
C_RED      = "#7B241C"
C_GREEN    = "#1E6B3C"
C_DARK_RED = "#6B1A1A"
C_GOLD     = "#7D6608"
WHITE      = "white"
GRAY_TEXT  = "#888888"

def box(ax, x, y, w, h, title, subtitle=None, fc="#4A4A4A", fs=11, sfs=8.5):
    ax.add_patch(FancyBboxPatch((x, y), w, h,
        boxstyle="round,pad=0.05", facecolor=fc,
        edgecolor="none", zorder=2))
    ty = y + h/2 + (0.12 if subtitle else 0)
    ax.text(x + w/2, ty, title, ha="center", va="center",
            fontsize=fs, fontweight="bold", color=WHITE, zorder=3)
    if subtitle:
        ax.text(x + w/2, y + h/2 - 0.22, subtitle,
                ha="center", va="center", fontsize=sfs,
                color=WHITE, alpha=0.85, zorder=3)

def arrow(ax, x, y1, y2):
    ax.annotate("", xy=(x, y2), xytext=(x, y1),
                arrowprops=dict(arrowstyle="-|>", color=GRAY_TEXT, lw=1.3))

# 1. Dataset pool
box(ax, 1.0, 12.6, 8, 0.9, "Dataset pool",
    "SP-DocVQA + MP-DocVQA  ·  train + val", fc=C_GRAY)
arrow(ax, 5, 12.6, 12.15)

# 2. Corruption engine, box esterno
ax.add_patch(FancyBboxPatch((1.0, 10.2), 8, 1.85,
    boxstyle="round,pad=0.05", facecolor=C_PURPLE,
    edgecolor="none", zorder=2))
ax.text(5, 11.85, "Corruption engine",
        ha="center", va="center", fontsize=11,
        fontweight="bold", color=WHITE, zorder=3)

# sub-boxes
sub_w, sub_h = 2.1, 1.1
starts = [1.2, 3.95, 6.7]
labels = [("Entity", "spaCy NER"), ("Element", "regex"), ("Layout", "opposti spaziali")]
for bx, (t, s) in zip(starts, labels):
    ax.add_patch(FancyBboxPatch((bx, 10.3), sub_w, sub_h,
        boxstyle="round,pad=0.05", facecolor=C_GREEN,
        edgecolor="none", zorder=3))
    ax.text(bx + sub_w/2, 10.3 + sub_h*0.65, t,
            ha="center", va="center", fontsize=9,
            fontweight="bold", color=WHITE, zorder=4)
    ax.text(bx + sub_w/2, 10.3 + sub_h*0.28, s,
            ha="center", va="center", fontsize=7.5,
            color=WHITE, alpha=0.85, zorder=4)

arrow(ax, 5, 10.2, 9.75)

# 3. Campionamento
box(ax, 1.0, 8.85, 8, 0.85, "Campionamento",
    "~34 entity  +  ~33 element  +  ~33 layout  +  clean", fc=C_BLUE)
arrow(ax, 5, 8.85, 8.4)

# 4. LLM-as-a-judge
box(ax, 1.0, 6.85, 8, 1.45, "LLM-as-a-judge",
    "Llama-3.3-70B via Groq  ·  well_formed? + needs_document?\nKEEP / REJECT  (max 1 riscrittura automatica)",
    fc=C_RED, sfs=8)

# branch
ax.plot([5, 3, 3], [6.85, 6.85, 6.35], color=GRAY_TEXT, lw=1.3)
ax.plot([5, 7, 7], [6.85, 6.85, 6.35], color=GRAY_TEXT, lw=1.3)
ax.annotate("", xy=(3, 6.35), xytext=(3, 6.4),
            arrowprops=dict(arrowstyle="-|>", color=GRAY_TEXT, lw=1.3))
ax.annotate("", xy=(7, 6.35), xytext=(7, 6.4),
            arrowprops=dict(arrowstyle="-|>", color=GRAY_TEXT, lw=1.3))

box(ax, 1.3, 5.7, 3.0, 0.6, "KEEP", fc=C_GREEN, fs=10)
box(ax, 5.7, 5.7, 3.0, 0.6, "REJECT", fc=C_DARK_RED, fs=10)

ax.plot([2.8, 2.8, 4.8], [5.7, 5.35, 5.35], color=GRAY_TEXT, lw=1.3)
ax.annotate("", xy=(4.8, 5.35), xytext=(4.7, 5.35),
            arrowprops=dict(arrowstyle="-|>", color=GRAY_TEXT, lw=1.3))
ax.text(7.2, 5.45, "→ rimosso un\ncampione clean",
        ha="center", va="center", fontsize=7.5,
        color=GRAY_TEXT, style="italic")

# 5. Bilanciamento
box(ax, 1.0, 4.3, 8, 0.95, "Bilanciamento 50/50",
    "corrotti (UNANSWERABLE)  +  clean (ANSWERABLE)", fc=C_GOLD)
arrow(ax, 5, 4.3, 3.85)

# 6. benchmark_final.json
box(ax, 1.0, 3.1, 8, 0.7, "benchmark_final.json", fc=C_GREEN, fs=10)
arrow(ax, 5, 3.1, 2.65)

# 7. Output
box(ax, 1.0, 1.9, 8, 0.7, "Parte 2 — Inferenza modelli", fc=C_GRAY, fs=10)

plt.tight_layout(pad=1.0)
plt.savefig("/content/plot_workflow.png", dpi=180,
            bbox_inches="tight", facecolor="white")
plt.show()
print("Salvato: /content/plot_workflow.png")