"""Rebuild a merge from its parents' control vectors: apply a mergekit recipe to the vectors, layer by layer,
and bake the result into one GGUF.

A linear-family merge on stock moves the model by sum_i w_i(L) * delta_i in expectation, so the merge's vector
is the same weighted sum of the parents' vectors at each layer:

    direction.L = sum_i w_i(L) * parent_i.direction.L

w_i(L) is resolved as mergekit resolves it (mergekit/config.py evaluate_setting and mergekit/plan.py): a list
is a gradient interpolated at t = L / 59, conditional values take the first matching filter, then each
method's normalize default and lambda. The parent files are released GGUFs, which already carry their fitted
strength, so the baked file loads at 1.0 like any of them.

    python -m ecce_vector.bake recipe.yaml merge.gguf \\
        --parent Org/Parent-A=vectors/parent-a.gguf --parent Org/Parent-B=vectors/parent-b.gguf

Per-layer weights are what a single llama.cpp --control-vector-scaled weight per file cannot express; a
recipe with one constant weight per parent is just a launch-time blend and needs no baking.

Handled: linear, task_arithmetic, della_linear, dare_linear, and ties / dare_ties / della / breadcrumbs(_ties)
without their sign-election step (a weighted sum cannot reproduce it: a warning, and untested). Refused: other
methods (slerp, model_stock, ...), a base other than stock gemma-4-31B-it, and a recipe that weights a layer's
projections differently (one vector per layer cannot carry that).
"""
from __future__ import annotations

import argparse

import numpy as np
import yaml

from .io import N_LAYERS, read_gguf, write_gguf

STOCK = "gemma-4-31b-it"
LINEAR = {"linear", "task_arithmetic", "ties", "dare_ties", "dare_linear", "della", "della_linear", "breadcrumbs",
          "breadcrumbs_ties"}
# mergekit/merge_methods/registry.py: each method's default for `normalize` (linear.py defaults it to true)
NORMALIZE_DEFAULT = {"linear": True, "task_arithmetic": False, "ties": True, "dare_ties": False, "dare_linear": False,
                     "breadcrumbs": False, "breadcrumbs_ties": False, "della": True, "della_linear": False}
SIGN_ELECTION = {"ties", "dare_ties", "breadcrumbs_ties", "della"}
# one representative tensor per projection family (v_proj is absent on the full-attention layers)
TENSORS = ("self_attn.q_proj.weight", "self_attn.k_proj.weight", "self_attn.o_proj.weight", "mlp.gate_proj.weight",
           "mlp.up_proj.weight", "mlp.down_proj.weight")
PREFIX = "model.language_model.layers."


def evaluate(name: str, setting, t: float):
    """mergekit's evaluate_setting, numeric settings only."""
    if isinstance(setting, (int, float)):
        return float(setting)
    if isinstance(setting, list):
        if all(isinstance(e, (int, float)) for e in setting):
            scaled = t * (len(setting) - 1)
            i0 = int(scaled)
            i1 = min(len(setting) - 1, i0 + 1)
            return (1 - (scaled - i0)) * setting[i0] + (scaled - i0) * setting[i1]
        for cond in setting:
            f = cond.get("filter")
            if f is None or f == "*" or f in name:
                return evaluate(name, cond["value"], t)
        return None
    raise ValueError(f"unsupported setting {setting!r}")


def layer_weights(cfg: dict) -> tuple[dict[str, list[float]], list[tuple[str, int]]]:
    """Per parent, its weight at each layer; and the (parent, layer) pairs whose projections disagree."""
    method = cfg["merge_method"]
    if method not in LINEAR:
        raise SystemExit(f"{method}: not a linear-family merge; a weighted sum of vectors cannot reproduce it")
    params = cfg.get("parameters") or {}
    lam = float(params.get("lambda", 1.0))
    normalize = params.get("normalize", NORMALIZE_DEFAULT[method])
    out, mixed = {}, []
    for m in cfg["models"]:
        if m["model"] == cfg.get("base_model"):
            continue            # the base contributes no delta and is not in the normalizing sum
        weights = []
        for L in range(N_LAYERS):
            vals = [evaluate(f"{PREFIX}{L}.{tn}", (m.get("parameters") or {}).get("weight", 1.0), L / (N_LAYERS - 1))
                    for tn in TENSORS]
            if len(set(round(v, 9) for v in vals)) > 1:
                mixed.append((m["model"], L))
            weights.append(vals[-1])
        out[m["model"]] = weights
    if normalize:
        for L in range(N_LAYERS):
            s = sum(out[k][L] for k in out)
            for k in out:
                out[k][L] = out[k][L] / s if s else 0.0
    return {k: [lam * w for w in v] for k, v in out.items()}, mixed


def main(argv=None):
    ap = argparse.ArgumentParser(description=__doc__.split("\n\n")[0])
    ap.add_argument("recipe", help="the merge's mergekit YAML")
    ap.add_argument("out", help="output .gguf")
    ap.add_argument("--parent", action="append", required=True,
                    help="MODEL=path.gguf, MODEL exactly as the recipe names it")
    ap.add_argument("--name", help="general.name metadata")
    a = ap.parse_args(argv)
    cfg = yaml.safe_load(open(a.recipe))
    base = str(cfg.get("base_model", "")).rstrip("/").split("/")[-1].lower()
    if base != STOCK:
        raise SystemExit(f"base_model {cfg.get('base_model')!r}: parent vectors are deltas from stock gemma-4-31B-it")
    if cfg["merge_method"] in SIGN_ELECTION:
        print(f"warning: {cfg['merge_method']} elects signs per element; the weighted sum ignores that (untested)")
    w, mixed = layer_weights(cfg)
    if mixed:
        raise SystemExit(f"projections weighted differently within layers {sorted({L for _, L in mixed})}")
    files = dict(p.split("=", 1) for p in a.parent)
    if set(files) != set(w):
        raise SystemExit(f"recipe models {sorted(w)} vs --parent {sorted(files)}")
    V = np.zeros_like(read_gguf(next(iter(files.values()))))
    for m, f in files.items():
        V += np.asarray(w[m], np.float32)[:, None] * read_gguf(f)
    write_gguf(a.out, V, name=a.name)
    for m in w:
        print(f"{m:50s} weight at L1 {w[m][1]:.3f}  L30 {w[m][30]:.3f}  L59 {w[m][59]:.3f}")
    print(f"wrote {a.out}: {cfg['merge_method']}, {len(w)} parents; load it at 1.0")


if __name__ == "__main__":
    main()