"""Per-joint, per-dimension, per-episode signal diagnostics for robot training data.

Provides detailed health checks for action and state trajectories, identifying
dead joints, clipping, directional bias, gripper issues, and per-episode quality.
"""

from __future__ import annotations

import logging
from dataclasses import dataclass, field

import numpy as np

logger = logging.getLogger(__name__)

# ---------------------------------------------------------------------------
# Configuration defaults
# ---------------------------------------------------------------------------

DEFAULT_ZERO_VELOCITY_THRESHOLD = 0.01
DEFAULT_DEAD_JOINT_RATIO = 0.80
DEFAULT_CLIP_FRACTION = 0.20
DEFAULT_CLIP_MARGIN = 0.05
DEFAULT_MIN_EPISODES = 30
DEFAULT_CONSISTENCY_CV = 0.30
DEFAULT_MIN_TRANSITIONS_PER_EPISODE = 1.0
DEFAULT_INSUFFICIENT_TRANSITIONS = 0.5


# ---------------------------------------------------------------------------
# Dataclasses
# ---------------------------------------------------------------------------


@dataclass
class JointDiagnostic:
    """Health diagnostic for a single action dimension (joint)."""

    joint_index: int
    joint_name: str | None
    range_used: tuple[float, float]
    range_limit: tuple[float, float] | None
    range_utilization: float
    zero_velocity_ratio: float
    smoothness: float
    directional_bias: float
    is_dead: bool
    is_clipping: bool
    issues: list[str] = field(default_factory=list)


@dataclass
class GripperAnalysis:
    """Analysis of the gripper action dimension."""

    gripper_dim: int | None
    open_ratio: float
    close_ratio: float
    transitions_per_episode: float
    has_sufficient_grasps: bool
    issues: list[str] = field(default_factory=list)


@dataclass
class EpisodeDiagnostic:
    """Quality diagnostic for a single episode."""

    episode_index: int
    quality_score: float
    length: int
    is_truncated: bool
    jerk_spike_count: int
    jerk_spike_timesteps: list[int] = field(default_factory=list)
    zero_velocity_ratio: float = 0.0
    dead_joints_in_episode: list[int] = field(default_factory=list)
    recommendation: str = "keep"
    issues: list[str] = field(default_factory=list)


@dataclass
class DirectionalBalance:
    """Analysis of directional bias across joints."""

    per_joint_bias: list[float]
    overall_bias: float
    dominant_direction: str | None
    issues: list[str] = field(default_factory=list)


@dataclass
class SignalDiagnosticsReport:
    """Full signal diagnostics report for a robot training dataset."""

    joint_diagnostics: list[JointDiagnostic]
    gripper_analysis: GripperAnalysis
    episode_diagnostics: list[EpisodeDiagnostic]
    directional_balance: DirectionalBalance
    global_zero_velocity_ratio: float
    dead_dimensions: list[int]
    episodes_to_remove: list[int]
    episodes_to_review: list[int]
    blockers: list[str] = field(default_factory=list)
    warnings: list[str] = field(default_factory=list)
    passing: list[str] = field(default_factory=list)
    prescription: list[str] = field(default_factory=list)


# ---------------------------------------------------------------------------
# Helper: normalize variable-length episodes to list of 2D arrays
# ---------------------------------------------------------------------------


def _normalize_episodes(
    actions: np.ndarray | list[np.ndarray],
    episode_lengths: list[int] | None,
) -> list[np.ndarray]:
    """Convert actions into a list of per-episode 2D arrays.

    Args:
        actions: Either shape (num_episodes, max_timesteps, action_dim),
            or a list of arrays each with shape (T_i, action_dim).
        episode_lengths: If actions is a 3D array, the actual length of each
            episode (to handle padding).

    Returns:
        List of 2D arrays, one per episode.
    """
    if isinstance(actions, list):
        result = []
        for a in actions:
            arr = np.asarray(a, dtype=np.float64)
            if arr.ndim == 1:
                arr = arr.reshape(-1, 1)
            result.append(arr)
        return result

    actions = np.asarray(actions, dtype=np.float64)

    if actions.ndim == 1:
        # Single episode, single dimension
        return [actions.reshape(-1, 1)]

    if actions.ndim == 2:
        # Single episode, multiple dimensions (T, D)
        return [actions]

    # 3D: (num_episodes, max_timesteps, action_dim)
    episodes = []
    for i in range(actions.shape[0]):
        ep = actions[i]
        if episode_lengths is not None and i < len(episode_lengths):
            length = episode_lengths[i]
            ep = ep[:length]
        episodes.append(ep)
    return episodes


