#!/usr/bin/env python
"""CSMorgan v10 — question-routed answer correction on top of v1 pipeline.

Pipeline:
1. Load same SimulaMet adapter as v1
2. Greedy inference, max_tokens=24 (deterministic, reproducible)
3. v1 normalization.py runs first (existing canonical mapping preserved)
4. v10 route_answer runs after — uses qbank for per-question canonical snapping
"""

# Defensive hf_transfer install
import subprocess, sys
try:
    import hf_transfer
except ImportError:
    subprocess.run([sys.executable, "-m", "pip", "install", "-q", "hf_transfer"], check=False)
    try:
        import hf_transfer
    except ImportError:
        import os as _os
        _os.environ.pop("HF_HUB_ENABLE_HF_TRANSFER", None)

import os, json, re, time, tempfile, importlib.util
from collections import Counter
from difflib import SequenceMatcher
from tqdm import tqdm

import torch
from datasets import load_dataset
from evaluate import load
from transformers import BitsAndBytesConfig

os.environ.setdefault("PYTORCH_CUDA_ALLOC_CONF", "expandable_segments:True")

bleu, rouge, meteor = load("bleu"), load("rouge"), load("meteor")
val_dataset = load_dataset("SimulaMet/Kvasir-VQA-test", split="validation")
predictions = []

gpu_name = torch.cuda.get_device_name(0) if torch.cuda.is_available() else "cpu"
device = "cuda" if torch.cuda.is_available() else "cpu"
def get_mem():
    return torch.cuda.memory_allocated(device) / (1024**2) if torch.cuda.is_available() else 0
initial_mem = get_mem()

HF_REPO_ID = "sageofai/Qwen25VL-MEDVQA-GI-S1-subtask1-v10"
V1_REPO = "sageofai/Qwen25VL-MEDVQA-GI-S1-subtask1"

SUBMISSION_INFO = {
    "Participant_Names": "Peter Ojonugwa Ejiga",
    "Affiliations": "Morgan State University",
    "Contact_emails": ["ojeji1@morgan.edu"],
    "Team_Name": "CSMorgan-MEDVQA",
    "Country": "USA",
    "Notes_to_organizers": (
        "v10: v1 pipeline + question-routed canonical correction. Greedy decoding "
        "(max=24, T=0). After v1's normalization, predictions are routed by detected "
        "question type (yes_no, count, absence, color, location, size, generic) and "
        "snapped to canonical training answers via qbank fuzzy match (threshold 0.85). "
        "Absence patterns ('no X identified', 'not visible') mapped to qbank's canonical "
        "absence answer ('none' or 'not relevant') for that specific question."
    ),
}

# ===========================================================================
# Fetch companion files at runtime
# ===========================================================================
_HERE = os.path.dirname(os.path.abspath(__file__))
if _HERE not in sys.path: sys.path.insert(0, _HERE)

def _try_fetch(repo_id, filename):
    p = os.path.join(_HERE, filename)
    if os.path.exists(p): return p
    try:
        from huggingface_hub import hf_hub_download
        return hf_hub_download(repo_id=repo_id, filename=filename, repo_type="model")
    except Exception as e:
        print(f"[fetch] {filename} from {repo_id} failed: {e}")
        return None

# v1's normalization layer (first-stage canonical mapping)
_norm_path = _try_fetch(V1_REPO, "normalization.py")
if _norm_path:
    spec = importlib.util.spec_from_file_location("normalization", _norm_path)
    _norm_mod = importlib.util.module_from_spec(spec)
    spec.loader.exec_module(_norm_mod)
    normalize_v1 = _norm_mod.normalize_task1_answer
    print("[init] v1 normalization.py loaded")
else:
    def normalize_v1(answer, question=None, answer_bank=None):
        return (answer or "").strip()
    print("[init] WARN: no v1 normalization — using identity")

# v1's answer_bank.json
ANSWER_BANK = {}
_ab_path = _try_fetch(V1_REPO, "answer_bank.json")
if _ab_path:
    try:
        with open(_ab_path) as f: ANSWER_BANK = json.load(f)
    except Exception as e:
        print(f"[init] answer_bank parse failed: {e}")

