"""HPC spectral objective: folded error pyramid, DC/vesica penalty, shaping.

The C engine scores a block with::

    E = sum_i w_i e_i^2
        + (dc_lambda / n) * (sum_i e_i)^2
        + (vw_lambda / n) * sum_{pairs} (e_i + e_{i+n/2})^2

and the fold tree generalises the vesica term to every periodic level.  V2
keeps the exact same objective and adds a vectorised implementation of the
D6/vesica greedy code descent.
"""

from __future__ import annotations

from functools import lru_cache

import numpy as np

from .config import HPCConfig
from .q2k import QK_K, dequantize_params


@lru_cache(maxsize=64)
def _fold_weights_cached(n: int, dc_lambda: float, vw_lambda: float, depth: int, gamma: float) -> np.ndarray:
    levels = int(np.log2(n))
    if (1 << levels) != n:
        raise ValueError("fold pyramid requires a power-of-two length")
    weights = np.zeros(levels, dtype=np.float64)
    g = 1.0
    for k in range(1, levels + 1):
        if k == levels:
            weights[k - 1] = dc_lambda / n
        elif k <= depth:
            weights[k - 1] = (vw_lambda / n) * g
        g *= gamma
    weights.setflags(write=False)
    return weights


def fold_weights(n: int, cfg: HPCConfig) -> np.ndarray:
    return _fold_weights_cached(
        n, cfg.dc_lambda, cfg.vw_lambda, cfg.fold_depth, cfg.fold_gamma
    )


def fold_build(errors: np.ndarray) -> np.ndarray:
    e = np.asarray(errors, dtype=np.float64)
    n = e.size
    py = np.empty(n - 1, dtype=np.float64)
    src = e
    offset = 0
    m = n // 2
    while m >= 1:
        dst = src[:m] + src[m : 2 * m]
        py[offset : offset + m] = dst
        src = dst
        offset += m
        m //= 2
    return py


def fold_energy(pyramid: np.ndarray, weights: np.ndarray) -> float:
    acc = 0.0
    offset = 0
    m = pyramid.size // 2 + 1
    level = 0
    while m >= 1:
        if weights[level] != 0.0:
            v = pyramid[offset : offset + m]
            acc += float(weights[level] * np.dot(v, v))
        offset += m
        m //= 2
        level += 1
    return acc


def spectral_penalty(
    errors: np.ndarray, cfg: HPCConfig, dc_carry: float = 0.0
) -> float:
    if not cfg.use_spectral or (cfg.dc_lambda == 0.0 and cfg.vw_lambda == 0.0):
        return 0.0
    pyramid = fold_build(errors)
    pyramid[-1] += dc_carry
    return fold_energy(pyramid, fold_weights(errors.size, cfg))


def objective(
    x: np.ndarray,
    w: np.ndarray | None,
    d,
    dmin,
    ls: np.ndarray,
    lm: np.ndarray,
    codes: np.ndarray,
    cfg: HPCConfig,
    dc_carry: float = 0.0,
) -> tuple[float, float]:
    xh = dequantize_params(d, dmin, ls, lm, codes)
    e = np.asarray(x, dtype=np.float64) - xh.astype(np.float64)
    sse = float(np.dot(np.asarray(w, dtype=np.float64), e * e)) if w is not None else float(np.dot(e, e))
    return sse + spectral_penalty(e, cfg, dc_carry), sse


