"""
finetune/train_modal.py — QLoRA fine-tune of MiniCPM4-8B for PaperProf.

Task: exam-question generation. The training data is SQuAD passages and
questions reformatted into the EXACT production prompt used by
core/questioner.py, so the model learns PaperProf's task, not generic QA.

Run:
    modal run finetune/train_modal.py

Output:
    Merged bf16 model pushed to hf.co/build-small-hackathon/MiniCPM4-8B-PaperProf
"""

import modal

APP_NAME = "paperprof-finetune-41"
BASE_MODEL = "openbmb/MiniCPM4.1-8B"
HUB_REPO = "build-small-hackathon/MiniCPM4.1-8B-PaperProf"
N_SAMPLES = 3000
MAX_LEN = 1024

app = modal.App(APP_NAME)

image = (
    modal.Image.debian_slim(python_version="3.12")
    .pip_install(
        "torch==2.6.0",
        "transformers==4.57.1",
        "datasets==3.2.0",
        "peft==0.14.0",
        "bitsandbytes==0.46.0",
        "accelerate>=1.8.0",
        "huggingface_hub",
        "sentencepiece",
    )
)

# Mirror of core/questioner.py _PROMPT_TEMPLATE (Normal difficulty)
PROMPT_TEMPLATE = """\
You are a university professor creating exam questions.
Given the following excerpt from a course, write ONE focused question.

Difficulty — Ask for conceptual understanding (Explain X. Why does X happen?).

Rules:
- ONE question only, on ONE concept
- Maximum 25 words
- No sub-questions, no "and", no compound questions
- IMPORTANT: Always write the question in English, even if the source text is in another language
- Output only the question, nothing else

Excerpt:
{chunk}

Question:"""

# Mirror of core/questioner.py _MCQ_TEMPLATE
MCQ_TEMPLATE = """\
You are a university professor creating a multiple choice exam question.
Given the following excerpt, write ONE multiple choice question.

Rules:
- One clear question, maximum 25 words
- Exactly 4 options (A, B, C, D), only ONE is correct
- All wrong options must be plausible — no obviously wrong answers
- IMPORTANT: Always write everything in English, even if the source text is in another language
- For each option, write a 1-sentence explanation of why it is correct or incorrect

Output format (use EXACTLY these labels, one per line, nothing else):
QUESTION: <question>
A) <option>
B) <option>
C) <option>
D) <option>
CORRECT: <A, B, C or D>
EXPLAIN_A: <explanation>
EXPLAIN_B: <explanation>
EXPLAIN_C: <explanation>
EXPLAIN_D: <explanation>

Excerpt:
{chunk}
"""

# Mirror of core/evaluator.py _PROMPT_EN
EVAL_TEMPLATE = """\
You are a patient and constructive university tutor.
IMPORTANT: Always write your entire response in English, even if the source material or student answer is in another language.

Source material:
{chunk}

Question asked to the student:
{question}

Student's answer:
{answer}

Evaluate the answer using this exact structure:
1. Verdict: Correct / Partially correct / Incorrect
2. What was good about the answer.
3. What was missing or imprecise.
4. A concise model answer (2-4 sentences).

Be encouraging and specific. Write in English only."""

MODEL_CARD = f"""---
license: apache-2.0
base_model: {BASE_MODEL}
tags:
  - question-generation
  - education
  - lora
  - paperprof
datasets:
  - squad
  - allenai/sciq
language:
  - en
---

# MiniCPM4.1-8B-PaperProf

Fine-tuned from [{BASE_MODEL}](https://huggingface.co/{BASE_MODEL}) for
**exam-question generation** in [PaperProf](https://huggingface.co/spaces/build-small-hackathon/PaperProf),
an AI study buddy that turns course PDFs into interactive quiz sessions.

## Training

- **Method:** QLoRA (4-bit NF4, r=16, alpha=32, all-linear targets), merged to bf16
- **Data:** ~3500 multi-task pairs in PaperProf's three production formats:
  open question generation (SQuAD), MCQ with distractors and per-option
  explanations (SciQ), and structured answer evaluation (SQuAD-derived),
  so the model is optimized for the exact tasks it serves.
- **Epochs:** 1, lr 2e-4 cosine, bf16 compute

## Usage

Drop-in replacement for the base model:

```python
from transformers import AutoTokenizer, AutoModelForCausalLM
tok = AutoTokenizer.from_pretrained("{HUB_REPO}", trust_remote_code=True)
model = AutoModelForCausalLM.from_pretrained("{HUB_REPO}", trust_remote_code=True, torch_dtype="bfloat16")
```

Built for the Build Small Hackathon, June 2026, by Team PaperProf (EPITA).
"""


