# -*- coding: utf-8 -*-
"""Sparse Model Nexus — the operator algebra over LoRA quanta.

This is the executable form of the Pharosiraptor graph
`1c0cc1b9-d95a-4448-b7d7-a08e1e347820`. A "quantum" is a LoRA adapter over a
shared frozen base (distilgpt2). The four operators from the graph's
`Operator::` nodes are implemented here:

  ⊕  mixture(width)      -> route/blend several quanta at once (weighted combine)
                            assoc + commut; identity = null-expert
  ∘  sequential(depth)   -> chain quanta: output of one feeds the next
                            assoc, non-commut
  ▷  route(conditional)  -> gated dispatch: pick a quantum by a predicate
  ⊗  merge(weight-space) -> task-arithmetic merge into a single new quantum

Every composition is itself deployable, which is the whole point of the
sparse-limit-model continuum: the same algebra spans one quantum to a dense mix.
"""
import os, re
import torch
from transformers import AutoModelForCausalLM, AutoTokenizer
from peft import PeftModel

BASE = "distilgpt2"
ADAPTERS_DIR = os.path.join(os.path.dirname(__file__), "adapters")
QUANTA = ["shared_common", "code_specialist", "biomed_specialist"]


class Nexus:
    """Loads the frozen base once and hot-swaps / combines LoRA quanta on top."""

    def __init__(self, adapters_dir=ADAPTERS_DIR, quanta=QUANTA):
        self.tok = AutoTokenizer.from_pretrained(BASE)
        self.tok.pad_token = self.tok.eos_token
        base = AutoModelForCausalLM.from_pretrained(BASE)
        first, *rest = quanta
        # Load the first quantum as the PeftModel spine, then attach the others.
        self.model = PeftModel.from_pretrained(base, os.path.join(adapters_dir, first),
                                               adapter_name=first)
        for q in rest:
            self.model.load_adapter(os.path.join(adapters_dir, q), adapter_name=q)
        self.model.eval()
        self.quanta = list(quanta)
        self._derived = []

    # -- generation primitive --------------------------------------------------
    @torch.no_grad()
    def _gen(self, prompt, max_new_tokens=48):
        ids = self.tok(prompt, return_tensors="pt")
        out = self.model.generate(**ids, max_new_tokens=max_new_tokens, do_sample=False,
                                  temperature=1.0, pad_token_id=self.tok.eos_token_id,
                                  repetition_penalty=1.3)
        return self.tok.decode(out[0][ids["input_ids"].shape[1]:], skip_special_tokens=True).strip()

    def activate(self, name):
        self.model.set_adapter(name)
        return name

    @staticmethod
    def _safe(s):
        # PEFT adapter names become nn.Module keys; no "." allowed.
        return re.sub(r"[^0-9a-zA-Z_]+", "_", s).strip("_")

    # -- ⊕ mixture(width): weighted blend of several quanta (MoE-style) ---------
    def mixture(self, prompt, weights, name="mix", max_new_tokens=48):
        active = [q for q, w in weights.items() if w != 0]           # null-expert = identity
        coeffs = [weights[q] for q in active]
        combo = self._safe(f"{name}__" + "_".join(f"{w:g}x{q}" for q, w in weights.items() if w))
        if combo not in self._derived:
            self.model.add_weighted_adapter(active, coeffs, combo,
                                            combination_type="linear")
            self._derived.append(combo)
        self.activate(combo)
        return {"operator": "⊕ mixture", "weights": weights, "output": self._gen(prompt, max_new_tokens)}

    # -- ∘ sequential(depth): chain quanta, threading text through -------------
    def sequential(self, prompt, chain, max_new_tokens=40):
        steps, text = [], prompt
        for q in chain:                                              # non-commutative on purpose
            self.activate(q)
            gen = self._gen(text, max_new_tokens)
            steps.append({"quantum": q, "in": text[-80:], "out": gen})
            text = (text + " " + gen).strip()
        return {"operator": "∘ sequential", "chain": chain, "steps": steps, "output": text}

    # -- ▷ route(conditional): gated dispatch by predicate ---------------------
    def route(self, prompt, max_new_tokens=48):
        chosen = self._predicate(prompt)
        self.activate(chosen)
        return {"operator": "▷ route", "predicate_choice": chosen,
                "output": self._gen(prompt, max_new_tokens)}

    @staticmethod
    def _predicate(prompt):
        p = prompt.lower()
        if re.search(r"\b(def|class|function|sort|array|sql|select|import|async|api)\b", p):
            return "code_specialist"
        if re.search(r"\b(cell|protein|gene|patient|clinical|dna|enzyme|disease|insulin|trial)\b", p):
            return "biomed_specialist"
        return "shared_common"

    # -- ⊗ merge(weight-space): task-arithmetic merge into ONE new quantum -----
    def merge(self, sources, name="merged", max_new_tokens=48, prompt=None):
        coeffs = [1.0 / len(sources)] * len(sources)
        if name not in self._derived:
            self.model.add_weighted_adapter(sources, coeffs, name, combination_type="linear")
            self._derived.append(name)
        self.activate(name)
        result = {"operator": "⊗ merge", "sources": sources, "new_quantum": name}
        if prompt is not None:
            result["output"] = self._gen(prompt, max_new_tokens)
        return result
