import itertools
import numpy as np
from scipy.optimize import linear_sum_assignment
import Rhino.Geometry as rg

PERMS = list(itertools.permutations(range(3)))

def kabsch(A, B):
    """R (3x3, det=+1) that minimizes sum ||R a_i - b_i||^2. Rows = vectors."""
    H = A.T @ B
    U, S, Vt = np.linalg.svd(H)
    d = np.sign(np.linalg.det(Vt.T @ U.T)) or 1.0
    return Vt.T @ np.diag([1.0, 1.0, d]) @ U.T

def pair_cost(F, N, stem_idx=None, w_dir=0.0):
    """F, N: 3x3 unit vectors (as rows). Return (cost, R, perm)."""
    best = (np.inf, None, None)
    for p in PERMS:
        Fp = F[list(p)]
        R = kabsch(Fp, N)
        rot = Fp @ R.T
        ang = np.degrees(np.arccos(np.clip(np.sum(rot * N, axis=1), -1.0, 1.0)))
        c = float(np.sum(ang ** 2))
        if w_dir and stem_idx is not None:
            stem_world = R @ F[stem_idx]
            if stem_world[2] > 0:           # stem points upward -> penalty
                c += w_dir * stem_world[2]
        if c < best[0]:
            best = (c, R, p)
    return best

def cost_matrix(Fu, Nu, stem, w_dir=0.0):
    """Vectorized (fast) version of pair_cost for all forks x nodes at once.
    Return C (n,m) and the best R (n,m,3,3)."""
    P = np.array(PERMS)                                   # (6,3)
    Fp = Fu[:, P]                                         # (n,6,3,3)
    H = np.einsum('ipka,jkb->ijpab', Fp, Nu)              # (n,m,6,3,3)
    U, S, Vt = np.linalg.svd(H)
    d = np.sign(np.linalg.det(U) * np.linalg.det(Vt))
    d[d == 0] = 1.0
    D = np.zeros(H.shape); D[..., 0, 0] = 1; D[..., 1, 1] = 1; D[..., 2, 2] = d
    R = np.swapaxes(Vt, -1, -2) @ D @ np.swapaxes(U, -1, -2)
    rot = np.einsum('ijpab,ipkb->ijpka', R, Fp)
    cos = np.clip(np.einsum('ijpka,jka->ijpk', rot, Nu), -1.0, 1.0)
    cost = np.sum(np.degrees(np.arccos(cos)) ** 2, axis=-1)   # (n,m,6)
    if w_dir:
        fs = Fu[np.arange(len(Fu)), stem]                 # (n,3) original stem
        z = np.einsum('ijpb,ib->ijp', R[..., 2, :], fs)   # Z component of stem after rotation
        cost = cost + w_dir * np.maximum(z, 0.0)
    best = np.argmin(cost, axis=-1)                       # (n,m)
    C = np.take_along_axis(cost, best[..., None], -1)[..., 0]
    Rb = np.take_along_axis(R, best[..., None, None, None], 2)[:, :, 0]
    return C, Rb

def to_np(vs, name):
    if hasattr(vs, "X") or not hasattr(vs, "__len__"):
        raise ValueError("Input '%s' is still set to Item Access. Right-click input '%s' -> choose List Access." % (name, name))
    bad = [v for v in vs if not hasattr(v, "X")]
    if bad:
        raise ValueError("Input '%s' contains %s, not vectors. Check the wire & Type Hint (must be Vector3d / No Type Hint)." % (name, type(bad[0]).__name__))
    return np.array([[v.X, v.Y, v.Z] for v in vs], dtype=float)

# ---------------- main ----------------
_g = globals()
fmeshes = _g.get("fmeshes")
exclude = _g.get("exclude")
w_dir = _g.get("w_dir")

assign, cost, xform, second, placed = [], [], [], [], []
total, mean = 0.0, 0.0
summary = "Incomplete input: needs fvecs, fnodes, nvecs, nnodes (List Access)"

if fvecs and nvecs and fnodes and nnodes:
    Fr = to_np(fvecs, 'fvecs').reshape(-1, 3, 3)                 # (n,3,3) raw
    lens = np.linalg.norm(Fr, axis=2)                    # chord lengths
    Fu = Fr / lens[:, :, None]
    Nu = to_np(nvecs, 'nvecs').reshape(-1, 3, 3)
    Nu = Nu / np.linalg.norm(Nu, axis=2)[:, :, None]
    n, m = Fu.shape[0], Nu.shape[0]
    excl = set(exclude or [])
    w = float(w_dir or 0.0)

    BIG = 1e9
    stem = np.argmax(lens, axis=1)                       # assumption: longest chord = stem
    C, Rb = cost_matrix(Fu, Nu, stem, w)
    for i in excl:
        if 0 <= i < n:
            C[i, :] = BIG

    rows, cols = linear_sum_assignment(C)               # can be rectangular (ideally n>=m)
    col2row = {c: r for r, c in zip(rows, cols) if C[r, c] < BIG}
    used = set(col2row.values())

    for j in range(m):
        i = col2row.get(j, -1)
        assign.append(i)
        if i < 0:
            cost.append(-1.0); xform.append(rg.Transform.Identity); second.append(-1)
            continue
        cost.append(float(C[i, j]))
        R = Rb[i, j]
        src = fnodes[i]; dst = nnodes[j]
        p_src = rg.Plane(src, rg.Vector3d.XAxis, rg.Vector3d.YAxis)
        p_dst = rg.Plane(dst, rg.Vector3d(*R[:, 0]), rg.Vector3d(*R[:, 1]))
        xform.append(rg.Transform.PlaneToPlane(p_src, p_dst))
        spare = [k for k in range(n) if k not in used and k not in excl]
        second.append(min(spare, key=lambda k: C[k, j]) if spare else -1)

    valid = [c for c in cost if c >= 0]
    total = float(sum(valid))
    mean = total / len(valid) if valid else 0.0

    if fmeshes:
        for j, i in enumerate(assign):
            if i >= 0 and i < len(fmeshes) and fmeshes[i] is not None:
                mm = fmeshes[i].DuplicateMesh()
                mm.Transform(xform[j])
                placed.append(mm)

    summary = ("%d forks -> %d nodes | filled %d | total cost %.0f | mean/node %.0f"
               % (n, m, len(valid), total, mean))