# qbank.json from v10 repo (per-question training answer distribution)
QBANK = {}
_qb_path = _try_fetch(HF_REPO_ID, "qbank.json")
if _qb_path:
    try:
        with open(_qb_path) as f: QBANK = json.load(f)
        print(f"[init] qbank loaded: {len(QBANK)} questions")
    except Exception as e:
        print(f"[init] qbank parse failed: {e}")

# ===========================================================================
# QBANK FUZZY LOOKUP
# ===========================================================================
def _norm_q(q):
    q = (q or "").lower().strip()
    q = re.sub(r"[^\w\s]", " ", q)
    q = re.sub(r"\s+", " ", q).strip()
    return q

QBANK_NORM_KEYS = {_norm_q(k): k for k in QBANK.keys()}

def _find_qbank_entry(question, threshold=0.7):
    qn = _norm_q(question)
    if not qn: return []
    if qn in QBANK_NORM_KEYS: return QBANK[QBANK_NORM_KEYS[qn]]
    for nk, orig in QBANK_NORM_KEYS.items():
        if qn in nk or nk in qn: return QBANK[orig]
    best_s, best_k = 0.0, None
    for nk, orig in QBANK_NORM_KEYS.items():
        s = SequenceMatcher(None, qn, nk).ratio()
        if s > best_s:
            best_s, best_k = s, orig
    return QBANK[best_k] if (best_k and best_s >= threshold) else []

# ===========================================================================
# QUESTION-TYPE DETECTION
# ===========================================================================
def detect_question_type(question):
    """Detect question type from text. Conservative heuristics."""
    q = (question or "").lower().strip()
    # Yes/no — questions starting with auxiliary verbs
    if re.match(r"^(is|are|does|do|has|have|was|were|did|can|will|should|would|could|may|might|must)\s+", q):
        return "yes_no"
    # Count
    if re.search(r"how many|number of|count of|count are", q):
        return "count"
    # Color
    if "color" in q or "colour" in q:
        return "color"
    # Location/anatomy
    if re.search(r"where in|where is|location|located|region|site|area", q):
        return "location"
    # Size
    if "size" in q or "measure" in q or "how big" in q or "how large" in q or "diameter" in q:
        return "size"
    # Type/finding
    if re.search(r"^what type|what kind|what is the type|what.*polyp.*present|what category", q):
        return "finding_type"
    # Generic "what" — finding/object
    if q.startswith(("what ", "which ")):
        return "finding"
    return "generic"

# ===========================================================================
# ANSWER ROUTING
# ===========================================================================
NUMBER_WORDS = {
    "zero": "0", "none": "0", "no": "0",
    "one": "1", "single": "1",
    "two": "2", "three": "3", "four": "4", "five": "5",
    "six": "6", "seven": "7", "eight": "8", "nine": "9", "ten": "10",
}

ABSENCE_PHRASES = [
    "no polyp", "no polyps", "no instrument", "no instruments",
    "no visible", "no abnormalit", "no finding", "no lesion",
    "not visible", "not present", "not visualized", "not identified",
    "none identified", "none observed", "none present", "none visible",
    "absent",
]

def _is_absence_phrase(text):
    """Detect 'the thing isn't there' patterns."""
    t = (text or "").lower()
    return any(phrase in t for phrase in ABSENCE_PHRASES)

def _fuzzy_match_in_qbank(pred, qbank_entries, threshold=0.85):
    """Try to snap pred to a qbank candidate. Returns canonical answer or None."""
    if not pred or not qbank_entries:
        return None
    pn = (pred or "").lower().strip()
    if not pn:
        return None
    # Exact match check
    for entry in qbank_entries[:50]:
        if entry["answer"].lower().strip() == pn:
            return entry["answer"]
    # Fuzzy match — weighted by frequency
    best_score, best_ans = 0.0, None
    max_count = max((e["count"] for e in qbank_entries[:30]), default=1)
    for entry in qbank_entries[:30]:
        cn = entry["answer"].lower().strip()
        s = SequenceMatcher(None, pn, cn).ratio()
        score = s + 0.05 * (entry["count"] / max_count)  # slight frequency bias
        if score > best_score:
            best_score, best_ans = score, entry["answer"]
    return best_ans if best_score >= threshold else None