# ---------------------------------------------------------------------------
# Joint-level diagnostics
# ---------------------------------------------------------------------------


def _compute_joint_diagnostics(
    episodes: list[np.ndarray],
    joint_names: list[str] | None,
    joint_limits: list[tuple[float, float]] | None,
    zero_vel_threshold: float | np.ndarray,
) -> list[JointDiagnostic]:
    """Compute per-joint health diagnostics across all episodes.

    Args:
        episodes: List of per-episode action arrays (T_i, D).
        joint_names: Optional names for each joint dimension.
        joint_limits: Optional (min, max) limits for each joint.
        zero_vel_threshold: Threshold below which abs(action) counts as zero
            velocity.  May be a scalar or a per-dimension array of shape (D,).

    Returns:
        List of JointDiagnostic, one per action dimension.
    """
    if not episodes:
        return []

    action_dim = episodes[0].shape[1] if episodes[0].ndim == 2 else 1

    # Concatenate all episodes for global statistics
    all_actions = np.concatenate(episodes, axis=0)
    if all_actions.ndim == 1:
        all_actions = all_actions.reshape(-1, 1)

    # Allow per-dimension or scalar threshold
    _thresh = np.broadcast_to(
        np.asarray(zero_vel_threshold, dtype=np.float64),
        (action_dim,),
    )

    diagnostics = []
    for dim in range(action_dim):
        dim_thresh = float(_thresh[dim])
        col = all_actions[:, dim]

        # Filter NaN/inf for statistics
        valid = col[np.isfinite(col)]
        if len(valid) == 0:
            diagnostics.append(
                JointDiagnostic(
                    joint_index=dim,
                    joint_name=(
                        joint_names[dim]
                        if joint_names and dim < len(joint_names)
                        else f"joint_{dim}"
                    ),
                    range_used=(0.0, 0.0),
                    range_limit=joint_limits[dim]
                    if joint_limits and dim < len(joint_limits)
                    else None,
                    range_utilization=0.0,
                    zero_velocity_ratio=1.0,
                    smoothness=0.0,
                    directional_bias=0.5,
                    is_dead=True,
                    is_clipping=False,
                    issues=["All values are NaN/inf"],
                )
            )
            continue

        range_min, range_max = float(np.min(valid)), float(np.max(valid))
        range_used = (range_min, range_max)

        # Range utilization
        limit = joint_limits[dim] if joint_limits and dim < len(joint_limits) else None
        if limit is not None and (limit[1] - limit[0]) > 1e-8:
            range_utilization = (range_max - range_min) / (limit[1] - limit[0])
            range_utilization = float(np.clip(range_utilization, 0.0, 1.0))
        elif (range_max - range_min) > 1e-8:
            range_utilization = 1.0  # No limits known, assume full range used
        else:
            range_utilization = 0.0

        # Zero velocity ratio
        zero_vel_mask = np.abs(valid) < dim_thresh
        zero_velocity_ratio = float(np.mean(zero_vel_mask))

        # Smoothness: compute per-episode, then average
        smoothness_values = []
        for ep in episodes:
            ep_col = ep[:, dim] if ep.ndim == 2 else ep
            ep_valid = ep_col[np.isfinite(ep_col)]
            if len(ep_valid) >= 4:
                jerk = np.diff(ep_valid, n=3)
                mean_abs_jerk = np.mean(np.abs(jerk))
                # Normalize by range
                ep_range = np.ptp(ep_valid)
                if ep_range > 1e-8:
                    mean_abs_jerk /= ep_range
                smoothness_values.append(1.0 / (1.0 + mean_abs_jerk))

        smoothness = float(np.mean(smoothness_values)) if smoothness_values else 0.0

        # Directional bias
        positive_count = np.sum(valid > dim_thresh)
        negative_count = np.sum(valid < -dim_thresh)
        total_nonzero = positive_count + negative_count
        if total_nonzero > 0:
            directional_bias = float(max(positive_count, negative_count) / total_nonzero)
        else:
            directional_bias = 0.5

        # Dead joint check — a joint is dead if it never moves.
        # Two detection paths:
        #   1. High zero-velocity ratio AND near-zero range (original)
        #   2. Near-zero range alone (constant non-zero value, e.g.,
        #      a servo stuck at position 512 in raw tick units)
        # A joint with high ZVR but nonzero range is a fixed/locked joint
        # (common in scripted demos and bimanual setups where a base joint
        # is held at a near-constant position).  These are informational,
        # not hardware faults, and should not incur a penalty.
        observed_range_check = range_max - range_min
        is_dead = bool(
            observed_range_check < 1e-8  # constant value — never moves
            or (
                zero_velocity_ratio > DEFAULT_DEAD_JOINT_RATIO
                and observed_range_check < 1e-8
            )
        )

        # Clipping check — skip:
        #  - binary/gripper-like joints (naturally cluster at extremes)
        #  - joints with tiny observed range (< 0.10): these are
        #    near-stationary joints, not actuator-limit clipping
        observed_range = range_max - range_min
        n_unique = len(np.unique(np.round(valid, decimals=2)))
        is_binary_like = n_unique <= 5
        is_small_range = observed_range < 0.10
        if observed_range > 1e-8 and not is_binary_like and not is_small_range:
            margin = observed_range * DEFAULT_CLIP_MARGIN
            near_min = np.sum(valid < range_min + margin)
            near_max = np.sum(valid > range_max - margin)
            clip_fraction = (near_min + near_max) / len(valid)
            is_clipping = bool(clip_fraction > DEFAULT_CLIP_FRACTION)
        else:
            is_clipping = False

        # Build issues
        issues = []
        name = joint_names[dim] if joint_names and dim < len(joint_names) else f"joint_{dim}"
        if is_dead:
            issues.append(f"{name} is dead ({zero_velocity_ratio:.0%} zero velocity)")
        if is_clipping:
            issues.append(f"{name} is clipping (values concentrated at range limits)")
        if directional_bias > 0.90:
            direction = "positive" if positive_count > negative_count else "negative"
            issues.append(
                f"{name} has strong directional bias ({directional_bias:.0%} {direction})"
            )

        diagnostics.append(
            JointDiagnostic(
                joint_index=dim,
                joint_name=name,
                range_used=range_used,
                range_limit=limit,
                range_utilization=range_utilization,
                zero_velocity_ratio=zero_velocity_ratio,
                smoothness=smoothness,
                directional_bias=directional_bias,
                is_dead=is_dead,
                is_clipping=is_clipping,
                issues=issues,
            )
        )

    return diagnostics


