"""DELLA magnitude pruning + task-vector combine over voice head tensors.

Clean-room port of the DELLA spec (arXiv:2406.11617) in numpy. No mergekit
dependency, no vendored code: mergekit is LGPL-3.0 and the authors' snapshot
carries no license grant, so both served as read-only references only.

One deliberate fork from upstream mergekit's ``della_magprune``, recorded
here rather than collapsed: upstream rescales L1 against the whole-tensor
sum, while ``magprune`` rescales each row to its own pre-prune ``sum|x|``.
Per-row keeps every output neuron's scale (each head row is one token's
direction), stays bounded at aggressive sparsity, and keeps row-blocking
exactly correct (a global rescale would need the full tensor at once).
Equivalence spot-checks against mergekit therefore expect "close", not
bit-exact: same family, different rescale locality.

Determinism: ``seed`` spawns one child stream per voice
(``SeedSequence(seed)``), so a seeded run is bit-identical regardless of
``block_rows`` — each voice's Bernoulli draws are consumed in row order on
its own stream. ``seed=None`` gives fresh entropy per voice
(nondeterministic, as upstream).
"""

VALID_RESCALES = ("l1", "inv_p")


def _validate_density_epsilon(density, epsilon):
    if density + epsilon >= 1 or density - epsilon <= 0:
        raise ValueError(
            "della density±epsilon must stay in (0,1): "
            f"got density={density} epsilon={epsilon}"
        )


def _check_rescale(rescale):
    if rescale not in VALID_RESCALES:
        raise ValueError(f"della rescale must be 'l1' or 'inv_p', got {rescale!r}")


def magprune(delta, density=0.5, epsilon=0.1, rescale="l1", rng=None):
    """Rank-weighted Bernoulli prune of one delta, row-wise. Always F32 out.

    Keep-probability ramps linearly with within-row magnitude rank, from
    ``density-epsilon`` (smallest) to ``density+epsilon`` (largest), then a
    seeded Bernoulli draw decides survivors. ``rescale="l1"`` restores each
    surviving row's ``sum|x|``; ``rescale="inv_p"`` divides each survivor by
    its own keep-prob (paper-faithful, heavy-tailed).
    """
    import numpy as np
    F32 = np.float32
    _check_rescale(rescale)
    a = np.asarray(delta, dtype=F32)
    if density >= 1:
        return a
    if density <= 0:
        return np.zeros_like(a)
    _validate_density_epsilon(density, epsilon)
    orig_shape = a.shape
    w = np.atleast_2d(a)
    order = np.argsort(np.abs(w), axis=-1, kind="stable")
    ranks = np.argsort(order, axis=-1).astype(F32) + F32(1.0)
    rmin = ranks.min(axis=-1, keepdims=True)
    rmax = ranks.max(axis=-1, keepdims=True)
    denom = rmax - rmin
    # Single-column rows have no ordering; fall back to flat density.
    rank_norm = np.divide(ranks - rmin, denom,
                          out=np.full_like(ranks, F32(0.5)),
                          where=denom > 0)
    rank_norm = np.clip(rank_norm, F32(0.0), F32(1.0))
    d = F32(density)
    e = F32(epsilon)
    probs = (d - e) + rank_norm * (F32(2.0) * e)
    if rng is None:
        rng = np.random.default_rng()
    mask = rng.binomial(1, probs).astype(F32)
    masked = w * mask
    if rescale == "l1":
        before = np.abs(w).sum(axis=-1, keepdims=True)
        after = np.abs(masked).sum(axis=-1, keepdims=True)
        ok = (before >= 1e-7) & (after >= 1e-7)
        scale = np.where(ok, before / np.maximum(after, F32(1e-30)), F32(1.0))
        out = masked * scale
    else:  # inv_p
        out = masked / probs
        if not np.all(np.isfinite(out)):
            raise ValueError("della inv_p rescale produced NaN/Inf "
                             "- refusing to write a corrupt voice")
    return np.asarray(out.reshape(orig_shape), dtype=F32)


