# -*- coding: utf-8 -*-
"""Train three LoRA "quanta" for the Sparse Model Nexus.

Each quantum is a minimal, self-contained, composable unit (a LoRA adapter over
a shared frozen base) — the operational realization of the graph's
`Definition::SparseLimitModel`. Base is distilgpt2 (CPU-trainable in minutes).

Seed quanta mirror the graph nodes:
  SLM-quantum: code-specialist   -> code_specialist
  SLM-quantum: biomed-specialist -> biomed_specialist
  SLM-quantum: shared-common     -> shared_common
"""
import os, sys, json, time
os.environ.setdefault("HF_HUB_DISABLE_TELEMETRY", "1")
os.environ.setdefault("TOKENIZERS_PARALLELISM", "false")

import torch
from transformers import AutoModelForCausalLM, AutoTokenizer
from peft import LoraConfig, get_peft_model

BASE = "distilgpt2"
OUT = os.path.join(os.path.dirname(__file__), "adapters")
os.makedirs(OUT, exist_ok=True)
torch.manual_seed(42)

# Tiny domain corpora — enough for a LoRA to measurably specialize on CPU.
CORPORA = {
    "code_specialist": [
        "def binary_search(arr, target):\n    lo, hi = 0, len(arr) - 1\n    while lo <= hi:\n        mid = (lo + hi) // 2\n        if arr[mid] == target: return mid\n        if arr[mid] < target: lo = mid + 1\n        else: hi = mid - 1\n    return -1",
        "class Stack:\n    def __init__(self): self._items = []\n    def push(self, x): self._items.append(x)\n    def pop(self): return self._items.pop()\n    def peek(self): return self._items[-1]\n    def is_empty(self): return not self._items",
        "async def fetch_json(url):\n    async with aiohttp.ClientSession() as session:\n        async with session.get(url) as resp:\n            resp.raise_for_status()\n            return await resp.json()",
        "import functools\n@functools.lru_cache(maxsize=None)\ndef fib(n):\n    if n < 2: return n\n    return fib(n - 1) + fib(n - 2)",
        "def quicksort(xs):\n    if len(xs) <= 1: return xs\n    pivot = xs[len(xs) // 2]\n    left = [x for x in xs if x < pivot]\n    mid = [x for x in xs if x == pivot]\n    right = [x for x in xs if x > pivot]\n    return quicksort(left) + mid + quicksort(right)",
        "SELECT user_id, COUNT(*) AS orders\nFROM orders\nWHERE created_at >= '2026-01-01'\nGROUP BY user_id\nHAVING COUNT(*) > 5\nORDER BY orders DESC;",
        "def merge_sorted(a, b):\n    result, i, j = [], 0, 0\n    while i < len(a) and j < len(b):\n        if a[i] <= b[j]: result.append(a[i]); i += 1\n        else: result.append(b[j]); j += 1\n    return result + a[i:] + b[j:]",
        "try:\n    value = int(user_input)\nexcept ValueError:\n    raise SystemExit('expected an integer')\nfinally:\n    logging.info('parse attempt complete')",
    ],
    "biomed_specialist": [
        "The mitochondrion is the primary site of oxidative phosphorylation, generating ATP through the electron transport chain across the inner membrane.",
        "CRISPR-Cas9 introduces a double-strand break at a target locus specified by a guide RNA, enabling precise genome editing via non-homologous end joining or homology-directed repair.",
        "In a randomized controlled trial, participants are allocated to treatment and control arms to minimize confounding and estimate the causal effect of an intervention.",
        "Hemoglobin binds oxygen cooperatively; the sigmoidal dissociation curve reflects allosteric transitions between the tense and relaxed quaternary states.",
        "Antibiotic resistance spreads through horizontal gene transfer, including conjugation, transformation, and transduction of plasmid-borne resistance genes.",
        "The blood-brain barrier restricts paracellular diffusion via tight junctions between endothelial cells, limiting delivery of many therapeutic compounds to the central nervous system.",
        "Insulin secreted by pancreatic beta cells promotes glucose uptake in peripheral tissues by stimulating translocation of GLUT4 transporters to the plasma membrane.",
        "Apoptosis is a regulated form of programmed cell death mediated by caspase activation, distinct from the inflammatory cascade of necrosis.",
    ],
    "shared_common": [
        "The quick brown fox jumps over the lazy dog near the riverbank at dawn.",
        "In summary, the results indicate a consistent pattern across all measured conditions.",
        "First, we outline the background. Next, we describe the method. Finally, we present conclusions.",
        "This document provides an overview of the system and explains how each component fits together.",
        "Please review the attached material and share any feedback before the meeting on Thursday.",
        "The weather today is mild, with a gentle breeze and scattered clouds over the coast.",
        "Thank you for your patience; the requested information is included in the section below.",
        "A clear explanation helps the reader follow the argument from premise to conclusion.",
    ],
}


def train_one(name, texts, epochs=8, lr=2e-3, rank=8):
    tok = AutoTokenizer.from_pretrained(BASE)
    tok.pad_token = tok.eos_token
    base = AutoModelForCausalLM.from_pretrained(BASE)
    cfg = LoraConfig(r=rank, lora_alpha=16, lora_dropout=0.05,
                     target_modules=["c_attn"], bias="none", task_type="CAUSAL_LM")
    model = get_peft_model(base, cfg)
    model.train()
    opt = torch.optim.AdamW((p for p in model.parameters() if p.requires_grad), lr=lr)

    batch = tok(texts, return_tensors="pt", padding=True, truncation=True, max_length=96)
    labels = batch["input_ids"].clone()
    labels[batch["attention_mask"] == 0] = -100

    first = last = None
    for ep in range(epochs):
        opt.zero_grad()
        out = model(input_ids=batch["input_ids"], attention_mask=batch["attention_mask"], labels=labels)
        out.loss.backward()
        opt.step()
        if first is None:
            first = out.loss.item()
        last = out.loss.item()
        print(f"  [{name}] epoch {ep+1}/{epochs} loss={last:.4f}", flush=True)

    dest = os.path.join(OUT, name)
    model.save_pretrained(dest)
    tok.save_pretrained(dest)
    trainable = sum(p.numel() for p in model.parameters() if p.requires_grad)
    return {"name": name, "loss_start": round(first, 4), "loss_end": round(last, 4),
            "trainable_params": trainable, "path": dest}


def main():
    report = {"base": BASE, "quanta": [], "started": time.strftime("%Y-%m-%dT%H:%M:%SZ", time.gmtime())}
    for name, texts in CORPORA.items():
        print(f"Training quantum: {name}", flush=True)
        report["quanta"].append(train_one(name, texts))
    report["finished"] = time.strftime("%Y-%m-%dT%H:%M:%SZ", time.gmtime())
    with open(os.path.join(OUT, "train_report.json"), "w", encoding="utf-8") as f:
        json.dump(report, f, indent=2)
    print(json.dumps(report, indent=2))
    # every quantum must have actually learned something
    assert all(q["loss_end"] < q["loss_start"] for q in report["quanta"]), "a quantum failed to converge"


if __name__ == "__main__":
    main()
