# Self-Contained Cell 3: Pass 2 (Shuffled Options + 4096 Token Budget on Hard & Disputed Questions)
import json, os, re, time, zipfile, shutil, subprocess, urllib.request
from pathlib import Path
from concurrent.futures import ThreadPoolExecutor, as_completed

# 1. Ensure both Ollama servers (11434 and 11435) are alive
for idx, port in [(0, 11434), (1, 11435)]:
    try:
        urllib.request.urlopen(f"http://127.0.0.1:{port}/api/tags", timeout=3)
    except Exception:
        print(f"Restarting Ollama server on GPU #{idx} (port {port})...")
        env = {**os.environ, "CUDA_VISIBLE_DEVICES": str(idx), "OLLAMA_HOST": f"127.0.0.1:{port}", "OLLAMA_NUM_PARALLEL": "2", "OLLAMA_FLASH_ATTENTION": "1", "OLLAMA_KEEP_ALIVE": "24h"}
        subprocess.Popen(["ollama", "serve"], env=env, stdout=open(f"/tmp/ollama{idx}.log", "w"), stderr=subprocess.STDOUT)
        time.sleep(4)

# 2. Ensure data directory and helper functions are defined
work_dir = Path("/kaggle/working/work")
work_dir.mkdir(parents=True, exist_ok=True)
data_dir = work_dir / "data"
if not data_dir.exists():
    for p in Path("/kaggle/input").rglob("*"):
        if p.name == "kaggle_data.zip":
            with zipfile.ZipFile(p, "r") as z:
                z.extractall(work_dir)
            break
        elif p.name == "predictions_gemma4_cot.jsonl":
            shutil.copytree(p.parent.parent, work_dir, dirs_exist_ok=True)
            break

def load_all_questions(d_dir: Path):
    qs = []
    for lvl in ["easy", "medium", "hard"]:
        for f in sorted((d_dir / lvl).glob("*.json")):
            for item in json.load(open(f, encoding="utf-8")):
                item["_level"] = lvl
                qs.append(item)
    return qs

def load_ckpt(path: Path):
    m = {}
    if path.exists():
        for line in open(path, encoding="utf-8"):
            if line.strip():
                r = json.loads(line)
                k = f"{str(r['level']).lower().strip()}_{str(r['subject']).strip()}_{str(r['q_no']).strip()}"
                m[k] = r
    return m

canonical_qs = load_all_questions(data_dir)
cot_map = load_ckpt(data_dir / "predictions_gemma4_cot.jsonl")
dir_map = load_ckpt(data_dir / "predictions_gemma4_direct.jsonl")
p1_map = load_ckpt(Path("/kaggle/working/predictions.jsonl"))

lvl_rank = {"hard": 0, "medium": 1, "easy": 2}
def sort_key(q):
    k = f"{q['_level'].lower()}_{q['subject'].strip()}_{q['q_no']}"
    disputed = 0 if cot_map.get(k, {}).get("answer") != dir_map.get(k, {}).get("answer") else 1
    return (disputed, lvl_rank.get(q["_level"].lower(), 3))

priority_qs = sorted(canonical_qs, key=sort_key)

SYSTEM_THINK = (
    "You are an expert teacher of the Sri Lankan national school curriculum (grades 6-13), "
    "fluent in Sinhala. You will receive a multiple-choice question written in Sinhala with four "
    "numbered options. Exactly one option is correct. Read the Sinhala carefully, including "
    "negations such as 'නොවන' / 'නොවේ' (NOT). You may reason in English. "
    "Finish with the final answer in the form 'Answer: <number>'."
)
ANS_RE = re.compile(r"(?:Answer|පිළිතුර)\s*[:：]?\s*\**\s*\(?([1-4])\)?", re.I)
HOSTS = ["http://127.0.0.1:11434", "http://127.0.0.1:11435"]

p2_ckpt_path = Path("/kaggle/working/predictions_gemma4_think_pass2_ckpt.jsonl")