def _absence_canonical(qbank_entries):
    """Find the canonical 'absence' answer for this question in qbank.
    Returns 'none' / 'not relevant' / similar — whichever the training set uses."""
    for entry in qbank_entries[:10]:
        a = entry["answer"].lower().strip()
        if a in ("none", "not relevant", "not visible", "not present",
                 "no instrument", "no polyp", "not applicable"):
            return entry["answer"]
    # Fallback: if no canonical absence in top-10, return 'none'
    return "none"

_ROUTE_STATS = Counter()

def route_answer(raw_pred, v1_normalized, question):
    """Apply question-type-specific routing on top of v1's normalized output.
    Returns the final canonical answer."""
    qtype = detect_question_type(question)
    qbank_entries = _find_qbank_entry(question)

    raw_l = (raw_pred or "").lower().strip()
    v1_l = (v1_normalized or "").lower().strip()
    candidate = v1_normalized or raw_pred or ""

    # ----- yes_no -----
    if qtype == "yes_no":
        # Look for definite signals
        if re.search(r"\byes\b", raw_l) or re.search(r"\byes\b", v1_l):
            _ROUTE_STATS["yes_no->yes"] += 1
            return "yes"
        if (re.search(r"\bno\b", raw_l) or re.search(r"\bno\b", v1_l)
                or _is_absence_phrase(raw_l)):
            _ROUTE_STATS["yes_no->no"] += 1
            return "no"
        # Uncertain — try qbank
        snap = _fuzzy_match_in_qbank(candidate, qbank_entries, threshold=0.7)
        if snap and snap.lower() in ("yes", "no"):
            _ROUTE_STATS[f"yes_no->{snap.lower()}_qbank"] += 1
            return snap
        # Fallback: pick whichever (yes/no) is more common in qbank
        for entry in qbank_entries[:5]:
            if entry["answer"].lower() in ("yes", "no"):
                _ROUTE_STATS[f"yes_no->fallback_{entry['answer']}"] += 1
                return entry["answer"]
        _ROUTE_STATS["yes_no->uncertain"] += 1
        return candidate

    # ----- count -----
    if qtype == "count":
        # Already a digit?
        m = re.match(r"^\s*(\d+)\s*$", v1_l)
        if m:
            _ROUTE_STATS["count->digit"] += 1
            return m.group(1)
        m = re.match(r"^\s*(\d+)\s*$", raw_l)
        if m:
            _ROUTE_STATS["count->digit_raw"] += 1
            return m.group(1)
        # Number word → digit
        first_word = (v1_l.split() or [""])[0]
        if first_word in NUMBER_WORDS:
            _ROUTE_STATS[f"count->word_{first_word}"] += 1
            return NUMBER_WORDS[first_word]
        # Absence → 0
        if _is_absence_phrase(v1_l) or _is_absence_phrase(raw_l):
            _ROUTE_STATS["count->absence_0"] += 1
            return "0"
        # Extract first digit in text
        m = re.search(r"\b(\d+)\b", v1_l + " " + raw_l)
        if m:
            _ROUTE_STATS["count->digit_in_text"] += 1
            return m.group(1)
        # Fall back to qbank fuzzy
        snap = _fuzzy_match_in_qbank(candidate, qbank_entries)
        if snap:
            _ROUTE_STATS["count->qbank"] += 1
            return snap
        _ROUTE_STATS["count->keep"] += 1
        return candidate

    # ----- color -----
    if qtype == "color":
        # Absence in color question is unusual but treat as canonical
        if _is_absence_phrase(raw_l):
            absence = _absence_canonical(qbank_entries)
            _ROUTE_STATS[f"color->absence_{absence}"] += 1
            return absence
        # Snap to qbank color
        snap = _fuzzy_match_in_qbank(candidate, qbank_entries, threshold=0.65)
        if snap:
            if snap.lower() != candidate.lower():
                _ROUTE_STATS["color->snapped"] += 1
            else:
                _ROUTE_STATS["color->kept"] += 1
            return snap
        _ROUTE_STATS["color->no_match"] += 1
        return candidate

    # ----- location -----
    if qtype == "location":
        if _is_absence_phrase(raw_l):
            absence = _absence_canonical(qbank_entries)
            _ROUTE_STATS[f"location->absence"] += 1
            return absence
        snap = _fuzzy_match_in_qbank(candidate, qbank_entries, threshold=0.75)
        if snap:
            _ROUTE_STATS["location->snapped" if snap.lower() != candidate.lower() else "location->kept"] += 1
            return snap
        _ROUTE_STATS["location->no_match"] += 1
        return candidate

    # ----- size -----
    if qtype == "size":
        # Common Kvasir size patterns
        for pattern in ["<5mm", "5-10mm", "10-20mm", "11-20mm", ">20mm"]:
            if pattern in raw_l or pattern in v1_l:
                _ROUTE_STATS[f"size->{pattern}"] += 1
                return pattern
        if _is_absence_phrase(raw_l):
            absence = _absence_canonical(qbank_entries)
            _ROUTE_STATS["size->absence"] += 1
            return absence
        snap = _fuzzy_match_in_qbank(candidate, qbank_entries, threshold=0.80)
        if snap:
            _ROUTE_STATS["size->snapped"] += 1
            return snap
        _ROUTE_STATS["size->keep"] += 1
        return candidate

    # ----- finding_type (What type of polyp / category) -----
    if qtype == "finding_type":
        if _is_absence_phrase(raw_l):
            absence = _absence_canonical(qbank_entries)
            _ROUTE_STATS["finding_type->absence"] += 1
            return absence
        snap = _fuzzy_match_in_qbank(candidate, qbank_entries, threshold=0.75)
        if snap:
            _ROUTE_STATS["finding_type->snapped" if snap.lower() != candidate.lower() else "finding_type->kept"] += 1
            return snap
        _ROUTE_STATS["finding_type->no_match"] += 1
        return candidate

    # ----- finding (generic What X) -----
    if qtype == "finding":
        if _is_absence_phrase(raw_l):
            absence = _absence_canonical(qbank_entries)
            _ROUTE_STATS["finding->absence"] += 1
            return absence
        snap = _fuzzy_match_in_qbank(candidate, qbank_entries, threshold=0.80)
        if snap:
            _ROUTE_STATS["finding->snapped" if snap.lower() != candidate.lower() else "finding->kept"] += 1
            return snap
        _ROUTE_STATS["finding->no_match"] += 1
        return candidate

    # ----- generic -----
    snap = _fuzzy_match_in_qbank(candidate, qbank_entries, threshold=0.85)
    if snap and snap.lower() != candidate.lower():
        _ROUTE_STATS["generic->snapped"] += 1
        return snap
    _ROUTE_STATS["generic->kept"] += 1
    return candidate

