import sys
import numpy as np
sys.path.insert(0, "scripts")
from common import UP, fcp, load_mesh
P_ = sys.modules["fcp._parts"]
tol = float(sys.argv[1]) if len(sys.argv) > 1 else 0.01
cap = {}
orig = P_.split_by_fit
def spy(parts, coeffs, areas, M, tol=None, follows=None):
    cap["in"] = (parts, coeffs, areas)
    out = orig(parts, coeffs, areas, M, tol, follows)
    cap["out"] = out
    raise KeyboardInterrupt
P_.split_by_fit = spy
try:
    fcp.encode(load_mesh("dog"), atoms=10, layers=2048, N=128, M=64, up=UP, device="cuda", fit_tol=tol)
except KeyboardInterrupt:
    pass
parts, coeffs, areas = cap["out"]
pin = cap["in"][0]
print(f"fit_tol {tol}: tracks {len(pin)} -> parts {len(parts)}; longest track {max(p['h1']-p['h0']+1 for p in pin)}, longest part {max(p['h1']-p['h0']+1 for p in parts)}")
i = max(range(len(parts)), key=lambda j: parts[j]["h1"] - parts[j]["h0"])
p, C, A = parts[i], np.asarray(coeffs[i], float), areas[i]
K = C.shape[2]; w = np.full(K, 0.5); w[0] = 1
Cw = (P_.continuous_start(C, A) * np.sqrt(w)).reshape(len(C), -1)
Q, _ = P_.orth_basis(len(C), 64)
r = np.sqrt(((Cw - Q @ (Q.T @ Cw)) ** 2).sum(1))
print(f"longest part {p}: degree {Q.shape[1]-1}, per-ring RMS miss: median {np.median(r):.4f} max {r.max():.4f} at layer {p['h0'] + int(r.argmax())}")
for L in (700, 760, 780, 784, 786, 790, 800):
    j = L - p["h0"]
    if 0 <= j < len(r): print(f"  layer {L}: RMS miss {r[j]:.4f}")
# without continuity (+x-ray start)
Cw0 = (C * np.sqrt(w)).reshape(len(C), -1)
r0 = np.sqrt(((Cw0 - Q @ (Q.T @ Cw0)) ** 2).sum(1))
print(f"  same with raw starts: max {r0.max():.4f}")
jump = np.linalg.norm(np.diff(Cw, axis=0), axis=1)
print("  largest ring-to-ring coefficient steps at layers", (p["h0"] + 1 + np.argsort(-jump)[:6]).tolist(), np.sort(jump)[::-1][:6].round(4))
def err(a, b):
    Q, _ = P_.orth_basis(b - a, 64)
    X = Cw[a:b]
    return ((X - Q @ (Q.T @ X)) ** 2).sum(1)
e = err(0, len(Cw))
cs = np.arange(P_.FIT_MIN, len(Cw) - P_.FIT_MIN + 1)
sses = np.array([err(0, c).sum() + err(c, len(Cw)).sum() for c in cs])
jump_c = jump[cs - 1]
k = int(np.argmax(np.where(sses <= (1 + P_.STEP_SLACK) * sses.min(), jump_c, -1.0)))
print(f"cut test on the kept part: total SSE {e.sum():.4f}; best single cut at layer {p['h0'] + cs[sses.argmin()]} -> {sses.min():.4f} "
      f"({sses.min() / e.sum():.2f} of it); rule's pick at layer {p['h0'] + cs[k]} -> {sses[k]:.4f} ({sses[k] / e.sum():.2f}); needs <= {P_.FIT_GAIN}")
# two cuts bracketing the worst ring
j = int(r.argmax())
for half in (4, 8, 16):
    a, b = max(j - half, 0), min(j + half, len(Cw))
    print(f"  cuts at layers {p['h0']+a} and {p['h0']+b}: SSE {(err(0, a).sum() + err(a, b).sum() + err(b, len(Cw)).sum()) / e.sum():.2f} of it, "
          f"max ring miss {np.sqrt(max(err(0, a).max(), err(a, b).max(), err(b, len(Cw)).max())):.4f}")