#!/usr/bin/env python3
# Historical source snapshot. Comments about accuracy, data origin, or compliance
# are author claims, not verification; see training/README.md and REPRODUCIBILITY.md.
"""
Few-Shot In-Context Inference Runner for Sinhala MMLU Shared Task
Uses 100% external, balanced exemplars (2, 4, 1, 3) to eliminate position bias and ground the LLM.
Supports:
- Rapid local evaluation on the gold validation set (--val_only)
- Cross-Lingual Translate-and-Reason CoT mode (--cot) in English (--cot_lang en) or Tamil (--cot_lang ta)
- Thinking mode control (--think / default think=False) for Qwen 3 / 3.5 models
- Cyclic option-permutation debiasing (--passes 1, 2, or 4)
- Reliable checkpoint resuming
"""

import json
import argparse
import sys
from collections import Counter
from pathlib import Path

# Add src to sys.path
sys.path.insert(0, str(Path(__file__).parent))
from few_shot_prompts import build_few_shot_prompt
from baseline_runner import (
    extract_answer_digit,
    load_all_questions,
    load_val_questions,
    load_existing_checkpoint,
    query_ollama_chat,
    SYSTEM_EXAM_PROMPT,
    SYSTEM_COT_PROMPT,
)

SYSTEM_COT_TA_PROMPT = (
    "You are an expert examiner for the Sri Lankan school curriculum (Sinhala & Tamil medium). "
    "For each multiple-choice question, write a concise 1-2 sentence Tamil translation and fact-check ('Analysis:'), "
    "then end on a new line with strictly 'Answer: <digit>' (where <digit> is 1, 2, 3, or 4)."
)


def get_choices_list(q: dict) -> list:
    choices = q.get("choices")
    if choices is None:
        choices = q.get("options", [])
    if isinstance(choices, dict):
        keys = sorted(choices.keys(), key=lambda x: str(x))
        return [choices[k] for k in keys]
    elif isinstance(choices, list):
        return list(choices)
    return []


def predict_single_or_debiased(
    q: dict,
    model_name: str,
    num_shots: int = 4,
    think: bool = False,
    cot: bool = False,
    cot_lang: str = "en",
    passes: int = 1
) -> tuple[str, str]:
    """
    Runs 1-pass or multi-pass cyclic option-permutation inference on question q.
    Returns (raw_debug_str, chosen_original_digit).
    """
    if cot:
        sys_prompt = SYSTEM_COT_TA_PROMPT if cot_lang == "ta" else SYSTEM_COT_PROMPT
    else:
        sys_prompt = SYSTEM_EXAM_PROMPT

    orig_choices = get_choices_list(q)
    if passes <= 1 or len(orig_choices) != 4:
        prompt = build_few_shot_prompt(q, num_shots=num_shots, cot=cot, cot_lang=cot_lang)
        raw, ans_digit = query_ollama_chat(model_name, sys_prompt, prompt, think=think, cot=cot)
        return raw, ans_digit

    # Multi-pass cyclic permutation debiasing
    shifts = [0, 2] if passes == 2 else [0, 1, 2, 3][:passes]
    votes = []
    raw_logs = []

    for s in shifts:
        rotated_choices = orig_choices[s:] + orig_choices[:s]
        q_rot = dict(q)
        q_rot["choices"] = rotated_choices
        prompt = build_few_shot_prompt(q_rot, num_shots=num_shots, cot=cot, cot_lang=cot_lang)
        raw, rot_digit = query_ollama_chat(model_name, sys_prompt, prompt, think=think, cot=cot)
        rot_idx = int(rot_digit) - 1 if rot_digit in {"1", "2", "3", "4"} else 0
        orig_idx = (rot_idx + s) % 4
        orig_digit = str(orig_idx + 1)
        votes.append(orig_digit)
        raw_logs.append(f"s{s}:{rot_digit}->orig={orig_digit}")

    counts = Counter(votes)
    best_digit = counts.most_common(1)[0][0]
    return ",".join(raw_logs), best_digit


