#!/usr/bin/env python3
# mmlu_eval_sft.py — eval loop for JiRack models + automatic SFT set from the failures.
#
#   1. cais/mmlu ("all", split "test": 14042 questions, 57 subjects)
#   2. your model behind an OpenAI-compatible endpoint (Ollama: http://localhost:11434/v1)
#   3. the answer letter is checked against the MMLU key HERE (no judge)
#   4. accuracy overall and per subject -> summary.json
#   5. only for wrong / unparsed answers: DeepSeek writes the step-by-step reasoning
#      for the KNOWN correct answer -> sft.jsonl ("question -> reasoning -> answer")
#
# Both endpoints use the same openai SDK client; only base_url / api_key / model differ.
# Keys come from the environment only:
#   MODEL_API_KEY     (Ollama ignores it; default "ollama")
#   DEEPSEEK_API_KEY  (required for the sft stage)
#
# Install:  pip install openai datasets
# Examples:
#   # quick smoke test, 50 questions, eval only
#   python3 mmlu_eval_sft.py --model jirack-llama2b --limit 50 --stage eval
#   # full run: eval, then SFT from the failures
#   export DEEPSEEK_API_KEY=...        # in the terminal, never in chat
#   python3 mmlu_eval_sft.py --model jirack-llama2b --stage all --workers 4
#   # later the ternary 16B: only the config changes
#   python3 mmlu_eval_sft.py --model jirack-16b-ternary --out runs/16b
#
# Restart-safe: results.jsonl and sft.jsonl are appended per question; a rerun with the
# same --out skips everything already done.
#
# License note: sft.jsonl rows carry "generator" = the DeepSeek model. Output of the
# DeepSeek API falls under DeepSeek's terms (distillation allowed, AI-generated content
# must be disclosed) — keep that field when the data goes into training.

import argparse
import json
import os
import random
import re
import sys
import threading
import time
from collections import defaultdict
from concurrent.futures import ThreadPoolExecutor, as_completed

LETTERS = "ABCD"

ANSWER_SYSTEM = ("You are taking a multiple-choice test. Reply with the letter of the "
                 "correct option (A, B, C or D) and nothing else.")

REASON_SYSTEM = ("You are an expert tutor. You are given a multiple-choice question and its "
                 "correct answer. Explain step by step why that answer is correct and, briefly, "
                 "why the other options are wrong. Use facts, not guesses. Be concise "
                 "(at most about 200 words). The last line must be exactly: Answer: <letter>")


# ------------------------------------------------------------------ helpers
def format_question(q, choices):
    lines = [q.strip(), ""]
    for letter, text in zip(LETTERS, choices):
        lines.append(f"{letter}. {str(text).strip()}")
    return "\n".join(lines)


def parse_letter(text, choices):
    """Letter A-D from a model reply; None if it cannot be read unambiguously."""
    if not text:
        return None
    t = text.strip()
    # reasoning models: drop a <think>...</think> block
    t = re.sub(r"<think>.*?</think>", "", t, flags=re.S).strip()
    m = re.search(r"(?:answer|ответ)\s*(?:is|:)?\s*\(?([ABCD])\b", t, flags=re.I)
    if m:
        return m.group(1).upper()
    m = re.match(r"^\(?([ABCD])(?:[\).:,\s]|$)", t)
    if m:
        return m.group(1).upper()
    found = set(re.findall(r"\b([ABCD])\b", t))
    if len(found) == 1:
        return found.pop()
    # model replied with the option text instead of the letter
    low = t.lower().rstrip(".")
    hits = [LETTERS[i] for i, c in enumerate(choices) if str(c).strip().lower() == low]
    return hits[0] if len(hits) == 1 else None


def final_letter(text):
    m = re.findall(r"Answer:\s*\(?([ABCD])\b", text or "", flags=re.I)
    return m[-1].upper() if m else None


def load_done(path, key="idx"):
    done = {}
    if os.path.exists(path):
        with open(path, encoding="utf-8") as f:
            for line in f:
                line = line.strip()
                if line:
                    row = json.loads(line)
                    done[row[key]] = row
    return done


class JsonlWriter:
    def __init__(self, path):
        self.f = open(path, "a", encoding="utf-8")
        self.lock = threading.Lock()

    def write(self, row):
        with self.lock:
            self.f.write(json.dumps(row, ensure_ascii=False) + "\n")
            self.f.flush()

    def close(self):
        self.f.close()