@app.function(
    image=image,
    gpu="A100-80GB",
    timeout=3 * 60 * 60,
    secrets=[modal.Secret.from_name("paperprof-hf")],
)
def train():
    import os
    import torch
    from datasets import load_dataset
    from transformers import (
        AutoTokenizer, AutoModelForCausalLM, BitsAndBytesConfig,
        Trainer, TrainingArguments,
    )
    from peft import LoraConfig, get_peft_model, prepare_model_for_kbit_training
    from huggingface_hub import HfApi

    token = os.environ["HF_TOKEN"]

    # ── Resolve a writable repo BEFORE burning GPU time ──────────────────
    from huggingface_hub import create_repo
    from huggingface_hub.utils import HfHubHTTPError

    api = HfApi(token=token)
    try:
        create_repo(HUB_REPO, token=token, exist_ok=True)
        hub_repo = HUB_REPO
    except HfHubHTTPError:
        user = api.whoami()["name"]
        hub_repo = f"{user}/{HUB_REPO.split('/')[-1]}"
        create_repo(hub_repo, token=token, exist_ok=True)
        print(f"[push] no write access to {HUB_REPO}, falling back to {hub_repo}")
    print(f"[push] target repo: {hub_repo}")

    # ── Data: multi-task — the THREE production formats ───────────────────
    # A narrow single-task fine-tune degraded the other formats (MCQ came
    # back empty), so we train on question-gen + MCQ + evaluation together.
    import random as rnd
    rnd.seed(42)
    pairs = []  # (user_prompt, assistant_output)

    # Task A — open question generation (SQuAD)
    print("[data] loading SQuAD…")
    squad = load_dataset("squad", split="train")
    seen, qa_rows = set(), []
    for ex in squad:
        ctx = ex["context"].strip()
        if not (300 <= len(ctx) <= 1500) or ctx in seen:
            continue
        seen.add(ctx)
        q = ex["question"].strip()
        if not q.endswith("?") or len(q.split()) > 25:
            continue
        qa_rows.append(ex)
        pairs.append((PROMPT_TEMPLATE.format(chunk=ctx), q))
        if len(pairs) >= 1500:
            break
    print(f"[data] task A (question gen): {len(pairs)}")

    # Task B — MCQ generation (SciQ: question + 3 distractors + support text)
    # Explanations are grounded in the support text (quote the sentence that
    # justifies the answer) with varied phrasing — generic templates taught
    # the model to parrot one hollow sentence per option.
    import re as _re

    def _evidence(sup: str, answer: str) -> str:
        sents = [s.strip() for s in _re.split(r"(?<=[.!?])\s+", sup) if len(s.strip()) > 20]
        for s in sents:
            if answer.lower() in s.lower():
                return s[:160].rstrip(".") + ("…" if len(s) > 160 else ".")
        return (sents[0][:160] + "…") if sents else sup[:160]

    CORRECT_TPL = [
        'Correct — the excerpt states: "{ev}"',
        'This is the right answer. The source text says: "{ev}"',
        'Correct: the passage directly supports this — "{ev}"',
        'Yes — as the excerpt explains, "{ev}"',
    ]
    WRONG_TPL = [
        'Incorrect — the excerpt instead supports {truth}: "{ev}"',
        "Plausible, but the passage points to {truth}, not {opt}.",
        'Not supported by the text, which states: "{ev}"',
        "Incorrect: {opt} is not what the excerpt describes — the answer is {truth}.",
    ]

    print("[data] loading SciQ…")
    sciq = load_dataset("allenai/sciq", split="train")
    n_mcq = 0
    for ex in sciq:
        sup = (ex["support"] or "").strip()
        if not (250 <= len(sup) <= 1500):
            continue
        options = [ex["correct_answer"], ex["distractor1"], ex["distractor2"], ex["distractor3"]]
        options = [o.strip() for o in options]
        if any(not o for o in options):
            continue
        ev = _evidence(sup, options[0])
        order = rnd.sample(range(4), 4)
        letters = ["A", "B", "C", "D"]
        correct_letter = letters[order.index(0)]
        lines = [f"QUESTION: {ex['question'].strip()}"]
        for idx, l in zip(order, letters):
            lines.append(f"{l}) {options[idx]}")
        lines.append(f"CORRECT: {correct_letter}")
        for idx, l in zip(order, letters):
            if idx == 0:
                tpl = rnd.choice(CORRECT_TPL)
                lines.append(f"EXPLAIN_{l}: {tpl.format(ev=ev)}")
            else:
                tpl = rnd.choice(WRONG_TPL)
                lines.append(f"EXPLAIN_{l}: {tpl.format(ev=ev, truth=options[0], opt=options[idx])}")
        pairs.append((MCQ_TEMPLATE.format(chunk=sup), "\n".join(lines)))
        n_mcq += 1
        if n_mcq >= 1200:
            break
    print(f"[data] task B (MCQ): {n_mcq}")

    # Task C — answer evaluation (SQuAD answers, half correct half wrong)
    n_eval = 0
    all_answers = [ex["answers"]["text"][0] for ex in qa_rows if ex["answers"]["text"]]
    for ex in qa_rows:
        if not ex["answers"]["text"]:
            continue
        truth = ex["answers"]["text"][0].strip()
        ctx, q = ex["context"].strip(), ex["question"].strip()
        ev = _evidence(ctx, truth)
        if n_eval % 2 == 0:
            student, verdict = truth, "Correct"
            good = rnd.choice([
                f"You correctly identified {truth}, which is exactly what the source material states.",
                f"Your answer matches the excerpt — {truth} is precisely the point the passage makes.",
                f"Spot on: {truth} is the answer, and you stated it clearly.",
            ])
            missing = rnd.choice([
                "Nothing essential is missing — your answer covers the key point.",
                f'You could strengthen it by citing the supporting passage: "{ev}"',
                "Your answer is complete for this question; adding brief context would make it even stronger.",
            ])
        else:
            student = rnd.choice(all_answers)
            if student == truth:
                continue
            verdict = "Incorrect"
            good = rnd.choice([
                "You gave a direct, focused answer to the question.",
                "Your answer shows you engaged with the question seriously.",
                "You committed to a clear answer rather than hedging, which is good practice.",
            ])
            missing = rnd.choice([
                f'The answer does not match the source. The excerpt states: "{ev}"',
                f"The passage points to {truth} instead — re-read the relevant section: \"{ev}\"",
                f"Your answer is not supported by the text, which indicates {truth}.",
            ])
        out = (
            f"1. Verdict: {verdict}\n"
            f"2. {good}\n"
            f"3. {missing}\n"
            f'4. Based on the excerpt, the answer is {truth}: "{ev}"'
        )
        pairs.append((EVAL_TEMPLATE.format(chunk=ctx, question=q, answer=student), out))
        n_eval += 1
        if n_eval >= 800:
            break
    print(f"[data] task C (evaluation): {n_eval}")

    rnd.shuffle(pairs)
    print(f"[data] {len(pairs)} training pairs total")

    tokenizer = AutoTokenizer.from_pretrained(BASE_MODEL, trust_remote_code=True, token=token)
    if tokenizer.pad_token is None:
        tokenizer.pad_token = tokenizer.eos_token

    def to_text(pair):
        messages = [
            {"role": "user", "content": pair[0]},
            {"role": "assistant", "content": pair[1]},
        ]
        return tokenizer.apply_chat_template(messages, tokenize=False, enable_thinking=False)

    texts = [to_text(p) for p in pairs]

    def tokenize(batch):
        return tokenizer(batch["text"], truncation=True, max_length=MAX_LEN, padding=False)

    from datasets import Dataset
    train_ds = Dataset.from_dict({"text": texts}).map(
        tokenize, batched=True, remove_columns=["text"]
    )

    # ── Model: 4-bit base + LoRA ──────────────────────────────────────────
    print("[model] loading base in 4-bit…")
    bnb = BitsAndBytesConfig(
        load_in_4bit=True,
        bnb_4bit_quant_type="nf4",
        bnb_4bit_compute_dtype=torch.bfloat16,
        bnb_4bit_use_double_quant=True,
    )
    model = AutoModelForCausalLM.from_pretrained(
        BASE_MODEL, quantization_config=bnb, device_map="auto",
        trust_remote_code=True, token=token,
    )
    model = prepare_model_for_kbit_training(model)
    lora = LoraConfig(
        r=16, lora_alpha=32, lora_dropout=0.05,
        target_modules="all-linear", task_type="CAUSAL_LM",
    )
    model = get_peft_model(model, lora)
    model.print_trainable_parameters()

    # ── Train ─────────────────────────────────────────────────────────────
    args = TrainingArguments(
        output_dir="/tmp/out",
        num_train_epochs=1,
        per_device_train_batch_size=4,
        gradient_accumulation_steps=4,
        learning_rate=2e-4,
        lr_scheduler_type="cosine",
        warmup_ratio=0.03,
        bf16=True,
        logging_steps=10,
        save_strategy="no",
        report_to=[],
        gradient_checkpointing=True,
    )
    # Custom collator: pad_token == eos_token, so the stock LM collator would
    # mask every EOS in the labels and the model would never learn to stop
    # (v2 regression: runaway generation). Mask ONLY actual padding positions.
    pad_id = tokenizer.pad_token_id

    def collate(features):
        max_len = max(len(f["input_ids"]) for f in features)
        batch = {"input_ids": [], "attention_mask": [], "labels": []}
        for f in features:
            ids = list(f["input_ids"])
            n_pad = max_len - len(ids)
            batch["input_ids"].append(ids + [pad_id] * n_pad)
            batch["attention_mask"].append([1] * len(ids) + [0] * n_pad)
            batch["labels"].append(ids + [-100] * n_pad)
        return {k: torch.tensor(v) for k, v in batch.items()}

    trainer = Trainer(model=model, args=args, train_dataset=train_ds, data_collator=collate)
    print("[train] starting…")
    trainer.train()

    # ── Merge LoRA into bf16 base and push ───────────────────────────────
    print("[merge] reloading base in bf16 and merging adapter…")
    model.save_pretrained("/tmp/adapter")
    del model, trainer
    torch.cuda.empty_cache()

    from peft import PeftModel
    base = AutoModelForCausalLM.from_pretrained(
        BASE_MODEL, torch_dtype=torch.bfloat16, device_map="auto",
        trust_remote_code=True, token=token,
    )
    merged = PeftModel.from_pretrained(base, "/tmp/adapter")
    merged = merged.merge_and_unload()

    print(f"[push] uploading to {hub_repo}…")
    merged.push_to_hub(hub_repo, token=token, private=False)
    tokenizer.push_to_hub(hub_repo, token=token)
    api.upload_file(
        path_or_fileobj=MODEL_CARD.replace(HUB_REPO, hub_repo).encode(),
        path_in_repo="README.md",
        repo_id=hub_repo,
    )
    print(f"[done] https://huggingface.co/{hub_repo}")


@app.local_entrypoint()
def main():
    train.remote()