"""Collect the evidence for a fine-tune that scored worse than the model it started from.

Purpose: the diagnostic half of Part 13's challenge. Reads the run record from the lab notebook,
         the adapter configuration, the chat template the training tokeniser carries and the one
         the serving side carries, the merged model's precision, and the overlap between the
         training data and the evaluation set, then writes one report with a verdict line per
         fault. It changes nothing and needs no accelerator.
Platform: all (standard library only; it reads files and writes a report)
Minimum memory: 8 GB, and far less in practice
Assumes: Python 3.10 or newer. Whatever exists is read; whatever is missing is reported as
         missing rather than guessed, because "not recorded" is itself evidence.

Usage: python3 diagnose-fine-tune.py --adapter runs/format-qwen3-1.7b \
           --labbook labbook.md --train data/train.jsonl --tasks my-tasks.json \
           --serving-tokenizer models/format-qwen3-1.7b-merged --report diagnosis.md

       python3 diagnose-fine-tune.py --adapter runs/domain-qwen3-8b --labbook labbook.md
       # the minimum: the run record and the adapter, with everything else reported as absent

The four faults this collects evidence for, in the order the challenge page works through them:
  1. chat template mismatch between training and serving
  2. learning rate and overfitting
  3. evaluation contaminated by training data
  4. adapter merged into the wrong base or the wrong precision
"""
from __future__ import annotations

import argparse
import hashlib
import json
import re
from pathlib import Path
from typing import Any

WORD_RE = re.compile(r"[a-z0-9]+")
UNKNOWN = "not available"


def sha256_text(text: str) -> str:
    return hashlib.sha256(text.encode("utf-8")).hexdigest()


def normalise(text: str) -> list[str]:
    return WORD_RE.findall(text.lower())


def ngrams(words: list[str], n: int) -> set[tuple[str, ...]]:
    if len(words) < n:
        return {tuple(words)} if words else set()
    return {tuple(words[i:i + n]) for i in range(len(words) - n + 1)}


def containment(a: set[tuple[str, ...]], b: set[tuple[str, ...]]) -> float:
    return len(a & b) / len(a) if a else 0.0


def read_json(path: Path) -> dict[str, Any] | None:
    try:
        return json.loads(path.read_text(encoding="utf-8"))
    except (OSError, json.JSONDecodeError):
        return None


def find_run_record(labbook: Path, adapter: Path, run_id: str | None) -> dict[str, Any] | None:
    """The last training record for this adapter, or the one named by --run-id.

    The run log is JSON lines mixed into a Markdown file, so anything that does not parse is
    prose and is skipped.
    """
    if not labbook.is_file():
        return None
    matches = []
    for line in labbook.read_text(encoding="utf-8").splitlines():
        line = line.strip()
        if not line.startswith("{"):
            continue
        try:
            record = json.loads(line)
        except json.JSONDecodeError:
            continue
        if run_id:
            if record.get("run_id") == run_id:
                return record
            continue
        if "hyperparameters" not in record:
            continue
        output_dir = str(record.get("hyperparameters", {}).get("output_dir", ""))
        if output_dir and Path(output_dir).name == adapter.name:
            matches.append(record)
        elif not output_dir and str(record.get("lab", "")).startswith("part-13/train"):
            matches.append(record)
    return matches[-1] if matches else None


def chat_template_of(source: Path) -> tuple[str | None, str]:
    """The chat template a tokeniser directory carries, and where it came from.

    Newer tokenisers keep it in chat_template.jinja; older ones keep it inside
    tokenizer_config.json. Both are checked, because a directory saved by one version and read
    by another is exactly the situation this script exists to diagnose.
    """
    if source.is_file():
        return source.read_text(encoding="utf-8"), str(source)
    jinja = source / "chat_template.jinja"
    if jinja.is_file():
        return jinja.read_text(encoding="utf-8"), str(jinja)
    config = read_json(source / "tokenizer_config.json")
    if config and isinstance(config.get("chat_template"), str):
        return config["chat_template"], str(source / "tokenizer_config.json")
    return None, str(source)


def template_shape(template: str) -> dict[str, Any]:
    """A few properties that differ between templates, so a difference can be described."""
    special = sorted(set(re.findall(r"<\|[a-zA-Z0-9_]+\|>", template)))
    return {
        "sha256": sha256_text(template),
        "characters": len(template),
        "special_tokens": special[:12],
        "mentions_system": "system" in template,
        "mentions_generation_keyword": "generation" in template,
    }