# ---------------------------------------------------------------------------
# Gripper analysis
# ---------------------------------------------------------------------------


def _detect_gripper_dim(episodes: list[np.ndarray]) -> int | None:
    """Auto-detect which action dimension is the gripper.

    Heuristic: the gripper dimension has the fewest unique values or the
    most binary-like distribution (values clustered at two levels).
    Prefers the last dimension if multiple candidates tie.

    Args:
        episodes: List of per-episode action arrays.

    Returns:
        Dimension index of the gripper, or None if no binary-like dimension found.
    """
    if not episodes:
        return None

    action_dim = episodes[0].shape[1] if episodes[0].ndim == 2 else 1
    if action_dim < 2:
        return None

    all_actions = np.concatenate(episodes, axis=0)
    if all_actions.ndim == 1:
        return None

    best_dim = None
    best_score = float("inf")

    for dim in range(action_dim):
        col = all_actions[:, dim]
        valid = col[np.isfinite(col)]
        if len(valid) == 0:
            continue

        # Count unique values (rounded to 2 decimals to handle float noise)
        unique_count = len(np.unique(np.round(valid, decimals=2)))

        # Binary-like: check if most values cluster around just 2 levels
        if unique_count <= 5:
            # Strong candidate — prefer last dim for ties
            score = unique_count - (dim * 0.001)  # slight preference for later dims
            if score < best_score:
                best_score = score
                best_dim = dim

    # Only return if the detected dim is reasonably binary-like
    if best_dim is not None and best_score <= 5:
        return best_dim

    # Fallback: detect continuous grippers via bimodal distribution.
    # Many robots (e.g., ALOHA) use continuous gripper values that cluster
    # around two modes (open/closed). Check the last and second-to-last
    # dimensions, which are conventionally where grippers live.
    for dim in [action_dim - 1, action_dim - 2] if action_dim >= 2 else []:
        col = all_actions[:, dim]
        valid = col[np.isfinite(col)]
        if len(valid) < 10:
            continue

        # Check for bimodality: values cluster near min and max of range
        col_min, col_max = float(np.min(valid)), float(np.max(valid))
        col_range = col_max - col_min
        if col_range < 1e-6:
            continue
        normalized = (valid - col_min) / col_range  # [0, 1]
        near_low = float(np.mean(normalized < 0.25))
        near_high = float(np.mean(normalized > 0.75))
        in_middle = float(np.mean((normalized >= 0.25) & (normalized <= 0.75)))

        # Bimodal: most values at extremes, few in the middle
        if near_low + near_high > 0.70 and in_middle < 0.40:
            return dim

    return None