def chat(client, model, messages, max_tokens, temperature, retries, timeout):
    """chat.completions with retry on rate limits, timeouts, connection and 5xx errors."""
    import openai
    delay = 2.0
    for attempt in range(1, retries + 1):
        try:
            r = client.chat.completions.create(model=model, messages=messages,
                                               max_tokens=max_tokens, temperature=temperature,
                                               timeout=timeout)
            return r.choices[0].message.content or ""
        except (openai.RateLimitError, openai.APITimeoutError, openai.APIConnectionError,
                openai.InternalServerError) as e:
            err = e
        except openai.APIStatusError as e:
            if e.status_code < 500 and e.status_code != 429:
                raise                      # 400/401/404: retrying will not help
            err = e
        if attempt == retries:
            raise err
        time.sleep(delay + random.random())
        delay = min(delay * 2, 60)


# ------------------------------------------------------------------ data
def load_mmlu(args):
    from datasets import load_dataset
    ds = load_dataset("cais/mmlu", args.config, split=args.split)
    rows = [{"idx": i, "subject": r["subject"], "question": r["question"],
             "choices": list(r["choices"]), "answer": LETTERS[int(r["answer"])]}
            for i, r in enumerate(ds)]
    if args.subjects:
        keep = set(args.subjects.split(","))
        rows = [r for r in rows if r["subject"] in keep]
    if args.shuffle:
        random.Random(args.seed).shuffle(rows)
    if args.limit:
        rows = rows[:args.limit]
    return rows


# ------------------------------------------------------------------ stage 1: eval
def run_eval(args, rows):
    from openai import OpenAI
    client = OpenAI(base_url=args.model_base_url,
                    api_key=os.environ.get("MODEL_API_KEY", "ollama"), max_retries=0)
    path = os.path.join(args.out, "results.jsonl")
    done = load_done(path)
    todo = [r for r in rows if r["idx"] not in done]
    print(f"[eval] {args.model} @ {args.model_base_url}: {len(rows)} questions, "
          f"{len(rows) - len(todo)} already done, {len(todo)} to run")
    w = JsonlWriter(path)

    def one(r):
        prompt = format_question(r["question"], r["choices"])
        try:
            reply = chat(client, args.model,
                         [{"role": "system", "content": ANSWER_SYSTEM},
                          {"role": "user", "content": prompt}],
                         args.answer_max_tokens, args.temperature, args.retries, args.timeout)
            error = None
        except Exception as e:                       # keep going; the row is marked
            reply, error = "", f"{type(e).__name__}: {e}"
        pred = parse_letter(reply, r["choices"])
        row = {**r, "model": args.model, "reply": reply, "pred": pred,
               "correct": pred == r["answer"], "error": error}
        w.write(row)
        return row

    t0, n = time.time(), 0
    with ThreadPoolExecutor(max_workers=args.workers) as ex:
        for fut in as_completed([ex.submit(one, r) for r in todo]):
            fut.result()
            n += 1
            if n % 50 == 0 or n == len(todo):
                print(f"[eval] {n}/{len(todo)} | {(time.time() - t0) / n:.2f} s/question")
    w.close()


def summarize(args, rows):
    wanted = {r["idx"] for r in rows}
    res = [r for i, r in load_done(os.path.join(args.out, "results.jsonl")).items() if i in wanted]
    by = defaultdict(lambda: [0, 0])
    for r in res:
        by[r["subject"]][0] += r["correct"]
        by[r["subject"]][1] += 1
    n = len(res)
    correct = sum(r["correct"] for r in res)
    summary = {
        "model": args.model, "dataset": f"cais/mmlu/{args.config}/{args.split}",
        "questions": n, "correct": correct,
        "accuracy": round(correct / n, 4) if n else None,
        "unparsed": sum(r["pred"] is None and not r["error"] for r in res),
        "api_errors": sum(bool(r["error"]) for r in res),
        "per_subject": {s: {"accuracy": round(c / t, 4), "n": t}
                        for s, (c, t) in sorted(by.items(), key=lambda kv: kv[1][0] / kv[1][1])},
    }
    with open(os.path.join(args.out, "summary.json"), "w", encoding="utf-8") as f:
        json.dump(summary, f, ensure_ascii=False, indent=2)
    print(f"[summary] accuracy {summary['accuracy']} ({correct}/{n}) | unparsed "
          f"{summary['unparsed']} | api errors {summary['api_errors']} (25% = random guessing)")
    for s, v in list(summary["per_subject"].items())[:5]:
        print(f"   weakest: {s:40s} {v['accuracy']:.3f} (n={v['n']})")
    return res


