"""Faithful Python port of the original C factored solver.

This module is intentionally not on the production path.  It exists so the
V2 improvements can be measured against the exact reference trajectory:
fixed six-iteration 1-D Lloyd, single range initialisation, no fp16-domain
polish and no spectral shaping inside the solver.
"""

from __future__ import annotations

import numpy as np

from .config import HPCConfig
from .fp16 import from_fp16_bits, nearest_int, to_fp16_bits
from .q2k import QK_K, codes_from_params, dequantize_params, weighted_sse
from .spectral import objective, spectral_penalty


def reference_factor_1d(v, w, nmax: int = 15):
    v = np.maximum(np.asarray(v, dtype=np.float64), 0.0)
    w = np.maximum(np.asarray(w, dtype=np.float64), 0.0)
    vmax = float(np.max(v))
    if vmax < 1.0e-30:
        return 0.0, np.zeros(v.size, dtype=np.uint8)
    d = vmax / nmax
    for _ in range(6):
        codes = np.clip(nearest_int(v / d), 0, nmax)
        den = float(np.sum(w * codes * codes))
        if den < 1.0e-30:
            break
        dn = float(np.sum(w * v * codes) / den)
        conv = abs(dn - d) <= 1.0e-7 * abs(d) + 1.0e-30
        d = dn
        if conv:
            break
    codes = np.clip(nearest_int(v / d), 0, nmax)
    return d, codes.astype(np.uint8)


def reference_sub_ls(x, w, codes):
    NS = 16
    q = np.asarray(codes, dtype=np.float64).reshape(NS, 16)
    xb = np.asarray(x, dtype=np.float64).reshape(NS, 16)
    wb = np.asarray(w, dtype=np.float64).reshape(NS, 16)
    s = np.zeros(NS)
    t = np.zeros(NS)
    for j in range(NS):
        ww, qq, xx = wb[j], q[j], xb[j]
        Sqq = float(np.sum(ww * qq * qq))
        Swq = float(np.sum(ww * qq))
        Sw = float(np.sum(ww))
        Sqx = float(np.sum(ww * qq * xx))
        Sx = float(np.sum(ww * xx))
        det = Sqq * Sw - Swq * Swq
        if det > 1.0e-30:
            sj = (Sqx * Sw - Swq * Sx) / det
            tj = (sj * Swq - Sx) / Sw
        else:
            sj = tj = 0.0
        if sj < 0.0:
            sj = 0.0
        if tj < 0.0:
            tj = 0.0
            sj = Sqx / Sqq if Sqq > 1.0e-30 else 0.0
            if sj < 0.0:
                sj = 0.0
        s[j], t[j] = sj, tj
    return s, t


def reference_solve(x, w, cfg: HPCConfig, use_spectral: bool) -> tuple[float, float]:
    NS = 16
    xb = np.asarray(x, dtype=np.float64).reshape(NS, 16)
    wb = np.asarray(w, dtype=np.float64).reshape(NS, 16)
    s = np.zeros(NS)
    t = np.zeros(NS)
    codes = np.zeros(QK_K, dtype=np.uint8)
    for j in range(NS):
        sub = xb[j]
        xmin = min(float(sub.min()), 0.0)
        xmax = float(sub.max())
        span = xmax - xmin
        sj = span / 3.0 if span > 1.0e-30 else 0.0
        tj = -xmin
        s[j], t[j] = sj, tj
        if sj > 1.0e-30:
            codes[16 * j : 16 * (j + 1)] = np.clip(
                nearest_int((sub + tj) / sj), 0, 3
            ).astype(np.uint8)

    sw = wb.sum(axis=1)
    d, ls = reference_factor_1d(s, sw)
    m, lm = reference_factor_1d(t, sw)
    L = codes_from_params(x, d, m, ls, lm)
    obj, _ = objective(x, w, d, m, ls, lm, L, cfg)
    best = (obj, d, m, ls.copy(), lm.copy(), L.copy())
    stall = 0
    for _ in range(16):
        L = codes_from_params(x, d, m, ls, lm)
        s, t = reference_sub_ls(x, w, L)
        ws = np.sum(wb * L.astype(np.float64).reshape(NS, 16) ** 2, axis=1)
        wt = wb.sum(axis=1)
        d, ls = reference_factor_1d(s, ws)
        m, lm = reference_factor_1d(t, wt)
        L = codes_from_params(x, d, m, ls, lm)
        cur, _ = objective(x, w, d, m, ls, lm, L, cfg)
        if cur < best[0]:
            best = (cur, d, m, ls.copy(), lm.copy(), L.copy())
            stall = 0
        else:
            stall += 1
            if stall >= 2:
                break

    _, d, m, ls, lm, L = best
    d = float(from_fp16_bits(to_fp16_bits(np.float32(d))))
    m = float(from_fp16_bits(to_fp16_bits(np.float32(m))))
    L = codes_from_params(x, d, m, ls, lm)
    e = np.asarray(x, dtype=np.float64) - dequantize_params(d, m, ls, lm, L).astype(np.float64)
    sse = weighted_sse(x, w, d, m, ls, lm, L)
    obj = sse + (spectral_penalty(e, cfg) if use_spectral else 0.0)
    return obj, sse


def synthetic_cases(seed: int = 123, per_kind: int = 12):
    rng = np.random.default_rng(seed)
    cases = []
    for _ in range(per_kind):
        cases.append(("gaussian", rng.normal(0, 0.02, QK_K).astype(np.float32)))
    for _ in range(per_kind):
        cases.append(("student-t", (rng.standard_t(3, QK_K) * 0.01).astype(np.float32)))
    for _ in range(per_kind):
        mask = rng.random(QK_K) < 0.1
        cases.append(("sparse", (rng.normal(0, 0.02, QK_K) * mask).astype(np.float32)))
    for _ in range(per_kind):
        mask = rng.random(QK_K) < 0.5
        cases.append(("offset-sparse", (rng.normal(0, 0.02, QK_K) * mask + 0.01).astype(np.float32)))
    return cases