def refine_codes_hadamard(
    x: np.ndarray,
    w: np.ndarray | None,
    d,
    dmin,
    ls: np.ndarray,
    lm: np.ndarray,
    codes: np.ndarray,
    cfg: HPCConfig,
    dc_carry: float = 0.0,
) -> np.ndarray:
    """Deterministic greedy descent of ``4*vesica^2 + (DC + carry)^2``.

    Same acceptance rule as the C Phase-4 shaper, but the whole 256-element
    search is evaluated as four vectorised candidate matrices per pass.
    """
    x = np.asarray(x, dtype=np.float64)
    out = np.asarray(codes, dtype=np.uint8).copy()
    ww = np.ones(QK_K, dtype=np.float64) if w is None else np.asarray(w, dtype=np.float64)
    d_sub = np.repeat(np.float32(d) * np.asarray(ls, dtype=np.float32), 16).astype(np.float64)
    m_sub = np.repeat(np.float32(dmin) * np.asarray(lm, dtype=np.float32), 16).astype(np.float64)
    active = d_sub >= 1.0e-15
    if not np.any(active):
        return out

    xh = d_sub * out.astype(np.float64) - m_sub
    e_live = x - xh
    sse_live = float(np.dot(ww, e_live * e_live))
    sse_cap = sse_live * (1.0 + cfg.sse_budget)
    v_live = e_live[: QK_K // 2] + e_live[QK_K // 2 :]
    vesica = float(np.dot(v_live, v_live))
    dc = float(np.sum(e_live)) + dc_carry
    metric = 4.0 * vesica + dc * dc

    idx = np.arange(QK_K)
    pi = np.where(idx < QK_K // 2, idx, idx - QK_K // 2)
    for _ in range(QK_K):
        best_delta = 0.0
        best_k = -1
        best_q = 0
        for q_try in range(4):
            e_new = x - (d_sub * q_try - m_sub)
            delta_e = e_new - e_live
            sse_alt = sse_live + ww * (e_new * e_new - e_live * e_live)
            mask = (sse_alt <= sse_cap) & active & (out != q_try)
            if not np.any(mask):
                continue
            v_new = v_live[pi] + delta_e
            vesica_alt = vesica - v_live[pi] * v_live[pi] + v_new * v_new
            dc_alt = dc + delta_e
            delta = metric - (4.0 * vesica_alt + dc_alt * dc_alt)
            delta = np.where(mask, delta, -np.inf)
            k = int(np.argmax(delta))
            if delta[k] > best_delta:
                best_delta = float(delta[k])
                best_k = k
                best_q = q_try
        if best_k < 0:
            break
        old_e = e_live[best_k]
        new_e = x[best_k] - (d_sub[best_k] * best_q - m_sub[best_k])
        de = new_e - old_e
        p = pi[best_k]
        v_live[p] += de
        vesica += v_live[p] * v_live[p] - (v_live[p] - de) * (v_live[p] - de)
        dc += de
        metric = 4.0 * vesica + dc * dc
        sse_live += ww[best_k] * (new_e * new_e - old_e * old_e)
        e_live[best_k] = new_e
        out[best_k] = np.uint8(best_q)
    return out


def dc_nudge_codes(
    x: np.ndarray,
    d,
    dmin,
    ls: np.ndarray,
    lm: np.ndarray,
    codes: np.ndarray,
    w: np.ndarray | None,
    carry: float,
    cfg: HPCConfig,
    max_passes: int = 4,
) -> np.ndarray:
    """Spend a bounded SSE budget to cancel the rolling DC residual.

    Equivalent in spirit to Phase 4.8 of the C engine: the reconstruction
    target is always the true weight, only the codeword choice changes.
    """
    x = np.asarray(x, dtype=np.float64)
    out = np.asarray(codes, dtype=np.uint8).copy()
    ww = np.ones(QK_K, dtype=np.float64) if w is None else np.asarray(w, dtype=np.float64)
    d_sub = np.repeat(np.float32(d) * np.asarray(ls, dtype=np.float32), 16).astype(np.float64)
    m_sub = np.repeat(np.float32(dmin) * np.asarray(lm, dtype=np.float32), 16).astype(np.float64)
    active = d_sub >= 1.0e-15
    if not np.any(active):
        return out
    e = x - (d_sub * out.astype(np.float64) - m_sub)
    sse = float(np.dot(ww, e * e))
    cap = sse * (1.0 + cfg.sse_budget)
    residual = float(np.sum(e)) + carry
    for _ in range(max_passes):
        best_improvement = 0.0
        best_k = -1
        best_q = 0
        for q_try in range(4):
            e_new = x - (d_sub * q_try - m_sub)
            mask = active & (out != q_try)
            sse_alt = sse + ww * (e_new * e_new - e * e)
            mask &= sse_alt <= cap
            if not np.any(mask):
                continue
            improvement = np.abs(residual) - np.abs(residual + (e_new - e))
            improvement = np.where(mask, improvement, -np.inf)
            k = int(np.argmax(improvement))
            if improvement[k] > best_improvement:
                best_improvement = float(improvement[k])
                best_k = k
                best_q = q_try
        if best_k < 0:
            break
        old_e = e[best_k]
        new_e = x[best_k] - (d_sub[best_k] * best_q - m_sub[best_k])
        de = new_e - old_e
        residual += de
        sse += ww[best_k] * (new_e * new_e - old_e * old_e)
        e[best_k] = new_e
        out[best_k] = np.uint8(best_q)
        if abs(residual) < 1.0e-12:
            break
    return out


def pin_dc_codes_uniform(
    x: np.ndarray,
    d: np.ndarray,
    codes: np.ndarray,
    unit: np.ndarray,
    lo: int,
    hi: int,
    budget_frac: float = 0.05,
    max_rounds: int = 96,
) -> tuple[np.ndarray, int]:
    """Zero-DC pin for single-scale lattices (``Q4_0`` nibble, ``Q3_K`` signed grid).

    These formats carry only one free fp16 parameter per block, so the exact
    ``sum(e) == 0`` scale re-solve that gives ``Q2_K`` its free lunch is
    structurally degenerate (``sum(a) ~= 0`` for symmetric code vectors, and
    the ``fp16`` scale granularity turns any forced ``d`` move into a large
    SSE penalty).  The pin therefore spends ``budget_frac`` of each block's
    SSE (an exchange, never a free move) flipping codes toward the DC
    residual: each round picks, per block, the cheapest-to-DC flip direction
    and position by the ``SSE cost / |delta DC|`` ratio (cheap "free"
    strict-improvement flips first), applies it until the budget is spent or
    the residual is below half a flip quantum.  Measured on real tensors:
    ``Q3_K`` DC RMS −82% at ~0.9% extra SSE; ``Q4_0`` −43% at ~1.3% (the
    symmetric ``[-8, 7]`` grid has half the DC slack of ``Q3_K``'s signed
    ``sc`` units, so its residual lands at the flip-quantum floor rather
    than ~0).  Deterministic; safe on any input state.
    """
    X = np.asarray(x, dtype=np.float64)
    B, n = X.shape
    codes = np.asarray(codes, dtype=np.int64).copy()
    unit = np.broadcast_to(np.asarray(unit, dtype=np.float64), codes.shape)
    if B == 0 or budget_frac <= 0.0:
        return codes, 0
    d = np.asarray(d, dtype=np.float64)
    d_sub = np.repeat(d[:, None], n, axis=1)
    rows = np.arange(B)
    e = X - d_sub * codes * unit
    dc = e.sum(axis=1)
    sse0 = np.sum(e * e, axis=1)
    cap = sse0 * float(budget_frac)
    spent = np.zeros(B)
    touched = np.zeros(B, dtype=bool)
    tol = np.abs(d) * 0.5

    for _ in range(max_rounds):
        want_cost = np.full(B, np.inf)
        want_pos = np.full(B, -1, dtype=np.int64)
        want_dir = np.zeros(B, dtype=np.int64)
        for direction in (1, -1):
            if direction == 1:
                sel = codes < hi
            else:
                sel = codes > lo
            de = np.where(sel, float(direction), 0.0) * unit * -d_sub
            dE = 2.0 * e * de + de * de
            push = dc[:, None] * de < 0.0
            valid = sel & push
            ratio = np.where(valid & (np.abs(de) > 1.0e-300), dE / np.abs(np.where(valid, de, 1.0)), np.inf)
            pmin = np.argmin(ratio, axis=1)
            rmin = ratio[rows, pmin]
            cmin = dE[rows, pmin]
            take = np.isfinite(rmin) & (rmin < want_cost)
            want_cost = np.where(take, cmin, want_cost)
            want_pos = np.where(take, pmin, want_pos)
            want_dir = np.where(take, direction, want_dir)
        want = (np.abs(dc) > tol) & (want_pos >= 0) & (spent + want_cost <= cap)
        if not np.any(want):
            break
        bi = np.where(want)[0]
        pos = want_pos[bi]
        dirv = want_dir[bi]
        de = -d_sub[bi, pos] * dirv * unit[bi, pos]
        spent[bi] += 2.0 * e[bi, pos] * de + de * de
        e[bi, pos] += de
        codes[bi, pos] += dirv
        dc[bi] += de
        touched[bi] = True

    return codes, int(np.sum(touched))


def pin_dc_codes_uniform_cuda(
    x: "torch.Tensor",
    d: "torch.Tensor",
    codes: "torch.Tensor",
    unit: "torch.Tensor",
    lo: int,
    hi: int,
    budget_frac: float = 0.05,
    max_rounds: int = 96,
) -> tuple["torch.Tensor", int]:
    """GPU port of :func:`pin_dc_codes_uniform` (same greedy exchange).

    ``x``/``d``/``codes``/``unit`` all live on one CUDA device; ``codes`` is
    int64 and is updated in place (returned tensor is the same object).
    Float32 throughout — matches the CUDA pipeline's precision, and the
    exchange is self-correcting (any round that slightly mis-ranks a flip
    simply loses one candidate, the residual still never regresses SSE
    beyond ``budget_frac`` of the block's starting SSE).
    """
    import torch

    if budget_frac <= 0.0 or codes.shape[0] == 0:
        return codes, 0
    X = x.float()
    B, n = X.shape
    dev = X.device
    if codes.dtype != torch.int64:
        codes = codes.to(torch.int64)
    d = d.to(dev).float()
    if unit.dim() == 1 and unit.numel() == n:
        unit = unit[None, :]
    unit = unit.to(dev).float()
    if unit.shape[0] == 1 and B != 1:
        unit = unit.expand(B, n)

    d_col = d[:, None]
    e = X - d_col * codes.float() * unit
    dc = e.sum(dim=1)
    sse0 = (e * e).sum(dim=1)
    cap = sse0 * float(budget_frac)
    spent = torch.zeros(B, device=dev)
    rows = torch.arange(B, device=dev)
    tol = d.abs() * 0.5
    touched = 0

    for _ in range(max_rounds):
        want_cost = torch.full((B,), float("inf"), device=dev)
        want_pos = torch.full((B,), -1, dtype=torch.int64, device=dev)
        want_dir = torch.zeros(B, dtype=torch.int64, device=dev)
        for direction in (1, -1):
            sel = (codes < hi) if direction == 1 else (codes > lo)
            de = sel.float() * (float(direction) * unit * -d_col)
            dE = 2.0 * e * de + de * de
            valid = sel & (dc[:, None] * de < 0.0)
            ratio = torch.where(valid, dE.abs() / de.abs().clamp_min(1.0e-30), torch.inf)
            pmin = ratio.argmin(dim=1)
            rmin = ratio[rows, pmin]
            cmin = dE[rows, pmin]
            take = torch.isfinite(rmin) & (rmin < want_cost)
            want_cost = torch.where(take, cmin, want_cost)
            want_pos = torch.where(take, pmin, want_pos)
            want_dir = torch.where(take, torch.full_like(want_dir, direction), want_dir)
        want = (
            (dc.abs() > tol)
            & (want_pos >= 0)
            & (spent + want_cost <= cap)
        )
        if not bool(want.any()):
            break
        bi = want.nonzero(as_tuple=False).squeeze(1)
        pos = want_pos[bi]
        dirv = want_dir[bi]
        de = -d[bi] * dirv.float() * unit[bi, pos]
        spent[bi] += 2.0 * e[bi, pos] * de + de * de
        e[bi, pos] += de
        codes[bi, pos] += dirv
        dc[bi] += de
        touched += int(bi.numel())

    return codes, touched