"""In-process, for the study (the package files are untouched): the encoder's selection without the (blocks, longest
part, M+1) basis array. `_parts.parts_system` copies each part's height basis to its 2(2N+1) blocks, padded to the
longest part (6-13 GB at 2048 layers); `_select._project` only needs one basis per part. Here parts_system returns the
bases once per part (`PartBases`) and _project projects a part's blocks with a matrix product. The atoms and their
order are the package's up to float32 round-off of the sums (matmul against einsum). Import after `common`."""
import inspect
import sys

import numpy as np

import common  # noqa: F401  (puts the package on the path)

P = sys.modules["fcp._parts"]
S = sys.modules["fcp._select"]
F = P.F


class PartBases:
    """The blocks' height bases, one (Hmax, J) array per part, `share` consecutive blocks per part."""

    def __init__(self, bases, share, hmax, J):
        self.bases, self.share = bases, share
        self.shape = (len(bases) * share, hmax, J)


def _project(Y, Phi, w, steps, share=1):
    if not isinstance(Phi, PartBases):
        return _orig_project(Y, Phi, w, steps, share)
    T = np.asarray(Y, F)
    B, _, J = Phi.shape
    v = np.empty((B, J), F)
    norms = np.empty((B, J), F)
    for p, Tp in enumerate(Phi.bases):
        g = slice(p * Phi.share, (p + 1) * Phi.share)
        v[g] = T[g] @ Tp
        norms[g] = np.linalg.norm(Tp, axis=0)
    safe = np.where(norms > 1e-14, norms, np.inf)
    sc = np.abs(v) / safe * np.sqrt(np.asarray(w, F))[:, None]
    idx = np.argsort(-sc, axis=1, kind="stable")[:, :steps]
    return idx, np.take_along_axis(sc, idx, 1), np.take_along_axis(v, idx, 1)


_orig_project = S._project
S._project = _project

# parts_system with its last line changed: the per-part bases instead of np.array(Phis)
_src = inspect.getsource(P.parts_system)
_old = "return parts, np.array(Ys), np.array(Phis), np.array(Ws)"
if _old not in _src:
    raise RuntimeError("lowmem: parts_system changed; update the patch")
_src = _src.replace("Phis.append(Tp)", "pass").replace("for p, C in zip(parts, coeffs):",
                                                      "Bases = []\n    for p, C in zip(parts, coeffs):", 1)
_src = _src.replace("        Tp[:n, : Q.shape[1]] = Q\n", "        Tp[:n, : Q.shape[1]] = Q\n        Bases.append(Tp)\n", 1)
_src = _src.replace(_old, "return parts, np.array(Ys), PartBases(Bases, 2 * K, Hmax, M + 1), np.array(Ws)")
if __import__("os").environ.get("FCP_KEEP1") == "1":  # keep tracks of one layer (the package drops them: len(tr) < 2)
    if "len(tr) < 2 or " not in _src:
        raise RuntimeError("lowmem: the one-layer rule changed; update the patch")
    _src = _src.replace("len(tr) < 2 or ", "len(tr) < 1 or ")
P.PartBases = PartBases
exec(compile(_src, P.__file__ + " (lowmem)", "exec"), vars(P))  # in the module's namespace: its constants stay live