# ===========================================================================
# MODEL LOAD
# ===========================================================================
from swift.llm import PtEngine, RequestConfig, InferRequest

print("Loading model + adapter...")
engine = PtEngine(
    model_id_or_path="Qwen/Qwen2.5-VL-7B-Instruct",
    adapters=["SimulaMet/Qwen2.5-VL-KvasirVQA-x1-ft"],
    quantization_config=BitsAndBytesConfig(
        load_in_4bit=True, bnb_4bit_quant_type="nf4",
        bnb_4bit_use_double_quant=True, bnb_4bit_compute_dtype=torch.float16,
    ),
    attn_impl="sdpa", use_hf=True, max_length=2048,
)

# Greedy deterministic decoding (v1 style — proven 0.5423 baseline)
req_cfg = RequestConfig(
    max_tokens=24, temperature=0.0, top_k=1, top_p=1.0, repetition_penalty=1.0,
)
post_load_mem = get_mem()
print(f"Loaded. mem={post_load_mem:.0f} MB")

# ===========================================================================
# INFERENCE
# ===========================================================================
print("Inferring...")
_t0 = time.time()
_samples = []
_qtype_counts = Counter()

for idx, ex in enumerate(tqdm(val_dataset, desc="Validating")):
    question = ex["question"]
    image = ex["image"]
    if hasattr(image, "convert") and image.mode != "RGB":
        image = image.convert("RGB")
    tmp = tempfile.NamedTemporaryFile(suffix=".jpg", delete=False)
    tmp_path = tmp.name; tmp.close()
    image.save(tmp_path)

    try:
        req = InferRequest(messages=[{
            "role": "user",
            "content": [
                {"type": "image", "image": tmp_path},
                {"type": "text", "text": question},
            ],
        }])
        resp = engine.infer([req], req_cfg)
        raw = resp[0].choices[0].message.content or ""
        # Stage 1: v1 normalization
        v1_norm = normalize_v1(raw, question=question, answer_bank=ANSWER_BANK)
        # Stage 2: v10 question-routed correction
        final = route_answer(raw, v1_norm, question)
    except Exception as e:
        print(f"  idx={idx} error: {e}")
        raw, v1_norm, final = "", "", ""
    finally:
        try: os.unlink(tmp_path)
        except OSError: pass

    _qtype_counts[detect_question_type(question)] += 1

    predictions.append({"index": idx, "img_id": ex["img_id"],
                        "question": question, "answer": final})
    if len(_samples) < 30:
        _samples.append((question, raw, v1_norm, final, ex["answer"]))