def run_few_shot(
    data_dir: str,
    output_file: str,
    model_name: str,
    num_shots: int = 4,
    val_only: bool = False,
    gold_val_file: str = "data/gold_val_41.jsonl",
    resume: bool = True,
    think: bool = False,
    cot: bool = False,
    cot_lang: str = "en",
    passes: int = 1
):
    output_path = Path(output_file)
    output_path.parent.mkdir(parents=True, exist_ok=True)

    if val_only:
        val_path = Path(gold_val_file)
        if not val_path.exists():
            print(f"Error: Gold validation file '{gold_val_file}' not found!")
            sys.exit(1)
        questions = load_val_questions(val_path)
        print(f"Loaded {len(questions)} validation questions for rapid local benchmark.")
    else:
        questions = load_all_questions(Path(data_dir))
        print(f"Loaded {len(questions)} full exam questions across easy, medium, and hard.")

    existing_preds = {}
    if resume and output_path.exists():
        existing_preds = load_existing_checkpoint(output_path)
        print(f"Resuming: Found {len(existing_preds)} existing predictions in '{output_file}'.")

    predictions = []

    print(f"Running {num_shots}-Shot inference with Ollama model: '{model_name}' (think={think}, cot={cot}, cot_lang={cot_lang}, passes={passes})...")

    for i, q in enumerate(questions, 1):
        q_no = q["q_no"]
        subj = q.get("subject", "general")
        cat = q.get("category", "general")
        lvl = q.get("_level", "medium").lower()

        q_key = f"{lvl}_{subj}_{q_no}"
        if resume and q_key in existing_preds:
            predictions.append(existing_preds[q_key])
            continue

        raw_info = ""
        try:
            raw_info, ans_digit = predict_single_or_debiased(
                q, model_name=model_name, num_shots=num_shots, think=think, cot=cot, cot_lang=cot_lang, passes=passes
            )
        except Exception as e:
            print(f"[{i}/{len(questions)}] Error on Q#{q_no} ({subj}): {e}. Using fallback 1.")
            ans_digit = "1"

        record = {
            "q_no": q_no,
            "subject": subj,
            "category": cat,
            "level": lvl,
            "answer": str(ans_digit)
        }
        predictions.append(record)

        if i % 10 == 0 or i == len(questions):
            snippet = raw_info.strip().replace("\n", " ")[-60:]
            print(f"Progress: {i}/{len(questions)} | Q#{q_no} ({subj}) [{snippet}] => {ans_digit}")
            with open(output_path, "w", encoding="utf-8") as fp:
                for p in predictions:
                    fp.write(json.dumps(p, ensure_ascii=False) + "\n")

    with open(output_path, "w", encoding="utf-8") as fp:
        for p in predictions:
            fp.write(json.dumps(p, ensure_ascii=False) + "\n")

    print(f"\nDone! {len(predictions)} predictions written to {output_file}")


def main():
    parser = argparse.ArgumentParser(description="Sinhala MMLU Few-Shot Runner")
    parser.add_argument("--data_dir", default="data", help="Path to data directory")
    parser.add_argument("--output", default="data/predictions_few_shot.jsonl", help="Output path")
    parser.add_argument("--model", default="qwen3.5:4b", help="Model name in Ollama")
    parser.add_argument("--shots", type=int, default=4, help="Number of in-context exemplars (e.g. 4)")
    parser.add_argument("--val_only", action="store_true", help="Run only on the 41-question local gold benchmark")
    parser.add_argument("--gold_val_file", default="data/gold_val_41.jsonl", help="Gold validation benchmark file")
    parser.add_argument("--no_resume", action="store_true", help="Overwrite existing output file")
    parser.add_argument("--think", action="store_true", help="Enable <think> reasoning mode for Qwen 3 / 3.5 models")
    parser.add_argument("--cot", action="store_true", help="Enable concise Cross-Lingual Translate-and-Reason CoT mode")
    parser.add_argument("--cot_lang", default="en", choices=["en", "ta"], help="Target bridge language for CoT ('en' or 'ta')")
    parser.add_argument("--passes", type=int, default=1, choices=[1, 2, 4], help="Cyclic option-permutation passes (1, 2, or 4)")
    args = parser.parse_args()

    run_few_shot(
        args.data_dir,
        args.output,
        args.model,
        num_shots=args.shots,
        val_only=args.val_only,
        gold_val_file=args.gold_val_file,
        resume=not args.no_resume,
        think=args.think,
        cot=args.cot,
        cot_lang=args.cot_lang,
        passes=args.passes
    )


if __name__ == "__main__":
    main()