# Target the 749 Disputed questions + all remaining Hard questions (1,074 questions total)
target_qs = []
for q in priority_qs:
    k = f"{q['_level'].lower()}_{q['subject'].strip()}_{q['q_no']}"
    is_disp = cot_map.get(k, {}).get("answer") != dir_map.get(k, {}).get("answer")
    if is_disp or q["_level"].lower() == "hard":
        target_qs.append(q)

p2_map = load_ckpt(p2_ckpt_path)
todo_p2 = [(i, q) for i, q in enumerate(target_qs) if f"{q['_level'].lower()}_{q['subject'].strip()}_{q['q_no']}" not in p2_map]
print(f"Pass 2 (Shuffled + 4096 tokens) | Already done: {len(p2_map)} | Remaining: {len(todo_p2)}")

def solve_question_shuffled(idx_q):
    idx, q = idx_q
    host = HOSTS[idx % 2]
    lvl = q["_level"].lower()
    subj = q.get("subject", "General")
    q_no = q["q_no"]
    q_key = f"{lvl}_{subj.strip()}_{q_no}"
    grade = "6-9" if lvl == "easy" else ("10-11 (O/L)" if lvl == "medium" else "12-13 (A/L)")

    orig_choices = list(q["choices"])
    rot_choices = orig_choices[1:] + orig_choices[:1]
    opts = "\n".join(f"{i}. {c}" for i, c in enumerate(rot_choices, 1))
    user_msg = (
        f"විෂය (Subject): {subj} | ශ්‍රේණිය (Grade): {grade}\n\n"
        f"ප්‍රශ්නය (Question):\n{q['question'].strip()}\n\n"
        f"විකල්ප (Options):\n{opts}\n\n"
        "Which option (1, 2, 3 or 4) is correct? End your reply with 'Answer: <number>'."
    )
    payload = json.dumps({
        "model": "gemma4:e4b",
        "messages": [{"role": "system", "content": SYSTEM_THINK}, {"role": "user", "content": user_msg}],
        "stream": False,
        "think": True,
        "options": {"temperature": 0.0, "num_predict": 4096}
    }).encode("utf-8")
    try:
        req = urllib.request.Request(f"{host}/api/chat", data=payload, headers={"Content-Type": "application/json"})
        with urllib.request.urlopen(req, timeout=300) as resp:
            res = json.loads(resp.read().decode("utf-8"))
            msg = res.get("message", {})
            content, thinking = msg.get("content", ""), msg.get("thinking", "")
            m = ANS_RE.findall(content) or ANS_RE.findall(thinking)
            if m:
                rot_ans = m[-1]
            else:
                digits = re.findall(r"\b([1-4])\b", content if content.strip() else thinking)
                rot_ans = digits[-1] if digits else "1"
            rot_idx = int(rot_ans) - 1 if rot_ans in {"1", "2", "3", "4"} else 0
            orig_ans = str(((rot_idx + 1) % 4) + 1)
    except Exception:
        orig_ans = (p1_map.get(q_key) or cot_map.get(q_key, {})).get("answer", "1")

    return q_key, {
        "q_no": q_no,
        "subject": subj,
        "category": q.get("category", "general"),
        "level": lvl,
        "answer": orig_ans
    }

t0 = time.time()
done_p2 = 0
with ThreadPoolExecutor(max_workers=4) as ex:
    futures = [ex.submit(solve_question_shuffled, item) for item in todo_p2]
    for fut in as_completed(futures):
        k, rec = fut.result()
        p2_map[k] = rec
        done_p2 += 1
        if done_p2 % 20 == 0 or done_p2 == len(todo_p2):
            with open(p2_ckpt_path, "w", encoding="utf-8") as f:
                for r in p2_map.values():
                    f.write(json.dumps(r, ensure_ascii=False) + "\n")
            elapsed = time.time() - t0
            print(f"Pass 2 [{len(p2_map)}/{len(target_qs)}] ({elapsed/60:.1f}m elapsed, {elapsed/done_p2:.1f}s/q)")