# === STANDALONE CELL 4: PASS 3 (4096 TOKENS + SANSKRIT/TECHNICAL HINT ON 720 TARGET QS) ===
import os, re, json, time, zipfile, shutil, subprocess, urllib.request
from pathlib import Path
from concurrent.futures import ThreadPoolExecutor, as_completed

# 1. Ensure both Ollama servers (GPU 0 -> 11434, GPU 1 -> 11435) are running
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:
        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=subprocess.DEVNULL, stderr=subprocess.DEVNULL)
        time.sleep(3)

# 2. Locate /kaggle/working/work/data (or extract from /kaggle/input if needed)
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] = str(r["answer"])
    return m

# NOTE (archive edit): the original cell embedded the 2238 Pass-1 answers as a digit string (P1_DIGITS) to resume
# after a Kaggle session reset. They are identical to ../passes/pass1_think_2048_orig.jsonl and are loaded from there.
P1_DIGITS = ''.join(json.loads(l)['answer'] for l in open('/kaggle/working/pass1_think_2048_orig.jsonl', encoding='utf-8') if l.strip())

all_qs = load_all_questions(data_dir)
p1_map = {f"{q['_level'].lower()}_{q['subject'].strip()}_{q['q_no']}": P1_DIGITS[i] for i, q in enumerate(all_qs)}
p2_map = load_ckpt(Path("/kaggle/working/predictions_gemma4_think_pass2_ckpt.jsonl"))
cot_map = load_ckpt(data_dir / "predictions_gemma4_cot.jsonl")
dir_map = load_ckpt(data_dir / "predictions_gemma4_direct.jsonl")

# Target: (A) 348 questions where P1 != P2, plus (B) 372 questions not in P2 where P1 != CoT or P1 != Direct
target_qs = []
for q in all_qs:
    k = f"{q['_level'].lower()}_{q['subject'].strip()}_{q['q_no']}"
    if k in p2_map:
        if p1_map.get(k) != p2_map.get(k):
            target_qs.append(q)
    else:
        if p1_map.get(k) != cot_map.get(k) or p1_map.get(k) != dir_map.get(k):
            target_qs.append(q)

p3_ckpt_path = Path("/kaggle/working/predictions_gemma4_think_pass3_ckpt.jsonl")
done_p3 = load_ckpt(p3_ckpt_path)
rem_p3 = [q for q in target_qs if f"{q['_level'].lower()}_{q['subject'].strip()}_{q['q_no']}" not in done_p3]
print(f"Pass 3 Target: {len(target_qs)} questions | Already done: {len(done_p3)} | Remaining: {len(rem_p3)}")

SYSTEM_THINK_PROMPT = (
    "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). If the question or options contain Sanskrit, "
    "Hindustani, Pali, Tamil, or technical terms, carefully map them to their standard curriculum "
    "meaning. You may reason in English. Finish with the final answer in the form 'Answer: <number>'."
)
ANSWER_THINK_RE = re.compile(r"(?:Answer|පිළිතුර)\s*[:：]?\s*\**\s*\(?([1-4])\)?", re.IGNORECASE)

def solve_p3(idx_q):
    idx, q = idx_q
    port = 11434 if (idx % 2 == 0) else 11435
    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)")
    opts = "\n".join(f"{i}. {c}" for i, c in enumerate(q["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 = {
        "model": "gemma4:e4b",
        "messages": [
            {"role": "system", "content": SYSTEM_THINK_PROMPT},
            {"role": "user", "content": user_msg},
        ],
        "think": True,
        "stream": False,
        "options": {"temperature": 0.1, "top_p": 0.95, "num_predict": 4096},
    }
    for attempt in range(3):
        try:
            req = urllib.request.Request(
                f"http://127.0.0.1:{port}/api/chat",
                data=json.dumps(payload).encode("utf-8"),
                headers={"Content-Type": "application/json"},
            )
            with urllib.request.urlopen(req, timeout=300) as resp:
                data = json.loads(resp.read().decode("utf-8"))
            content = data.get("message", {}).get("content", "") or ""
            thinking = data.get("message", {}).get("thinking", "") or ""
            m = ANSWER_THINK_RE.findall(content) or ANSWER_THINK_RE.findall(thinking)
            if m:
                ans = m[-1]
            else:
                digits = re.findall(r"\b([1-4])\b", content if content.strip() else thinking)
                ans = digits[-1] if digits else "1"
            return {
                "q_no": q_no,
                "subject": subj,
                "category": q.get("category", "general"),
                "level": lvl,
                "answer": str(ans),
            }
        except Exception:
            time.sleep(2)
    return {
        "q_no": q_no,
        "subject": subj,
        "category": q.get("category", "general"),
        "level": lvl,
        "answer": p1_map.get(q_key, "1"),
    }

t0 = time.time()
done_count = len(done_p3)
with open(p3_ckpt_path, "a", encoding="utf-8") as f_ckpt:
    with ThreadPoolExecutor(max_workers=4) as ex:
        futures = {ex.submit(solve_p3, (i, q)): q for i, q in enumerate(rem_p3)}
        for fut in as_completed(futures):
            res = fut.result()
            f_ckpt.write(json.dumps(res, ensure_ascii=False) + "\n")
            f_ckpt.flush()
            done_count += 1
            if done_count % 20 == 0 or done_count == len(target_qs):
                el = time.time() - t0
                rate = (done_count - len(done_p3)) / max(el, 1e-5)
                eta = (len(target_qs) - done_count) / max(rate, 1e-5) / 60
                print(f"[Pass 3: {done_count}/{len(target_qs)}] {rate:.2f} q/s | ETA: {eta:.1f} min", flush=True)
print("PASS 3 COMPLETE!")