def combine(base, deltas, weights, lambda_=1.0, normalize=True):
    """TIES-sum combine of pruned deltas onto base. F32 in, F32 out.

    Weight-scale, elect the majority sign per position (ties go positive),
    average the elected (``normalize``) or sum them, scale by ``lambda_``,
    add base. Positions every voice pruned divide by 1, not 0.
    """
    import numpy as np
    F32 = np.float32
    b = np.asarray(base, dtype=F32)
    ds = [np.asarray(d, dtype=F32) for d in deltas]
    if not ds:
        return b
    if len(weights) != len(ds):
        raise ValueError(f"della weights ({len(weights)}) != deltas ({len(ds)})")
    for i, d in enumerate(ds):
        if d.shape != b.shape:
            raise ValueError(f"della delta {i} shape {d.shape} != base {b.shape}")
    wcol = np.asarray(list(weights), dtype=F32).reshape((-1,) + (1,) * b.ndim)
    W = np.stack(ds).astype(F32, copy=False) * wcol
    majority = np.where(W.sum(axis=0) >= 0, F32(1.0), F32(-1.0))
    elected = np.sign(W) == majority
    mixed = np.where(elected, W, F32(0.0)).sum(axis=0)
    divisor = np.where(elected, wcol, F32(0.0)).sum(axis=0)
    divisor = np.where(divisor == 0, F32(1.0), divisor)
    if normalize:
        mixed = mixed / divisor
    if lambda_ != 1:
        mixed = mixed * F32(lambda_)
    return b + mixed


def merge(base, voices, weights=None, density=0.5, epsilon=0.1, lambda_=1.0,
          rescale="l1", normalize=True, seed=None, block_rows=4096):
    """Stream one DELLA merge: per-block (voice-base) prune, then combine.

    Only ``block_rows`` rows per voice are pruned/combined at once, so the
    full ``n_voices x rows x cols`` delta stack is never materialized; peak
    besides the inputs/output is ``O(block_rows x cols x n_voices)``.
    Blocking is exactly correct: pruning ranks within rows and election
    votes per position, neither crosses a block boundary. Each voice draws
    on its own ``SeedSequence(seed)`` child stream, so repartitioning blocks
    never changes the result. Output is cast back to the base input's dtype.
    """
    import numpy as np
    F32 = np.float32
    _check_rescale(rescale)
    b = np.asarray(base)
    out_dtype = b.dtype
    bw = np.asarray(b, dtype=F32)
    vs = [np.asarray(v) for v in voices]
    if not vs:
        return np.asarray(b, dtype=out_dtype)
    wts = [1.0] * len(vs) if weights is None else list(weights)
    if len(wts) != len(vs):
        raise ValueError(f"della weights ({len(wts)}) != voices ({len(vs)})")
    for i, v in enumerate(vs):
        if v.shape != bw.shape:
            raise ValueError(f"della voice {i} shape {v.shape} != base {bw.shape}")
    if not (density >= 1 or density <= 0):
        _validate_density_epsilon(density, epsilon)
    squeeze = False
    if bw.ndim == 1:
        squeeze = True
        bw = bw[None, :]
        vs = [v[None, :] for v in vs]
    if bw.ndim < 1:
        raise ValueError("della merge needs at least 1-D tensors")
    if seed is None:
        rngs = [np.random.default_rng() for _ in vs]
    else:
        rngs = [np.random.default_rng(s)
                for s in np.random.SeedSequence(seed).spawn(len(vs))]
    vw = [np.asarray(v, dtype=F32) for v in vs]
    out = np.empty_like(bw)
    for start in range(0, bw.shape[0], block_rows):
        sl = slice(start, min(start + block_rows, bw.shape[0]))
        bb = bw[sl]
        pruned = [magprune(v[sl] - bb, density=density, epsilon=epsilon,
                           rescale=rescale, rng=rngs[vi])
                  for vi, v in enumerate(vw)]
        out[sl] = combine(bb, pruned, wts, lambda_=lambda_, normalize=normalize)
    if squeeze:
        out = out.reshape(b.shape)
    if out_dtype != F32:
        return out.astype(out_dtype, copy=False)
    return out