def _compute_gripper_analysis(
    episodes: list[np.ndarray],
    gripper_dim: int | None = None,
) -> GripperAnalysis:
    """Analyze gripper behavior across episodes.

    Args:
        episodes: List of per-episode action arrays.
        gripper_dim: Known gripper dimension, or None to auto-detect.

    Returns:
        GripperAnalysis with open/close ratios and transition counts.
    """
    if gripper_dim is None:
        gripper_dim = _detect_gripper_dim(episodes)

    if gripper_dim is None or not episodes:
        return GripperAnalysis(
            gripper_dim=None,
            open_ratio=0.0,
            close_ratio=0.0,
            transitions_per_episode=0.0,
            has_sufficient_grasps=False,
            issues=["No gripper dimension detected"],
        )

    all_gripper = []
    transitions_list = []

    for ep in episodes:
        if ep.ndim == 1:
            continue
        if gripper_dim >= ep.shape[1]:
            continue

        g = ep[:, gripper_dim]
        valid = g[np.isfinite(g)]
        if len(valid) == 0:
            continue

        all_gripper.append(valid)

        # Count transitions: threshold at median value
        median_val = np.median(valid)
        binary = (valid > median_val).astype(int)
        transitions = int(np.sum(np.abs(np.diff(binary))))
        transitions_list.append(transitions)

    if not all_gripper:
        return GripperAnalysis(
            gripper_dim=gripper_dim,
            open_ratio=0.0,
            close_ratio=0.0,
            transitions_per_episode=0.0,
            has_sufficient_grasps=False,
            issues=["No valid gripper data found"],
        )

    concat = np.concatenate(all_gripper)
    median_val = np.median(concat)
    open_ratio = float(np.mean(concat > median_val))
    close_ratio = 1.0 - open_ratio

    transitions_per_episode = float(np.mean(transitions_list)) if transitions_list else 0.0
    has_sufficient_grasps = transitions_per_episode >= DEFAULT_MIN_TRANSITIONS_PER_EPISODE

    issues = []
    if not has_sufficient_grasps:
        issues.append(
            f"Insufficient grasp transitions ({transitions_per_episode:.1f}/episode, "
            f"need >= {DEFAULT_MIN_TRANSITIONS_PER_EPISODE})"
        )
    if open_ratio > 0.9:
        issues.append(f"Gripper mostly open ({open_ratio:.0%} of time)")
    elif close_ratio > 0.9:
        issues.append(f"Gripper mostly closed ({close_ratio:.0%} of time)")

    return GripperAnalysis(
        gripper_dim=gripper_dim,
        open_ratio=open_ratio,
        close_ratio=close_ratio,
        transitions_per_episode=transitions_per_episode,
        has_sufficient_grasps=has_sufficient_grasps,
        issues=issues,
    )


# ---------------------------------------------------------------------------
# Episode-level diagnostics
# ---------------------------------------------------------------------------