def overlap_report(train_path: Path, tasks_path: Path, n: int, threshold: float) -> dict[str, Any]:
    """Exact and near-duplicate overlap between the training file and the evaluation set."""
    spec = read_json(tasks_path)
    if spec is None:
        return {"status": f"could not read {tasks_path}"}
    tasks = spec.get("tasks") if isinstance(spec, dict) else spec
    if not isinstance(tasks, list):
        return {"status": f"{tasks_path} has no tasks list"}

    task_entries = []
    for index, task in enumerate(tasks):
        text = "\n".join(str(task[k]) for k in ("prompt", "reference") if isinstance(task.get(k), str))
        words = normalise(text)
        task_entries.append({
            "id": str(task.get("id", f"task-{index:03d}")),
            "hash": sha256_text(" ".join(words)),
            "grams": ngrams(words, n),
        })
    task_hashes = {entry["hash"]: entry["id"] for entry in task_entries}

    exact, near, examples = [], [], 0
    try:
        handle = train_path.open("r", encoding="utf-8")
    except OSError:
        return {"status": f"could not read {train_path}"}
    with handle:
        for number, line in enumerate(handle, start=1):
            line = line.strip()
            if not line:
                continue
            try:
                record = json.loads(line)
            except json.JSONDecodeError:
                continue
            examples += 1
            parts: list[str] = []
            for key in ("prompt", "completion", "messages"):
                value = record.get(key)
                if isinstance(value, str):
                    parts.append(value)
                elif isinstance(value, list):
                    for message in value:
                        if isinstance(message, dict) and isinstance(message.get("content"), str):
                            parts.append(message["content"])
            if isinstance(record.get("text"), str):
                parts.append(record["text"])
            words = normalise("\n".join(parts))
            if not words:
                continue
            digest = sha256_text(" ".join(words))
            if digest in task_hashes:
                exact.append({"line": number, "task": task_hashes[digest]})
                continue
            grams = ngrams(words, n)
            best_id, best = None, 0.0
            for entry in task_entries:
                score = max(containment(grams, entry["grams"]), containment(entry["grams"], grams))
                if score > best:
                    best_id, best = entry["id"], score
            if best >= threshold:
                near.append({"line": number, "task": best_id, "containment": round(best, 3)})
    return {"status": "checked", "training_examples": examples, "evaluation_tasks": len(task_entries),
            "exact_matches": exact, "near_duplicates": near, "n": n, "threshold": threshold}


def section(title: str) -> str:
    return f"\n## {title}\n\n"