# ------------------------------------------------------------------ stage 2: SFT from failures
def run_sft(args, res):
    key = os.environ.get("DEEPSEEK_API_KEY")
    if not key:
        sys.exit("DEEPSEEK_API_KEY is not set (export it in the terminal)")
    from openai import OpenAI
    client = OpenAI(base_url=args.deepseek_base_url, api_key=key, max_retries=0)
    # API errors are excluded: the model never answered, that is not a knowledge failure
    fails = [r for r in res if not r["correct"] and not r["error"]]
    path = os.path.join(args.out, "sft.jsonl")
    done = load_done(path, key="idx")
    rej_path = os.path.join(args.out, "sft_rejected.jsonl")
    rejected = load_done(rej_path, key="idx")
    todo = [r for r in fails if r["idx"] not in done and r["idx"] not in rejected]
    print(f"[sft] {len(fails)} failures, {len(fails) - len(todo)} already handled, "
          f"{len(todo)} to send to {args.deepseek_model}")
    w, wr = JsonlWriter(path), JsonlWriter(rej_path)

    def one(r):
        prompt = format_question(r["question"], r["choices"])
        ans = r["answer"]
        ask = (f"{prompt}\n\nCorrect answer: {ans}. "
               f"{r['choices'][LETTERS.index(ans)]}\n\nExplain why.")
        try:
            text = chat(client, args.deepseek_model,
                        [{"role": "system", "content": REASON_SYSTEM},
                         {"role": "user", "content": ask}],
                        args.reason_max_tokens, 0.3, args.retries, args.timeout).strip()
        except Exception as e:
            wr.write({"idx": r["idx"], "reason": f"api: {type(e).__name__}: {e}"})
            return False
        if final_letter(text) != ans:   # explanation argues for another option -> drop it
            wr.write({"idx": r["idx"], "reason": "final letter mismatch", "text": text})
            return False
        body = re.sub(r"\n*\s*Answer:\s*\(?[ABCD]\)?\.?\s*$", "", text, flags=re.I).strip()
        assistant = f"{body}\n\nAnswer: {ans}. {r['choices'][LETTERS.index(ans)]}"
        w.write({
            "idx": r["idx"], "subject": r["subject"],
            "messages": [{"role": "user", "content": prompt},
                         {"role": "assistant", "content": assistant}],
            "model_pred": r["pred"], "answer": ans,
            "source": f"cais/mmlu/{args.config}/{args.split}",
            "generator": f"deepseek:{args.deepseek_model}",
        })
        return True

    ok = 0
    with ThreadPoolExecutor(max_workers=args.workers) as ex:
        for i, fut in enumerate(as_completed([ex.submit(one, r) for r in todo]), 1):
            ok += fut.result()
            if i % 25 == 0 or i == len(todo):
                print(f"[sft] {i}/{len(todo)} | written {ok}")
    w.close()
    wr.close()
    print(f"[sft] {path}: {len(load_done(path))} rows | rejected: {rej_path}")


# ------------------------------------------------------------------ main
def main():
    ap = argparse.ArgumentParser(description="MMLU eval + SFT-from-failures via OpenAI-compatible APIs")
    ap.add_argument("--stage", choices=["eval", "sft", "all"], default="all")
    ap.add_argument("--out", default="runs/llama2b", help="output folder")
    ap.add_argument("--model", required=True, help="model name on your server (ollama list)")
    ap.add_argument("--model-base-url", default=os.environ.get("MODEL_BASE_URL", "http://localhost:11434/v1"))
    ap.add_argument("--deepseek-base-url", default=os.environ.get("DEEPSEEK_BASE_URL", "https://api.deepseek.com"))
    ap.add_argument("--deepseek-model", default=os.environ.get("DEEPSEEK_MODEL", "deepseek-chat"))
    ap.add_argument("--config", default="all")
    ap.add_argument("--split", default="test")
    ap.add_argument("--subjects", default="", help="comma list, e.g. high_school_geography,world_religions")
    ap.add_argument("--limit", type=int, default=0, help="0 = all questions")
    ap.add_argument("--shuffle", action="store_true", help="random subset with --limit")
    ap.add_argument("--seed", type=int, default=42)
    ap.add_argument("--workers", type=int, default=1, help="parallel requests")
    ap.add_argument("--temperature", type=float, default=0.0)
    ap.add_argument("--answer-max-tokens", type=int, default=16)
    ap.add_argument("--reason-max-tokens", type=int, default=600)
    ap.add_argument("--retries", type=int, default=5)
    ap.add_argument("--timeout", type=float, default=120.0)
    args = ap.parse_args()

    os.makedirs(args.out, exist_ok=True)
    rows = load_mmlu(args)
    if args.stage in ("eval", "all"):
        run_eval(args, rows)
    res = summarize(args, rows)
    if args.stage in ("sft", "all"):
        run_sft(args, res)


if __name__ == "__main__":
    main()