def _compute_episode_diagnostics(
    episodes: list[np.ndarray],
    zero_vel_threshold: float | np.ndarray,
) -> list[EpisodeDiagnostic]:
    """Compute per-episode quality diagnostics.

    Args:
        episodes: List of per-episode action arrays.
        zero_vel_threshold: Threshold for zero velocity detection.
            May be a scalar or per-dimension array of shape (D,).

    Returns:
        List of EpisodeDiagnostic, one per episode.
    """
    if not episodes:
        return []

    lengths = [len(ep) for ep in episodes]
    median_length = float(np.median(lengths)) if lengths else 0.0
    action_dim = episodes[0].shape[1] if episodes[0].ndim == 2 else 1

    diagnostics = []
    for i, ep in enumerate(episodes):
        if ep.ndim == 1:
            ep = ep.reshape(-1, 1)

        length = len(ep)
        is_truncated = length < 0.5 * median_length if median_length > 0 else False

        # Replace NaN/inf for analysis
        ep_clean = np.nan_to_num(ep, nan=0.0, posinf=0.0, neginf=0.0)

        # Zero velocity ratio for this episode
        zero_vel = np.abs(ep_clean) < zero_vel_threshold
        ep_zero_vel_ratio = float(np.mean(zero_vel))

        # Dead joints in this episode
        _ep_thresh = np.broadcast_to(
            np.asarray(zero_vel_threshold, dtype=np.float64),
            (ep_clean.shape[1],),
        )
        dead_joints = []
        for dim in range(ep_clean.shape[1]):
            dim_zero = float(np.mean(np.abs(ep_clean[:, dim]) < _ep_thresh[dim]))
            if dim_zero > DEFAULT_DEAD_JOINT_RATIO:
                dead_joints.append(dim)

        # Jerk spikes
        jerk_spike_count = 0
        jerk_spike_timesteps = []
        if length >= 4:
            jerk = np.diff(ep_clean, n=3, axis=0)
            jerk_magnitude = np.linalg.norm(jerk, axis=1)
            if len(jerk_magnitude) > 1:
                mean_jerk = np.mean(jerk_magnitude)
                std_jerk = np.std(jerk_magnitude)
                if std_jerk > 1e-8:
                    spike_mask = jerk_magnitude > mean_jerk + 3 * std_jerk
                    jerk_spike_count = int(np.sum(spike_mask))
                    jerk_spike_timesteps = [int(t) for t in np.where(spike_mask)[0]]

        # Smoothness for this episode
        if length >= 4:
            jerk_all = np.diff(ep_clean, n=3, axis=0)
            mean_abs_jerk = np.mean(np.abs(jerk_all))
            ep_range = np.mean(np.ptp(ep_clean, axis=0))
            if ep_range > 1e-8:
                mean_abs_jerk /= ep_range
            smoothness = 1.0 / (1.0 + mean_abs_jerk)
        else:
            smoothness = 0.5

        # Active joints ratio
        active_joints = action_dim - len(dead_joints)
        active_ratio = active_joints / max(1, action_dim)

        # Quality score
        quality_score = (
            (1 - ep_zero_vel_ratio) * 0.3
            + smoothness * 0.3
            + (1 - jerk_spike_count / max(1, length)) * 0.2
            + active_ratio * 0.2
        )
        quality_score = float(np.clip(quality_score, 0.0, 1.0))

        # Recommendation
        issues = []
        if is_truncated:
            issues.append(f"Truncated ({length} frames, median is {median_length:.0f})")
        # Scale jerk spike threshold with episode length: longer episodes
        # naturally have more spikes. Use max(5, 3% of episode length).
        jerk_spike_threshold = max(5, int(length * 0.03))
        if jerk_spike_count > jerk_spike_threshold:
            issues.append(f"{jerk_spike_count} jerk spikes detected (threshold {jerk_spike_threshold})")
        if len(dead_joints) > action_dim * 0.5:
            issues.append(f"{len(dead_joints)}/{action_dim} joints dead in this episode")
        if ep_zero_vel_ratio > 0.8:
            issues.append(f"High zero velocity ({ep_zero_vel_ratio:.0%})")

        if is_truncated or quality_score < 0.3 or len(dead_joints) > action_dim * 0.5:
            recommendation = "remove"
        elif jerk_spike_count > jerk_spike_threshold or quality_score < 0.5:
            recommendation = "review"
        else:
            recommendation = "keep"

        diagnostics.append(
            EpisodeDiagnostic(
                episode_index=i,
                quality_score=quality_score,
                length=length,
                is_truncated=is_truncated,
                jerk_spike_count=jerk_spike_count,
                jerk_spike_timesteps=jerk_spike_timesteps,
                zero_velocity_ratio=ep_zero_vel_ratio,
                dead_joints_in_episode=dead_joints,
                recommendation=recommendation,
                issues=issues,
            )
        )

    return diagnostics


# ---------------------------------------------------------------------------
# Directional balance
# ---------------------------------------------------------------------------


