"""DINOv3 patch-feature PCA panel.

Project per-patch features (P×P×C) to the first three principal components,
normalize to [0, 255], upsample to the live image size, return a uint8
RGB panel. Shared between training (wandb), inference (rerun), and
paper-figure scripts.

Three APIs in increasing reusability:

- :func:`feature_pca_panel` — fit + project a single map (per-image PCA).
- :func:`shared_feature_pca_panels` — fit one basis across N maps and
  project each (visually comparable colors across N panels).
- :func:`fit_pca_basis` + :func:`apply_pca_basis` — two-step API: fit a
  ``PCABasis`` once on N sample maps, then project arbitrarily-many new
  maps onto the same basis. Required for video pipelines where the basis
  is fit on a small set of sample frames but applied to every frame.
"""
from __future__ import annotations

from dataclasses import dataclass
from typing import Optional, Sequence

import cv2
import numpy as np


def feature_pca_panel(
    patch_features: np.ndarray,    # (P, P, C) numpy
    out_wh: tuple[int, int] = (504, 504),
) -> np.ndarray:
    """Return RGB uint8 panel of the first 3 PCs of patch features (resized)."""
    P, _, C = patch_features.shape
    flat = patch_features.reshape(P * P, C).astype(np.float32)
    flat -= flat.mean(0, keepdims=True)
    _, _, Vt = np.linalg.svd(flat, full_matrices=False)
    pcs = flat @ Vt[:3].T                                  # (P*P, 3)
    pcs -= pcs.min(0, keepdims=True)
    pcs /= (pcs.max(0, keepdims=True) + 1e-8)
    img = (pcs.reshape(P, P, 3) * 255).astype(np.uint8)
    img = cv2.resize(img, out_wh, interpolation=cv2.INTER_NEAREST)
    return img


@dataclass(frozen=True)
class PCABasis:
    """Frozen PCA basis for projecting (P, P, C) feature maps to RGB.

    Produced by :func:`fit_pca_basis`, consumed by :func:`apply_pca_basis`.
    The ``lo``/``hi`` fields define the normalization range so that maps
    projected through the same basis use the **same** color mapping.

    Fields
    ------
    mean : (1, C) — centering vector for incoming features.
    pcs  : (3, C) — top-3 principal components (rows are PCs).
    lo   : (3,)   — per-PC min from the fit population.
    hi   : (3,)   — per-PC max from the fit population.
    """
    mean: np.ndarray
    pcs: np.ndarray
    lo: np.ndarray
    hi: np.ndarray


def fit_pca_basis(
    feats_list: Sequence[np.ndarray],          # each (P, P, C); same P/C
    sample_pixels_per_map: Optional[int] = None,
    seed: int = 0,
) -> PCABasis:
    """Fit a single PCA basis on ``feats_list``.

    SVD is computed on the pooled (optionally subsampled) features; ``lo``
    and ``hi`` are computed on the **full** projected pool so the color
    range is accurate even when subsampling for SVD speed.
    """
    if not feats_list:
        raise ValueError("feats_list is empty")
    P, P2, C = feats_list[0].shape
    assert P == P2, f"expected square feat map, got {P}×{P2}"

    rs = np.random.RandomState(seed)
    flats_sub = []
    for f in feats_list:
        x = f.reshape(P * P, C).astype(np.float32)
        if sample_pixels_per_map is not None and len(x) > sample_pixels_per_map:
            idx = rs.choice(len(x), sample_pixels_per_map, replace=False)
            flats_sub.append(x[idx])
        else:
            flats_sub.append(x)
    pooled_sub = np.concatenate(flats_sub, axis=0)
    mean = pooled_sub.mean(0, keepdims=True)
    _, _, Vt = np.linalg.svd(pooled_sub - mean, full_matrices=False)
    pcs = Vt[:3]                                            # (3, C)

    # lo/hi from full (un-subsampled) projection for accurate color range.
    full_proj = []
    for f in feats_list:
        x = f.reshape(P * P, C).astype(np.float32)
        full_proj.append((x - mean) @ pcs.T)
    pooled_full = np.concatenate(full_proj, axis=0)
    lo = pooled_full.min(0)
    hi = pooled_full.max(0)
    return PCABasis(mean=mean.astype(np.float32),
                    pcs=pcs.astype(np.float32),
                    lo=lo.astype(np.float32),
                    hi=hi.astype(np.float32))


def apply_pca_basis(
    feat_map: np.ndarray,          # (P, P, C)
    basis: PCABasis,
    out_wh: tuple[int, int] = (504, 504),
) -> np.ndarray:
    """Project one feature map onto ``basis``. Returns (out_h, out_w, 3) uint8 RGB.

    Values outside ``[basis.lo, basis.hi]`` are clipped to [0, 1] before
    multiplying to 255 — so out-of-distribution maps still produce a
    valid panel, they just saturate at the basis's color extremes.
    """
    P, P2, C = feat_map.shape
    x = feat_map.reshape(P * P, C).astype(np.float32)
    proj = (x - basis.mean) @ basis.pcs.T                  # (P*P, 3)
    rng = (basis.hi - basis.lo) + 1e-8
    norm = ((proj - basis.lo) / rng).clip(0, 1)
    rgb = (norm.reshape(P, P, 3) * 255).astype(np.uint8)
    return cv2.resize(rgb, out_wh, interpolation=cv2.INTER_NEAREST)


def shared_feature_pca_panels(
    feats_list: Sequence[np.ndarray],
    out_wh: tuple[int, int] = (504, 504),
    sample_pixels_per_map: Optional[int] = None,
    seed: int = 0,
) -> list[np.ndarray]:
    """Convenience: ``fit_pca_basis(...)`` + ``apply_pca_basis`` for every input.

    Use the two-step API directly when the basis needs to outlive the
    fit (e.g., video pipelines that project every frame onto a basis fit
    from a few samples).
    """
    basis = fit_pca_basis(feats_list, sample_pixels_per_map=sample_pixels_per_map,
                           seed=seed)
    return [apply_pca_basis(f, basis, out_wh) for f in feats_list]