"""In-process fix of fcp._parts.split_by_fit for the study (the package file is untouched): the cut at the largest step is
chosen only among the cuts that pass the gain test. In the package rule the step pick may fail FIT_GAIN while the
best cut passes, and the part is then never cut (dog, 2048 layers: a 719-layer part whose best cut lowers the squared
error to 0.65, the step pick to 0.95, rings missed by 14x fit_tol). Import this module before encoding."""
import sys

import numpy as np

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

P = sys.modules["fcp._parts"]


def split_by_fit(parts, coeffs, areas, M, tol=None, follows=None):
    tol = P.FIT_TOL if tol is None else tol
    out_p, out_c, out_a = [], [], []
    for i, (p, C, A) in enumerate(zip(parts, coeffs, areas)):
        C = np.asarray(C, float)
        K = C.shape[2]
        w = np.full(K, 0.5)
        w[0] = 1.0
        Cw = (P.continuous_start(C, A) * np.sqrt(w)).reshape(len(C), -1)

        def fit(a, b):
            Q, _ = P.orth_basis(b - a, M)
            X = Cw[a:b]
            return X, Q @ (Q.T @ X)

        def err(a, b):
            X, Xf = fit(a, b)
            return ((X - Xf) ** 2).sum(1)

        cuts, todo = [], [(0, len(C))]
        while todo:
            a, b = todo.pop()
            if b - a < 2 * P.FIT_MIN:
                continue
            e = err(a, b)
            if follows is None and np.sqrt(e.max()) <= tol:
                continue
            if follows is not None and follows(i, a, b, fit(a, b)[1].reshape(b - a, 2, K) / np.sqrt(w)):
                continue
            cs = np.arange(a + P.FIT_MIN, b - P.FIT_MIN + 1)
            sses = np.array([err(a, c).sum() + err(c, b).sum() for c in cs])
            ok = sses <= min((1 + P.STEP_SLACK) * sses.min(), P.FIT_GAIN * e.sum())  # the fix: the gain test first
            if not ok.any():
                continue
            jump = np.linalg.norm(np.diff(Cw[a:b], axis=0), axis=1)[cs - a - 1]
            k = int(np.argmax(np.where(ok, jump, -1.0)))
            cuts.append(int(cs[k]))
            todo += [(a, int(cs[k])), (int(cs[k]), b)]
        edges = [0, *sorted(cuts), len(C)]
        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], P.F))
            out_a.append(list(A[a:b]))
    return out_p, out_c, out_a


P.split_by_fit = split_by_fit