def _compute_directional_balance(
    episodes: list[np.ndarray],
    zero_vel_threshold: float | np.ndarray,
) -> DirectionalBalance:
    """Analyze directional bias across all joints.

    Args:
        episodes: List of per-episode action arrays.
        zero_vel_threshold: Threshold for counting positive/negative actions.
            May be a scalar or per-dimension array of shape (D,).

    Returns:
        DirectionalBalance with per-joint and overall bias metrics.
    """
    if not episodes:
        return DirectionalBalance(
            per_joint_bias=[],
            overall_bias=0.5,
            dominant_direction=None,
        )

    all_actions = np.concatenate(episodes, axis=0)
    if all_actions.ndim == 1:
        all_actions = all_actions.reshape(-1, 1)

    action_dim = all_actions.shape[1]
    _dir_thresh = np.broadcast_to(
        np.asarray(zero_vel_threshold, dtype=np.float64),
        (action_dim,),
    )
    per_joint_bias = []
    positive_total = 0
    negative_total = 0

    for dim in range(action_dim):
        col = all_actions[:, dim]
        valid = col[np.isfinite(col)]
        pos = np.sum(valid > _dir_thresh[dim])
        neg = np.sum(valid < -_dir_thresh[dim])
        total = pos + neg
        if total > 0:
            bias = float(max(pos, neg) / total)
        else:
            bias = 0.5
        per_joint_bias.append(bias)
        positive_total += pos
        negative_total += neg

    overall_bias = float(np.mean(per_joint_bias))

    # Determine dominant direction
    dominant_direction = None
    if overall_bias > 0.65:
        if positive_total > negative_total:
            dominant_direction = "mostly positive/forward"
        else:
            dominant_direction = "mostly negative/backward"

    issues = []
    for dim, bias in enumerate(per_joint_bias):
        if bias > 0.90:
            issues.append(f"Joint {dim} has strong directional bias ({bias:.0%})")

    return DirectionalBalance(
        per_joint_bias=per_joint_bias,
        overall_bias=overall_bias,
        dominant_direction=dominant_direction,
        issues=issues,
    )


# ---------------------------------------------------------------------------
# Blockers / Warnings / Passing / Prescription
# ---------------------------------------------------------------------------