def main() -> None:
    parser = argparse.ArgumentParser(description=__doc__.splitlines()[0])
    parser.add_argument("--adapter", required=True, help="the adapter directory a run saved")
    parser.add_argument("--labbook", default="labbook.md")
    parser.add_argument("--run-id", default=None, help="a specific run record instead of the last")
    parser.add_argument("--train", default=None, help="the training file, for the overlap check")
    parser.add_argument("--tasks", default=None, help="the evaluation task file it is checked against")
    parser.add_argument("--serving-tokenizer", default=None,
                        help="directory or .jinja file whose chat template the server is using; "
                             "usually the merged model directory")
    parser.add_argument("--merged", default=None,
                        help="the merged model directory, for the precision and base checks")
    parser.add_argument("--n", type=int, default=13)
    parser.add_argument("--threshold", type=float, default=0.5)
    parser.add_argument("--report", default="diagnosis.md")
    args = parser.parse_args()

    adapter = Path(args.adapter)
    lines = [f"# Fine-tune diagnosis: {adapter}\n"]
    verdicts: list[str] = []

    # ---------------------------------------------------------------- the run record
    lines.append(section("1. The run"))
    record = find_run_record(Path(args.labbook), adapter, args.run_id)
    if record is None:
        lines.append(f"No run record found in `{args.labbook}` for this adapter. "
                     "Without it, the hyperparameters and the losses are unknown and every "
                     "conclusion below is weaker.\n")
        verdicts.append("RUN RECORD: missing. Record every run; a result you cannot reproduce "
                        "is not a result.")
    else:
        hyper = record.get("hyperparameters", {})
        losses = record.get("losses", {})
        lines.append("| Field | Value |\n| --- | --- |\n")
        for key in ("run_id", "date", "model", "seed", "config_commit"):
            lines.append(f"| {key} | `{record.get(key, UNKNOWN)}` |\n")
        for key in ("method", "rank", "alpha", "target_modules", "learning_rate", "epochs",
                    "effective_batch", "max_length", "precision", "compute_precision",
                    "load_in_4bit", "completion_only_loss"):
            if key in hyper:
                lines.append(f"| {key} | `{hyper[key]}` |\n")
        dataset = record.get("dataset", {})
        lines.append(f"| dataset | `{dataset.get('path', UNKNOWN)}` |\n")
        lines.append(f"| dataset sha256 | `{dataset.get('sha256', UNKNOWN)}` |\n")
        lines.append(f"| train / validation examples | "
                     f"{dataset.get('train_examples', '?')} / "
                     f"{dataset.get('validation_examples', '?')} |\n")
        lines.append("\n### Losses\n\n```json\n" + json.dumps(losses, indent=2) + "\n```\n")

        # ------------------------------------------------- fault 2: rate and overfitting
        best_epoch = losses.get("best_epoch")
        first, final = losses.get("first_train_loss"), losses.get("final_train_loss")
        final_eval, best_eval = losses.get("final_eval_loss"), losses.get("best_eval_loss")
        if best_epoch is not None and best_epoch <= 1:
            verdicts.append("OVERFITTING: the best epoch was the first, so later epochs made the "
                            "evaluation loss worse. Fewer epochs, a lower learning rate, a lower "
                            "rank or more data, in that order of cost.")
        if final_eval is not None and best_eval is not None and final_eval > best_eval * 1.05:
            verdicts.append("OVERFITTING: the final evaluation loss is meaningfully above the best "
                            "one, so the curves diverged. Confirm load_best_model_at_end was set, "
                            "or the saved adapter is not the best checkpoint.")
        if first is not None and final is not None and abs(first - final) < 0.01:
            verdicts.append("LEARNING RATE: the training loss barely moved. Either the rate is too "
                            "low or the adapter attached to nothing; check the trainable-parameter "
                            "count from the run.")
        rate = hyper.get("learning_rate")
        if isinstance(rate, (int, float)) and rate >= 1e-3:
            verdicts.append(f"LEARNING RATE: {rate} is high for an adapter run, where roughly 1e-4 "
                            "is the documented starting point. A loss that spiked or went to nan "
                            "points here first.")

    # ---------------------------------------------------------------- the adapter
    lines.append(section("2. The adapter"))
    adapter_config = read_json(adapter / "adapter_config.json")
    recorded_base = None
    if adapter_config is None:
        lines.append(f"`{adapter / 'adapter_config.json'}` is missing, so this is not a PEFT "
                     "adapter directory.\n")
        verdicts.append("ADAPTER: no adapter_config.json. Point --adapter at what the run saved.")
    else:
        recorded_base = adapter_config.get("base_model_name_or_path")
        lines.append("| Field | Value |\n| --- | --- |\n")
        for key in ("base_model_name_or_path", "peft_type", "task_type", "r", "lora_alpha",
                    "lora_dropout", "target_modules", "use_dora", "use_rslora"):
            if key in adapter_config:
                lines.append(f"| {key} | `{adapter_config[key]}` |\n")
        if record is not None and record.get("model") and recorded_base:
            if str(record["model"]) != str(recorded_base):
                verdicts.append(f"WRONG BASE: the run record says the base was "
                                f"`{record['model']}` and the adapter records "
                                f"`{recorded_base}`. One of them is not what you think.")

    # ---------------------------------------------------------------- chat templates
    lines.append(section("3. Chat template, training against serving"))
    train_template, train_source = chat_template_of(adapter)
    serve_template, serve_source = (None, UNKNOWN)
    if args.serving_tokenizer:
        serve_template, serve_source = chat_template_of(Path(args.serving_tokenizer))

    if train_template is None:
        lines.append(f"No chat template found under `{adapter}`. The adapter was saved without a "
                     "tokeniser, so the template used in training cannot be recovered from the "
                     "artefact.\n")
        verdicts.append("CHAT TEMPLATE: the adapter carries no tokeniser, so the training-time "
                        "template is unknown. Save the tokeniser with every run.")
    else:
        lines.append(f"Training-side template from `{train_source}`:\n\n```json\n"
                     + json.dumps(template_shape(train_template), indent=2) + "\n```\n")
    if serve_template is None and args.serving_tokenizer:
        lines.append(f"No chat template found under `{serve_source}`.\n")
        verdicts.append("CHAT TEMPLATE: the serving side has no template file where one was "
                        "expected, so the server is using whatever the model file carries.")
    elif serve_template is not None:
        lines.append(f"Serving-side template from `{serve_source}`:\n\n```json\n"
                     + json.dumps(template_shape(serve_template), indent=2) + "\n```\n")

    if train_template and serve_template:
        if sha256_text(train_template) == sha256_text(serve_template):
            lines.append("\n**The two templates are byte-identical.** This fault is eliminated.\n")
        else:
            lines.append("\n**The two templates differ.** Every training example was formatted "
                         "with one and every request is formatted with the other.\n")
            verdicts.append("CHAT TEMPLATE MISMATCH: the training-side and serving-side templates "
                            "have different hashes. This is the first fault to fix, and it usually "
                            "explains a fine-tune that got worse at everything at once.")
    elif not args.serving_tokenizer:
        lines.append("\nNo `--serving-tokenizer` given, so the comparison was not made. Point it "
                     "at the merged model directory the server is loading.\n")

    # ---------------------------------------------------------------- precision and merge
    lines.append(section("4. Precision and the merge"))
    if args.merged:
        merged_config = read_json(Path(args.merged) / "config.json")
        if merged_config is None:
            lines.append(f"`{Path(args.merged) / 'config.json'}` is missing.\n")
        else:
            dtype = merged_config.get("dtype") or merged_config.get("torch_dtype") or UNKNOWN
            quant = merged_config.get("quantization_config")
            lines.append(f"| Field | Value |\n| --- | --- |\n| merged dtype | `{dtype}` |\n")
            lines.append(f"| merged base architecture | `{merged_config.get('model_type', UNKNOWN)}` |\n")
            lines.append(f"| quantization_config present | `{bool(quant)}` |\n")
            if quant:
                verdicts.append("WRONG PRECISION: the merged model still carries a "
                                "quantization_config, so the adapter was merged into a quantised "
                                "base. Merge into the base at bfloat16 and quantise afterwards.")
            if isinstance(dtype, str) and dtype.lower() in {"float16", "torch.float16"}:
                verdicts.append("PRECISION: the merged model is float16 while adapters here are "
                                "trained in bfloat16. The two have different ranges; convert at "
                                "bfloat16 unless you have a reason not to.")
    else:
        lines.append("No `--merged` directory given, so the precision of the exported model was "
                     "not checked.\n")

    if record is not None and record.get("hyperparameters", {}).get("load_in_4bit") and args.merged:
        lines.append("\nThis run trained against a 4-bit base. The merge must still be into a "
                     "bfloat16 copy of that base, never into the quantised one.\n")

    # ---------------------------------------------------------------- contamination
    lines.append(section("5. Overlap between training data and the evaluation set"))
    if args.train and args.tasks:
        overlap = overlap_report(Path(args.train), Path(args.tasks), args.n, args.threshold)
        lines.append("```json\n" + json.dumps(
            {k: v for k, v in overlap.items() if k != "status"} | {"status": overlap["status"]},
            indent=2, default=str) + "\n```\n")
        if overlap.get("exact_matches"):
            verdicts.append(f"CONTAMINATED EVALUATION: {len(overlap['exact_matches'])} training "
                            "example(s) are byte-identical to evaluation tasks after "
                            "normalisation. Any score measured on this set is not a measurement.")
        if overlap.get("near_duplicates"):
            verdicts.append(f"POSSIBLE CONTAMINATION: {len(overlap['near_duplicates'])} "
                            "near-duplicate pair(s) above the threshold. Read them; some will be "
                            "shared stock phrases and some will be paraphrased evaluation tasks.")
        if overlap.get("status") == "checked" and not overlap.get("exact_matches") \
                and not overlap.get("near_duplicates"):
            lines.append("\n**No overlap found at this n-gram size and threshold.** This fault is "
                         "eliminated for this pair of files.\n")
    else:
        lines.append("No `--train` and `--tasks` given, so the overlap check was not run. It is "
                     "the cheapest of the four checks and the one that invalidates the others.\n")

    # ---------------------------------------------------------------- settings
    lines.append(section("6. Sampling settings the evaluation used"))
    if args.tasks:
        spec = read_json(Path(args.tasks))
        settings = (spec or {}).get("settings") if isinstance(spec, dict) else None
        if settings:
            lines.append("```json\n" + json.dumps(settings, indent=2) + "\n```\n")
            lines.append("\nBoth models must be scored at these settings, at the same "
                         "quantisation, with the same system prompt. A comparison across a change "
                         "in any of them is measuring more than the fine-tune.\n")
        else:
            lines.append(f"`{args.tasks}` has no settings block, so the run used whatever the "
                         "harness defaults are. Record them.\n")
    else:
        lines.append("No `--tasks` given.\n")

    # ---------------------------------------------------------------- verdicts
    lines.append(section("Verdict"))
    if verdicts:
        for item in verdicts:
            lines.append(f"- **{item}**\n")
        lines.append("\nFix one thing, then measure again with the same command. Two changes at "
                     "once means you will not know which one worked.\n")
    else:
        lines.append("Nothing in the collected evidence points at one of the four faults. That is "
                     "a real result: the remaining candidates are that the fine-tune genuinely did "
                     "not help on these tasks, that the change is inside the run-to-run noise, or "
                     "that the fault is somewhere this script does not look. Run the base model "
                     "twice to establish the noise floor before concluding anything.\n")

    report = "".join(lines)
    Path(args.report).write_text(report, encoding="utf-8")
    print(report)
    print(f"\nwritten to {args.report}")


if __name__ == "__main__":
    main()
