#!/usr/bin/env python
"""CSMorgan Task 1 v4 — SimulaMet adapter + fuzzy-question constrained decoding."""

import json, os, re, sys, tempfile, time, subprocess, platform
from collections import Counter
from difflib import SequenceMatcher

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

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

bleu   = load("bleu")
rouge  = load("rouge")
meteor = 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()

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": (
        "SimulaMet/Qwen2.5-VL-KvasirVQA-x1-ft adapter (no extra fine-tuning). "
        "Greedy decoding, max 20 tokens. Per-question constrained decoding: "
        "test question is fuzzy-matched against training questions (substring + "
        "SequenceMatcher) to find the right answer-set, then model output is "
        "fuzzy-matched against the top-30 training answers for that question."
    ),
}

HF_REPO_ID       = "sageofai/Qwen25VL-MEDVQA-GI-S1-subtask1-v4"
FALLBACK_ADAPTER = "SimulaMet/Qwen2.5-VL-KvasirVQA-x1-ft"

_HERE = os.path.dirname(os.path.abspath(__file__))
if _HERE not in sys.path:
    sys.path.insert(0, _HERE)

# -------- fetch qbank.json from the repo at runtime --------
QUESTION_BANK = {}
_bank_path = os.path.join(_HERE, "qbank.json")
if not os.path.exists(_bank_path):
    try:
        from huggingface_hub import hf_hub_download
        _bank_path = hf_hub_download(repo_id=HF_REPO_ID, filename="qbank.json", repo_type="model")
        print(f"[fetch] qbank.json from HF: {_bank_path}")
    except Exception as e:
        print(f"[fetch] qbank.json unavailable: {e}")
        _bank_path = None

if _bank_path and os.path.exists(_bank_path):
    with open(_bank_path) as f:
        QUESTION_BANK = json.load(f)
    print(f"Loaded question bank: {len(QUESTION_BANK)} training questions")
else:
    print("WARNING: no qbank.json — constraint disabled, raw model output only")

# -------- adapter selection (use SimulaMet's directly) --------
def _adapter_exists(repo_id):
    if not repo_id: return False
    try:
        from huggingface_hub import HfApi
        info = HfApi().repo_info(repo_id=repo_id, repo_type="model")
        files = {s.rfilename for s in info.siblings}
        return "adapter_config.json" in files and any(f.startswith("adapter_model.") for f in files)
    except Exception:
        return False

if _adapter_exists(HF_REPO_ID):
    ADAPTER = HF_REPO_ID
    print(f"[adapter] using user-fine-tuned: {ADAPTER}")
else:
    ADAPTER = FALLBACK_ADAPTER
    print(f"[adapter] no user adapter → falling back to: {ADAPTER}")

# -------- normalization helpers --------
def _normalize_question(q):
    q = (q or "").lower().strip()
    q = re.sub(r"[^\w\s]", " ", q)
    q = re.sub(r"\s+", " ", q).strip()
    return q

def _normalize_for_match(text):
    if not isinstance(text, str): text = str(text)
    t = text.lower().strip()
    if "\n" in t:
        t = t.split("\n", 1)[0].strip()
    for pfx in ("answer:", "final answer:", "the answer is", "it is", "this is"):
        if t.startswith(pfx):
            t = t[len(pfx):].strip()
    t = re.sub(r"[^\w\s\-;]", " ", t)
    t = re.sub(r"\b(a|an|the)\b", " ", t)
    t = re.sub(r"\s+", " ", t).strip()
    return t

# Pre-build normalized index of training questions for fast fuzzy lookup
QBANK_NORM_KEYS = {_normalize_question(k): k for k in QUESTION_BANK.keys()}
print(f"  pre-indexed {len(QBANK_NORM_KEYS)} normalized question keys")

def _find_question_in_bank(question, threshold=0.7):
    """Locate the training question whose answer pool we should constrain to."""
    qn = _normalize_question(question)
    if not qn:
        return None
    # 1. Direct normalized hit
    if qn in QBANK_NORM_KEYS:
        return QBANK_NORM_KEYS[qn]
    # 2. Substring containment (test is in training, or training is in test)
    for nk, orig in QBANK_NORM_KEYS.items():
        if qn in nk or nk in qn:
            return orig
    # 3. Fuzzy match
    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 best_k if best_s >= threshold else None

def _constrain(pred, question, threshold=0.5):
    if pred is None: pred = ""
    pnorm = _normalize_for_match(pred)

    matched_q = _find_question_in_bank(question)
    candidates = QUESTION_BANK.get(matched_q, []) if matched_q else []

    if not candidates:
        return pred  # no matching training question — give up

    if not pnorm:
        return candidates[0]["answer"]

    # Direct normalized hit in candidates
    for c in candidates[:30]:
        if _normalize_for_match(c["answer"]) == pnorm:
            return c["answer"]

    # Fuzzy match weighted by frequency
    best_s, best_a = 0.0, pred
    max_cnt = max((c["count"] for c in candidates[:30]), default=1)
    for c in candidates[:30]:
        cn = _normalize_for_match(c["answer"])
        if not cn: continue
        s = SequenceMatcher(None, pnorm, cn).ratio() + 0.05 * (c["count"]/max_cnt)
        if s > best_s:
            best_s, best_a = s, c["answer"]

    # If no candidate is even moderately close, default to most-common
    return best_a if best_s >= threshold else candidates[0]["answer"]

# -------- model load --------
from swift.llm import PtEngine, RequestConfig, InferRequest

print(f"Loading model + adapter...")
model_hf = PtEngine(
    model_id_or_path="Qwen/Qwen2.5-VL-7B-Instruct",
    adapters=[ADAPTER],
    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,
)
req_cfg = RequestConfig(max_tokens=20, temperature=0.0, top_k=1, top_p=1.0, repetition_penalty=1.0)
post_load_mem = get_mem()
print("Model loaded.")

# -------- inference loop --------
print("Starting inference...")
_t0 = time.time()
_samples = []

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 = model_hf.infer([req], req_cfg)
        raw = resp[0].choices[0].message.content or ""
        final = _constrain(raw, question, threshold=0.5)
    finally:
        try: os.unlink(tmp_path)
        except OSError: pass
    predictions.append({"index": idx, "img_id": ex["img_id"],
                        "question": question, "answer": final})
    if len(_samples) < 25:
        _samples.append((question, raw, final))

# -------- diagnostic --------
if predictions:
    answers = [p["answer"] for p in predictions]
    wc = [len(a.split()) for a in answers]
    freq = Counter(answers)
    print()
    print("=========== DIAGNOSTIC ===========")
    print(f"predictions      : {len(predictions)}")
    print(f"empty            : {sum(1 for a in answers if not a.strip())}")
    print(f"avg length       : {sum(wc)/len(wc):.2f}")
    print(f"top-20:")
    for a, c in freq.most_common(20):
        print(f"  {c:5d}  {a!r}")
    print("Sample (Q | raw | constrained):")
    for q, r, c in _samples[:15]:
        print(f"  Q: {q[:50]!r:55s}  raw={r!r:30s} -> {c!r}")
    print("==================================")

# -------- 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():.2f}MB | GPU: {gpu_name}")
print("Done. Results saved to 'predictions_1.json'.")
print(f"Run:: medvqa validate_and_submit --competition=gi-2026 --task=1 --repo_id={HF_REPO_ID}")