"""The rate-distortion partition: each track cut where the descriptor's own objective says so, with no tolerances.

With orthonormal atoms, keeping those whose score (|value| sqrt(w)) is above a level lam is the best descriptor of its size,
and its error plus lam^2 per atom kept is sum over atoms of min(score^2, lam^2). The partition is chosen by the same rule: a
piece of a track (layers i..j-1) costs sum over its atoms of min(score^2, lam^2) (its rings projected on its own height basis,
`_parts.orth_basis`, start points by continuity, Parseval weights), plus the energy its basis cannot hold, plus lam^2 x HEADER
(a part's header in atoms: 3 bytes against 6); each track is cut into the pieces of least total cost (dynamic programming over
cut positions, pieces of at least 2 layers). lam comes from the data, not the budget, so prefixes stay exact: the noise of the
rings' coefficients, read from their second differences along the height (median |d2| / 0.6745 / sqrt(6), in score units,
over every track), times 2 ("rd2", as atoms="auto") or the universal threshold sqrt(2 ln N) ("rdU", N the coefficients);
or "rdT<T0>": the level of the T0-th atom (iterated with the partition), the partition best at that reference budget.
Swapped in for split_by_fit (no package change); the track filter of v6 (clouds) is left as it is.

    from rd_split import install; install("rd2")     # then fcp.encode(...) cuts by it
"""
import numpy as np

import fcp  # noqa: F401  (the package on the path: imported by the caller first)
from fcp import _parts

HEADER = 0.5  # a part's header (sign, two ends: 3 B) in atoms (index + value: 6 B)
MIN_LEN = 2
_split = _parts.split_by_fit
LOG = {}


def _weighted(C, A):
    K = C.shape[2]
    w = np.full(K, 0.5)
    w[0] = 1.0
    return (_parts.continuous_start(np.asarray(C, float), np.asarray(A, float)) * np.sqrt(w)).reshape(len(C), -1)


def noise_sigma(Xs):
    """The coefficients' noise in score units from their second differences along the height (every track of >= 5 layers)."""
    d2 = [np.diff(X, 2, axis=0).ravel() for X in Xs if len(X) >= 5]
    if not d2:
        return 0.0
    return float(np.median(np.abs(np.concatenate(d2))) / 0.6745 / np.sqrt(6.0))


class Track:
    """A track's candidate pieces, projected once: per piece length L, the squared scores of every window's atoms (float32)
    and the energy its basis cannot hold; `costs(lam2)` then only sums min(s^2, lam2)."""

    def __init__(self, X, M):
        self.n = n = len(X)
        X = np.asarray(X, np.float32)
        ce = np.r_[0, np.cumsum((X.astype(float) ** 2).sum(1))]
        self.by_len = []
        for L in range(MIN_LEN, n + 1) if n >= 2 * MIN_LEN else [n]:
            Q = np.asarray(_parts.orth_basis(L, M)[0], np.float32)
            V = np.lib.stride_tricks.sliding_window_view(X, L, axis=0) @ Q  # (windows, B, J)
            s2 = (V * V).reshape(len(V), -1)
            resid = np.maximum((ce[L:] - ce[:-L]) - s2.sum(1, dtype=float), 0)
            self.by_len.append((L, s2, resid))

    def costs(self, lam2):
        cost = np.full((self.n + 1, self.n + 1), np.inf)
        l2 = np.float32(lam2)
        for L, s2, resid in self.by_len:
            i = np.arange(len(s2))
            cost[i, i + L] = np.minimum(s2, l2).sum(1, dtype=float) + resid + lam2 * HEADER
        return cost

    def scores(self, cuts):
        """The squared scores of the atoms of the pieces between cuts."""
        edges = [0, *cuts, self.n]
        look = {L: s2 for L, s2, _ in self.by_len}
        return np.concatenate([look[b - a][a] for a, b in zip(edges[:-1], edges[1:])])


def best_cuts(cost, n):
    best = np.full(n + 1, np.inf)
    best[0] = 0.0
    arg = np.zeros(n + 1, int)
    for j in range(1, n + 1):
        c = best[:j] + cost[:j, j]
        arg[j] = int(np.argmin(c))
        best[j] = c[arg[j]]
    cuts, j = [], n
    while j > 0:
        cuts.append(arg[j])
        j = arg[j]
    return sorted(cuts)[1:]


def _partition(tracks, lam2):
    """Per track: (cuts, the squared scores of its pieces' atoms) at the level lam2."""
    out = []
    for T in tracks:
        cuts = best_cuts(T.costs(lam2), T.n) if T.n >= 2 * MIN_LEN else []
        out.append((cuts, T.scores(cuts)))
    return out


ITER = 4  # "rdT<T0>": the level of the T0-th atom, refound after each partition


def make(mode):
    def split(parts, coeffs, areas, M, tol=None, follows=None):
        Xs = [_weighted(C, A) for C, A in zip(coeffs, areas)]
        sigma = noise_sigma(Xs)
        N = sum(X.size for X in Xs)
        tracks = [Track(X, M) for X in Xs]
        if mode.startswith("rdT"):  # the partition that is best at a reference budget T0: lam = the T0-th atom's score
            T0 = int(mode[3:])
            s2 = np.sort(np.concatenate([T.scores([]) for T in tracks]))[::-1]  # start: every track one piece
            lam2 = float(s2[min(T0, len(s2)) - 1]) if len(s2) else 4 * sigma * sigma
            for _ in range(ITER):
                res = _partition(tracks, max(lam2, 1e-30))
                s2 = np.sort(np.concatenate([sc for _, sc in res]))[::-1]
                new = float(s2[min(T0, len(s2)) - 1]) if len(s2) else lam2
                if abs(new - lam2) <= 1e-6 * max(new, 1e-30):
                    break
                lam2 = new
        else:
            lam = sigma * (2.0 if mode == "rd2" else np.sqrt(2 * np.log(max(N, 2))))
            lam2 = lam * lam
        lam2 = max(lam2, 1e-30)
        LOG.update(sigma=sigma, lam=float(np.sqrt(lam2)))
        res = _partition(tracks, lam2)
        out_p, out_c, out_a = [], [], []
        for (cuts, _), p, C, A in zip(res, parts, coeffs, areas):
            n = len(C)
            edges = [0, *cuts, n]
            for a, b in zip(edges[:-1], edges[1:]):
                out_p.append({"sign": p["sign"], "h0": p["h0"] + a, "h1": p["h0"] + b - 1})
                out_c.append(np.asarray(C[a:b], _parts.F))
                out_a.append(list(A[a:b]))
        return out_p, out_c, out_a
    return split


def install(mode):
    """Swap the package's cut rule for the rate-distortion partition ("rd2" / "rdU"), or back ("package")."""
    _parts.split_by_fit = _split if mode == "package" else make(mode)