"""Dense point tracking via nearest-neighbour matching in canonical space.

The model predicts, for every pixel, a coordinate in a shared *canonical* facial
space (the ``deformation`` output). Two pixels in different frames that map to
the same canonical coordinate are in correspondence.

We seed a set of tracks on the first frame and assign each a distinct color.
For every frame, each seed's canonical nearest neighbours are recolored with the
seed's color, while all other points keep their RGB color. Because the seed
canonical coordinates are fixed, corresponding points get the *same* color in
every frame — a temporally consistent track visualization (no lines/trails).
"""
from __future__ import annotations

import numpy as np
from scipy.spatial import cKDTree

from .colorize import hsv_palette


def compute_track_colors(frames, n_tracks: int = 300, k: int = 20,
                         threshold: float = 0.02, seed: int = 0, seed_frame: int = 0,
                         regions=None):
    """Recolor canonical correspondences of seeded tracks, consistently per frame.

    Args:
        frames: list of dicts with ``canonical`` (M,3), ``rgb`` (M,3) uint8 and
            ``pix`` (M,2) for each frame (all 1:1 aligned).
        n_tracks: number of seed tracks selected on ``seed_frame``.
        k: recolor the ``k`` canonical nearest neighbours of each seed (a small
            visible blob); they share the seed's color.
        threshold: max canonical distance for a neighbour to count as a match.
        seed: RNG seed for reproducible track selection.
        regions: optional list (1:1 with ``frames``) of (M,) boolean arrays. When
            given, seeds are drawn only from ``regions[seed_frame]`` and only
            points inside ``regions[fi]`` are recolored on frame ``fi`` — used to
            keep the colorful tracks on a sub-region (e.g. face + hair) instead of
            the whole subject (shoulders / clothing).

    Returns:
        per_frame_colors: list of (M,3) uint8 — RGB with matched points recolored.
        per_frame_overlay: list of (pix (P,2) int32, col (P,3) uint8) — the
            recolored pixels for 2D overlays.
    """
    ref = frames[seed_frame]
    ref_can = np.asarray(ref["canonical"], dtype=np.float32)
    M0 = ref_can.shape[0]
    if M0 == 0:
        return ([np.asarray(f["rgb"], np.uint8) for f in frames],
                [(np.zeros((0, 2), np.int32), np.zeros((0, 3), np.uint8)) for _ in frames])

    rng = np.random.default_rng(seed)
    # candidate seeds: restricted to the allowed region on the seed frame, if any
    if regions is not None:
        cand = np.nonzero(np.asarray(regions[seed_frame], bool))[0]
        if cand.size == 0:
            cand = np.arange(M0)
    else:
        cand = np.arange(M0)
    n = min(n_tracks, cand.size)
    seed_idx = rng.choice(cand, size=n, replace=False)
    seed_can = ref_can[seed_idx]
    palette = hsv_palette(n)

    per_frame_colors, per_frame_overlay = [], []
    for fi, fr in enumerate(frames):
        can = np.asarray(fr["canonical"], dtype=np.float32)
        cols = np.asarray(fr["rgb"], dtype=np.uint8).copy()
        pix = np.asarray(fr["pix"], dtype=np.int32)
        region = None if regions is None else np.asarray(regions[fi], bool)
        ov_pix, ov_col = [], []
        if can.shape[0] > 0:
            tree = cKDTree(can)
            kk = min(k, can.shape[0])
            dist, idx = tree.query(seed_can, k=kk, workers=-1)
            if kk == 1:
                dist = dist[:, None]
                idx = idx[:, None]
            for ti in range(n):
                m = dist[ti] < threshold
                sel = idx[ti][m]
                if region is not None and sel.size:
                    sel = sel[region[sel]]            # keep only in-region matches
                if sel.size:
                    cols[sel] = palette[ti]            # 3D: recolor the k-NN blob
                    ov_pix.append(pix[sel[:1]])        # 2D overlay: just the nearest point
                    ov_col.append(palette[ti][None])
        per_frame_colors.append(cols)
        if ov_pix:
            per_frame_overlay.append((np.concatenate(ov_pix), np.concatenate(ov_col)))
        else:
            per_frame_overlay.append((np.zeros((0, 2), np.int32), np.zeros((0, 3), np.uint8)))
    return per_frame_colors, per_frame_overlay