# ===========================================================================
# DIAGNOSTIC
# ===========================================================================
print()
print("="*60)
print("QUESTION TYPE DISTRIBUTION")
print("="*60)
for qt, c in _qtype_counts.most_common():
    print(f"  {qt:15s} {c:5d}  ({100*c/len(val_dataset):.1f}%)")

print()
print("="*60)
print("ROUTING DECISIONS (top 25)")
print("="*60)
for decision, count in _ROUTE_STATS.most_common(25):
    print(f"  {decision:30s} {count:5d}")

answers = [p["answer"] for p in predictions]
wc = [len(a.split()) for a in answers if a]
freq = Counter(answers)
exact = sum(1 for p, e in zip(predictions, val_dataset)
            if str(p["answer"]).strip().lower() ==
               (e["answer"] if isinstance(e["answer"], str)
                else "; ".join(map(str, e["answer"]))).strip().lower())

print()
print("="*60)
print("FINAL DIAGNOSTIC")
print("="*60)
print(f"predictions    : {len(predictions)}")
print(f"empty answers  : {sum(1 for a in answers if not a.strip())}")
print(f"exact match    : {exact}/{len(predictions)} ({100*exact/len(predictions):.1f}%)")
print(f"avg length     : {sum(wc)/max(len(wc),1):.2f} words")
print()
print("Top-20 final predictions:")
for a, c in freq.most_common(20):
    print(f"  {c:5d}  {a!r}")
print()
print("Sample of 20 (Q | raw | v1_norm | final | ref):")
for q, raw, v1n, fin, ref in _samples[:20]:
    ref_str = ref if isinstance(ref, str) else "; ".join(map(str, ref))
    qt = detect_question_type(q)
    match = "✅" if fin.strip().lower() == ref_str.strip().lower() else "❌"
    print(f"  {match} [{qt:12s}] Q: {q[:35]!r:40s}")
    print(f"           raw={raw!r:30s} v1={v1n!r:25s} -> {fin!r:25s} | ref={ref_str!r}")
print("="*60)

# ===========================================================================
# SCORING
# ===========================================================================
preds_t = [p["answer"] for p in predictions]
refs_t = [ex["answer"] for ex in val_dataset]
def _flat(r): return "; ".join(str(x) for x in r) if isinstance(r, list) else str(r)
refs_f = [_flat(r) for r in refs_t]

bleu1 = bleu.compute(predictions=preds_t, references=[[r] for r in refs_f], max_order=1)
bleu4 = bleu.compute(predictions=preds_t, references=[[r] for r in refs_f])
rg = rouge.compute(predictions=preds_t, references=refs_f)
mt = meteor.compute(predictions=preds_t, references=refs_f)
scores = {
    "bleu":   round(bleu1["bleu"], 4),
    "bleu4":  round(bleu4["bleu"], 4),
    "rouge1": round(rg["rouge1"], 4),
    "rouge2": round(rg["rouge2"], 4),
    "rougeL": round(rg["rougeL"], 4),
    "meteor": round(mt["meteor"], 4),
}
print(f"\u2728Public scores: {scores}")

with open("predictions_1.json", "w") as f:
    json.dump(predictions, f, indent=2)

elapsed = time.time() - _t0
print(f"Time: {elapsed:.1f}s | Mem: {get_mem():.0f} MB | GPU: {gpu_name}")
print("Done.")
print(f"Run:: medvqa validate_and_submit --competition=gi-2026 --task=1 --repo_id={HF_REPO_ID}")