def _classify_findings(
    joint_diags: list[JointDiagnostic],
    gripper: GripperAnalysis,
    episode_diags: list[EpisodeDiagnostic],
    balance: DirectionalBalance,
    global_zero_vel: float,
    dead_dims: list[int],
    min_episodes: int,
) -> tuple[list[str], list[str], list[str], list[str]]:
    """Classify findings into blockers, warnings, passing, and prescriptions.

    Returns:
        (blockers, warnings, passing, prescription)
    """
    blockers = []
    warnings = []
    passing = []
    prescription = []

    num_episodes = len(episode_diags)
    remove_count = sum(1 for e in episode_diags if e.recommendation == "remove")
    remove_ratio = remove_count / max(1, num_episodes)

    # BLOCKERS
    if dead_dims:
        joint_names = []
        for d in dead_dims:
            jd = joint_diags[d] if d < len(joint_diags) else None
            name = jd.joint_name if jd else f"joint_{d}"
            joint_names.append(name)
        blockers.append(
            f"Dead dimensions: {', '.join(joint_names)} "
            f"({len(dead_dims)} of {len(joint_diags)} joints have >80% zero velocity)"
        )
        prescription.append(
            f"Record {min_episodes} episodes with more {', '.join(joint_names)} movement"
        )

    if (
        gripper.gripper_dim is not None
        and gripper.transitions_per_episode < DEFAULT_INSUFFICIENT_TRANSITIONS
    ):
        blockers.append(
            f"Gripper transitions too low ({gripper.transitions_per_episode:.1f}/episode, "
            f"need >= {DEFAULT_INSUFFICIENT_TRANSITIONS})"
        )
        prescription.append(
            "Record episodes with explicit grasp-release cycles "
            f"(currently {gripper.transitions_per_episode:.1f} transitions/episode)"
        )

    if remove_ratio > 0.20:
        remove_indices = [e.episode_index for e in episode_diags if e.recommendation == "remove"]
        preview = remove_indices[:5]
        suffix = f" and {len(remove_indices) - 5} more" if len(remove_indices) > 5 else ""
        blockers.append(
            f"{remove_count}/{num_episodes} episodes ({remove_ratio:.0%}) flagged for removal "
            f"(episodes {', '.join(str(i) for i in preview)}{suffix})"
        )
        prescription.append(
            f"Remove episodes {', '.join(str(i) for i in remove_indices)} before training "
            f"(truncated/jerky)"
        )

    # DIRECTIONAL BIAS WARNINGS
    # When most joints show directional bias, it indicates goal-directed motion
    # (e.g., pick-and-place, insertion) where bias is expected and normal.
    # Only flag as warnings when a small minority of joints are biased, which
    # suggests a hardware or collection issue on specific joints.
    biased_joints = [jd for jd in joint_diags if jd.directional_bias > 0.90 and not jd.is_dead]
    active_joints = [jd for jd in joint_diags if not jd.is_dead]
    biased_ratio = len(biased_joints) / max(1, len(active_joints))

    if biased_ratio >= 0.5:
        # Majority of joints biased → goal-directed task, NOT a data problem.
        # Emit a single informational note, not per-joint warnings.
        passing.append(
            f"Directional bias on {len(biased_joints)}/{len(active_joints)} joints "
            f"— consistent with goal-directed task (expected)"
        )
    elif biased_joints:
        # Minority of joints biased → possible collection issue on those joints
        for jd in biased_joints:
            warnings.append(f"{jd.joint_name} has directional bias ({jd.directional_bias:.0%})")
        if len(biased_joints) <= max(2, len(joint_diags) // 3):
            for jd in biased_joints:
                prescription.append(
                    f"Add {min_episodes // 3} demos with varied {jd.joint_name} direction "
                    f"({jd.directional_bias:.0%} of current demos go one direction)"
                )

    if global_zero_vel > 0.40:
        warnings.append(
            f"High global zero velocity ({global_zero_vel:.0%}) — many timesteps with no movement"
        )

    clipping_joints = [jd for jd in joint_diags if jd.is_clipping]
    if clipping_joints:
        names = [jd.joint_name or f"joint_{jd.joint_index}" for jd in clipping_joints]
        warnings.append(f"Clipping detected on {', '.join(names)}")

    # PASSING
    if num_episodes >= min_episodes:
        passing.append(f"Episode count ({num_episodes}) meets minimum ({min_episodes})")
    else:
        warnings.append(f"Episode count ({num_episodes}) below minimum ({min_episodes})")
        prescription.append(f"Record {min_episodes - num_episodes} more episodes to meet minimum")

    if num_episodes >= 2:
        lengths = [e.length for e in episode_diags]
        length_cv = float(np.std(lengths) / np.mean(lengths)) if np.mean(lengths) > 0 else 0
        if length_cv < DEFAULT_CONSISTENCY_CV:
            passing.append(f"Episode lengths consistent (CV={length_cv:.2f})")
        else:
            warnings.append(f"Episode lengths inconsistent (CV={length_cv:.2f})")

    avg_smoothness = float(np.mean([jd.smoothness for jd in joint_diags])) if joint_diags else 0
    if avg_smoothness > 0.6:
        passing.append(f"Average smoothness good ({avg_smoothness:.2f})")
    elif avg_smoothness > 0.3:
        warnings.append(f"Average smoothness moderate ({avg_smoothness:.2f})")

    if not dead_dims:
        passing.append("All joints active")

    if balance.overall_bias < 0.65:
        passing.append(f"Directional balance good (bias={balance.overall_bias:.2f})")

    return blockers, warnings, passing, prescription


# ---------------------------------------------------------------------------
# Main entry point
# ---------------------------------------------------------------------------


def compute_signal_diagnostics(
    actions: np.ndarray | list[np.ndarray],
    states: np.ndarray | None = None,
    joint_names: list[str] | None = None,
    joint_limits: list[tuple[float, float]] | None = None,
    episode_lengths: list[int] | None = None,
    zero_vel_threshold: float = DEFAULT_ZERO_VELOCITY_THRESHOLD,
    min_episodes: int = DEFAULT_MIN_EPISODES,
) -> SignalDiagnosticsReport:
    """Compute comprehensive signal diagnostics for robot training data.

    Args:
        actions: Action data — either shape (num_episodes, max_timesteps, action_dim),
            a list of per-episode arrays, or a single 2D array (T, D).
        states: Optional state data in the same format as actions.
        joint_names: Optional human-readable names for each action dimension.
        joint_limits: Optional (min, max) tuples for each action dimension.
        episode_lengths: Actual lengths per episode when actions is a padded 3D array.
        zero_vel_threshold: Absolute threshold for zero velocity detection.
            When left at the default, an adaptive per-dimension threshold is
            computed as 1% of each dimension's range (with a floor of 1e-10).
        min_episodes: Minimum number of episodes expected.

    Returns:
        SignalDiagnosticsReport with full diagnostics.
    """
    logger.info("Computing signal diagnostics")

    # Normalize to list of 2D episode arrays
    episodes = _normalize_episodes(actions, episode_lengths)

    if not episodes:
        logger.warning("No episodes provided for signal diagnostics")
        return SignalDiagnosticsReport(
            joint_diagnostics=[],
            gripper_analysis=GripperAnalysis(
                gripper_dim=None,
                open_ratio=0.0,
                close_ratio=0.0,
                transitions_per_episode=0.0,
                has_sufficient_grasps=False,
                issues=["No data"],
            ),
            episode_diagnostics=[],
            directional_balance=DirectionalBalance(
                per_joint_bias=[],
                overall_bias=0.5,
                dominant_direction=None,
            ),
            global_zero_velocity_ratio=0.0,
            dead_dimensions=[],
            episodes_to_remove=[],
            episodes_to_review=[],
            blockers=["No episodes to analyze"],
            warnings=[],
            passing=[],
            prescription=["Collect training episodes"],
        )

    # Filter out episodes with NaN/inf only (keep them but clean)
    valid_episodes = []
    for ep in episodes:
        if ep.size == 0:
            valid_episodes.append(ep)
            continue
        # Check if the entire episode is NaN/inf
        finite_mask = np.isfinite(ep)
        if not np.any(finite_mask):
            logger.warning("Episode entirely NaN/inf, keeping with zeros")
            valid_episodes.append(np.zeros_like(ep))
        else:
            valid_episodes.append(ep)
    episodes = valid_episodes

    # Compute adaptive per-dimension threshold when the caller did not
    # override the default.  Action scales vary 100,000x across robots
    # (delta-EEF ~0.001 vs raw servo ticks ~4095), so a single absolute
    # number silently misclassifies joints on most platforms.
    if zero_vel_threshold == DEFAULT_ZERO_VELOCITY_THRESHOLD:
        all_for_range = np.concatenate(episodes, axis=0)
        if all_for_range.ndim == 1:
            all_for_range = all_for_range.reshape(-1, 1)
        all_for_range = np.nan_to_num(all_for_range, nan=0.0, posinf=0.0, neginf=0.0)
        action_range = np.ptp(all_for_range, axis=0)  # per-dimension range
        zero_vel_threshold = np.maximum(action_range * 0.01, 1e-10)

    # Joint diagnostics
    joint_diags = _compute_joint_diagnostics(
        episodes, joint_names, joint_limits, zero_vel_threshold
    )

    # Gripper analysis
    gripper = _compute_gripper_analysis(episodes)

    # Episode diagnostics
    episode_diags = _compute_episode_diagnostics(episodes, zero_vel_threshold)

    # Directional balance
    balance = _compute_directional_balance(episodes, zero_vel_threshold)

    # Global metrics
    all_actions = np.concatenate(episodes, axis=0)
    if all_actions.ndim == 1:
        all_actions = all_actions.reshape(-1, 1)
    all_clean = np.nan_to_num(all_actions, nan=0.0, posinf=0.0, neginf=0.0)
    global_zero_vel = float(np.mean(np.abs(all_clean) < zero_vel_threshold))

    dead_dims = [jd.joint_index for jd in joint_diags if jd.is_dead]
    episodes_to_remove = [e.episode_index for e in episode_diags if e.recommendation == "remove"]
    episodes_to_review = [e.episode_index for e in episode_diags if e.recommendation == "review"]

    # Classify findings
    blockers, warnings, passing, prescription = _classify_findings(
        joint_diags,
        gripper,
        episode_diags,
        balance,
        global_zero_vel,
        dead_dims,
        min_episodes,
    )

    logger.info(
        "Signal diagnostics complete: %d blockers, %d warnings, %d passing",
        len(blockers),
        len(warnings),
        len(passing),
    )

    return SignalDiagnosticsReport(
        joint_diagnostics=joint_diags,
        gripper_analysis=gripper,
        episode_diagnostics=episode_diags,
        directional_balance=balance,
        global_zero_velocity_ratio=global_zero_vel,
        dead_dimensions=dead_dims,
        episodes_to_remove=episodes_to_remove,
        episodes_to_review=episodes_to_review,
        blockers=blockers,
        warnings=warnings,
        passing=passing,
        prescription=prescription,
    )
