"""Universal array profiler — infers what a robotics data array represents.

Given a raw numpy array of shape (num_timesteps, num_dimensions), this module
infers everything about it using statistical properties alone — no column names,
no metadata, no hardcoded format-specific logic.

Statistical signatures used:
- Autocorrelation: position (>0.95) vs velocity (<0.5)
- Value range: [-pi, pi] → joint angles, [-1, 1] → normalized
- Distribution shape: bimodal → gripper, gaussian → velocity
- Cross-dimension correlation: kinematic chain detection
- Frequency content: FFT power spectrum for smoothness
- Integration test: cumsum of deltas → bounded trajectory

References:
- Belkhale et al. (NeurIPS 2023): action divergence as primary quality axis
- Mandlekar et al. (CoRL 2021): demonstration quality > quantity
- Sakr et al. (2025): consistency CVs explain 70-91% of success variance
"""

from __future__ import annotations

import logging
from dataclasses import dataclass, field

import numpy as np

logger = logging.getLogger(__name__)


# ---------------------------------------------------------------------------
# Data models
# ---------------------------------------------------------------------------


@dataclass
class DimensionRole:
    """Inferred role and confidence for a single data dimension."""

    role: str  # joint_angle, eef_position, eef_rotation, velocity,
    # gripper, torque, padding, unknown
    confidence: float  # 0-1
    evidence: str  # human-readable explanation of why this role was inferred

    def __repr__(self) -> str:
        return f"DimensionRole({self.role}, conf={self.confidence:.2f})"


@dataclass
class DimensionProfile:
    """Statistical profile for a single dimension."""

    index: int
    mean: float
    std: float
    min_val: float
    max_val: float
    range_width: float
    lag1_autocorrelation: float
    is_bimodal: bool
    bimodality_coefficient: float
    unique_ratio: float  # fraction of unique values
    zero_crossing_rate: float
    spectral_low_power_ratio: float  # fraction of power in low frequencies
    kurtosis: float
    skewness: float
    is_dead: bool
    is_clipping: bool

    # Inferred role
    role: DimensionRole = field(default_factory=lambda: DimensionRole("unknown", 0.0, ""))


@dataclass
class ArrayProfile:
    """What we can infer about a data array from statistics alone."""

    # Per-dimension analysis
    dimension_profiles: list[DimensionProfile]
    dimension_roles: list[DimensionRole]

    # Array-level properties
    action_type: str  # absolute_position, delta, velocity, torque, normalized, unknown
    action_type_confidence: float  # 0-1
    is_normalized: bool  # values in [-1, 1]
    scale: float  # characteristic magnitude (median absolute value)
    estimated_frequency_hz: float  # estimated from autocorrelation
    smoothness: float  # 0-1, how smooth the trajectories are

    # Quality signals
    dead_dimensions: list[int]  # constant or near-constant
    clipping_dimensions: list[int]  # saturated at min/max
    noisy_dimensions: list[int]  # high-frequency noise dominates
    gripper_dimensions: list[int]  # binary/bimodal behavior

    # Structural inference
    num_dimensions: int
    is_bimanual: bool  # detected dual-arm structure
    kinematic_groups: list[list[int]]  # groups of correlated dimensions

    # Overall confidence
    confidence: float  # 0-1, how confident we are in the full inference


# ---------------------------------------------------------------------------
# Per-dimension statistical analysis
# ---------------------------------------------------------------------------


def _lag1_autocorrelation(x: np.ndarray) -> float:
    """Compute lag-1 autocorrelation for a 1D signal."""
    if len(x) < 3:
        return 0.0
    x_centered = x - np.mean(x)
    var = np.sum(x_centered**2)
    if var < 1e-15:
        return 0.0
    cov = np.sum(x_centered[:-1] * x_centered[1:])
    return float(np.clip(cov / var, -1.0, 1.0))


def _bimodality_coefficient(x: np.ndarray) -> float:
    """Compute bimodality coefficient B = (skewness^2 + 1) / kurtosis.

    B > 0.555 suggests bimodal distribution (Pfister et al. 2013).
    """
    n = len(x)
    if n < 4:
        return 0.0
    mean = np.mean(x)
    std = np.std(x, ddof=1)
    if std < 1e-15:
        return 0.0
    z = (x - mean) / std
    skew = float(np.mean(z**3))
    kurt = float(np.mean(z**4))
    if kurt < 1e-15:
        return 0.0
    return (skew**2 + 1) / kurt


def _spectral_low_power_ratio(x: np.ndarray, low_fraction: float = 0.25) -> float:
    """Fraction of total spectral power in the bottom frequencies."""
    if len(x) < 16:
        return 0.5
    fft = np.fft.rfft(x - np.mean(x))
    power = np.abs(fft) ** 2
    total = np.sum(power)
    if total < 1e-15:
        return 1.0  # constant signal
    n_low = max(1, int(len(power) * low_fraction))
    return float(np.sum(power[:n_low]) / total)


def _zero_crossing_rate(x: np.ndarray) -> float:
    """Rate of zero-crossings (relative to mean)."""
    if len(x) < 2:
        return 0.0
    centered = x - np.mean(x)
    crossings = np.sum(np.abs(np.diff(np.sign(centered))) > 0)
    return float(crossings / (len(x) - 1))


def _is_dead_dimension(x: np.ndarray, episodes: list[np.ndarray] | None, dim: int) -> bool:
    """Check if a dimension is effectively constant (dead joint).

    Uses the per-dimension range relative to the overall action scale,
    matching the approach in signal_diagnostics.py (relative thresholds).
    """
    range_width = float(np.ptp(x))
    std = float(np.std(x))

    # Absolute check: truly constant
    if range_width < 1e-10:
        return True

    # Relative check: range is tiny compared to mean magnitude
    magnitude = max(abs(float(np.mean(x))), abs(float(np.median(x))), 1e-10)
    if range_width / magnitude < 0.001 and std / magnitude < 0.0005:
        return True

    return False


def _is_clipping(x: np.ndarray, margin_fraction: float = 0.02) -> bool:
    """Check if values are saturating at boundaries."""
    if len(x) < 10:
        return False
    vmin, vmax = float(np.min(x)), float(np.max(x))
    range_width = vmax - vmin
    if range_width < 1e-10:
        return False
    margin = range_width * margin_fraction
    at_min = np.sum(np.abs(x - vmin) < margin) / len(x)
    at_max = np.sum(np.abs(x - vmax) < margin) / len(x)
    # Clipping if >10% of samples at either boundary
    return float(at_min) > 0.10 or float(at_max) > 0.10


def _profile_dimension(
    all_data: np.ndarray,
    dim: int,
    episodes: list[np.ndarray] | None = None,
) -> DimensionProfile:
    """Compute full statistical profile for one dimension."""
    x = all_data[:, dim] if all_data.ndim > 1 else all_data

    mean = float(np.mean(x))
    std = float(np.std(x))
    min_val = float(np.min(x))
    max_val = float(np.max(x))
    range_width = max_val - min_val

    # Autocorrelation (on first episode if available, else full)
    if episodes and len(episodes[0]) > 10:
        ep0 = episodes[0][:, dim] if episodes[0].ndim > 1 else episodes[0]
        autocorr = _lag1_autocorrelation(ep0)
    else:
        autocorr = _lag1_autocorrelation(x)

    # Distribution shape
    bimod_coeff = _bimodality_coefficient(x)
    is_bimodal = bimod_coeff > 0.555

    # Unique values ratio (for detecting discrete/gripper)
    n_unique = len(np.unique(np.round(x, 4)))
    unique_ratio = n_unique / max(len(x), 1)

    # Spectral properties
    spectral_ratio = _spectral_low_power_ratio(x)

    # Kurtosis and skewness
    if std > 1e-10:
        z = (x - mean) / std
        kurtosis = float(np.mean(z**4)) - 3.0  # excess kurtosis
        skewness = float(np.mean(z**3))
    else:
        kurtosis = 0.0
        skewness = 0.0

    is_dead = _is_dead_dimension(x, episodes, dim)
    is_clip = _is_clipping(x)

    return DimensionProfile(
        index=dim,
        mean=mean,
        std=std,
        min_val=min_val,
        max_val=max_val,
        range_width=range_width,
        lag1_autocorrelation=autocorr,
        is_bimodal=is_bimodal,
        bimodality_coefficient=bimod_coeff,
        unique_ratio=unique_ratio,
        zero_crossing_rate=_zero_crossing_rate(x),
        spectral_low_power_ratio=spectral_ratio,
        kurtosis=kurtosis,
        skewness=skewness,
        is_dead=is_dead,
        is_clipping=is_clip,
    )


# ---------------------------------------------------------------------------
# Dimension role inference
# ---------------------------------------------------------------------------


def _count_transitions(episodes: list[np.ndarray], dim: int) -> float:
    """Count average transitions per episode for a dimension.

    A "transition" is when the value crosses the midpoint between its
    two modes. Real grippers have 1-5 transitions per episode.
    Joint angles that happen to be bimodal have many more.
    """
    if not episodes:
        return 100.0

    transition_counts = []
    for ep in episodes[:20]:  # sample up to 20 episodes
        col = ep[:, dim] if ep.ndim > 1 else ep
        if len(col) < 5:
            continue
        midpoint = (float(np.min(col)) + float(np.max(col))) / 2.0
        above = col > midpoint
        transitions = int(np.sum(np.abs(np.diff(above.astype(int)))))
        transition_counts.append(transitions)

    return float(np.mean(transition_counts)) if transition_counts else 100.0


def _infer_dimension_role(
    profile: DimensionProfile,
    all_profiles: list[DimensionProfile],
    episodes: list[np.ndarray] | None = None,
) -> DimensionRole:
    """Infer what a single dimension represents from its statistics."""

    # Dead / padding — highest priority
    if profile.is_dead:
        return DimensionRole("padding", 0.95, f"constant (range={profile.range_width:.6f})")

    # Gripper detection — STRONGLY bimodal + few transitions per episode
    # Real grippers: bimodality coefficient >0.70 (strong two-mode distribution)
    # AND few transitions (1-5 open/close events per episode).
    # Joint angles with B=0.55-0.70 are just slightly bimodal motion patterns
    # (start position vs end position), not discrete gripper commands.
    # The 0.555 threshold from Pfister et al. is designed for general stats;
    # for robotics action spaces we need higher (0.70) to avoid false positives.
    is_strongly_bimodal = profile.bimodality_coefficient > 0.70
    if is_strongly_bimodal and episodes:
        avg_transitions = _count_transitions(episodes, profile.index)
        if avg_transitions <= 8 and profile.unique_ratio < 0.15:
            return DimensionRole(
                "gripper", 0.90,
                f"strongly bimodal (B={profile.bimodality_coefficient:.3f}), "
                f"few transitions ({avg_transitions:.1f}/ep), "
                f"unique_ratio={profile.unique_ratio:.3f}",
            )
        if avg_transitions <= 5:
            return DimensionRole(
                "gripper", 0.75,
                f"bimodal (B={profile.bimodality_coefficient:.3f}), "
                f"very few transitions ({avg_transitions:.1f}/ep)",
            )
        # Strongly bimodal but many transitions → not a gripper

    # Position-like (high autocorrelation, smooth)
    if profile.lag1_autocorrelation > 0.95:
        # Distinguish joint angle from EEF position by range
        if _is_joint_angle_range(profile):
            return DimensionRole(
                "joint_angle", 0.80,
                f"high autocorr={profile.lag1_autocorrelation:.3f}, "
                f"range [{profile.min_val:.2f}, {profile.max_val:.2f}] ~ [-pi, pi]",
            )
        elif abs(profile.max_val) < 2.0 and abs(profile.min_val) < 2.0:
            return DimensionRole(
                "eef_position", 0.70,
                f"high autocorr={profile.lag1_autocorrelation:.3f}, "
                f"small range [{profile.min_val:.3f}, {profile.max_val:.3f}]",
            )
        else:
            return DimensionRole(
                "joint_angle", 0.65,
                f"high autocorr={profile.lag1_autocorrelation:.3f}, "
                f"range [{profile.min_val:.2f}, {profile.max_val:.2f}]",
            )

    # Velocity-like (moderate autocorrelation, zero-centered)
    if profile.lag1_autocorrelation < 0.50 and abs(profile.mean) < 0.5 * profile.std:
        if profile.kurtosis > 3.0:
            return DimensionRole(
                "torque", 0.60,
                f"low autocorr={profile.lag1_autocorrelation:.3f}, "
                f"zero-centered, high kurtosis={profile.kurtosis:.2f}",
            )
        return DimensionRole(
            "velocity", 0.70,
            f"low autocorr={profile.lag1_autocorrelation:.3f}, "
            f"zero-centered (mean={profile.mean:.4f}, std={profile.std:.4f})",
        )

    # Delta-position (moderate autocorrelation, small values)
    if 0.50 <= profile.lag1_autocorrelation <= 0.95:
        if abs(profile.mean) < profile.std:
            return DimensionRole(
                "velocity", 0.55,
                f"moderate autocorr={profile.lag1_autocorrelation:.3f}, "
                f"near-zero mean",
            )
        return DimensionRole(
            "joint_angle", 0.50,
            f"moderate autocorr={profile.lag1_autocorrelation:.3f}",
        )

    return DimensionRole("unknown", 0.30, "no clear statistical signature")


def _is_joint_angle_range(profile: DimensionProfile) -> bool:
    """Check if value range is consistent with joint angles (~[-pi, pi])."""
    pi = np.pi
    # Standard revolute joint ranges
    if profile.min_val >= -pi - 0.5 and profile.max_val <= pi + 0.5:
        return True
    if profile.min_val >= -2 * pi - 0.5 and profile.max_val <= 2 * pi + 0.5:
        return True
    # Dynamixel-like ranges (0-4095, 0-1023)
    if profile.min_val >= -0.1 and profile.max_val <= 4100:
        if profile.range_width > 100:
            return True
    return False


# ---------------------------------------------------------------------------
# Array-level inference
# ---------------------------------------------------------------------------


def _detect_action_type(
    episodes: list[np.ndarray],
    profiles: list[DimensionProfile],
) -> tuple[str, float]:
    """Infer whether actions are absolute positions, deltas, velocities, etc.

    Strategy:
    1. Check if values are normalized (all in [-1, 1])
    2. Check if zero-centered with small magnitude → delta/velocity
    3. Check if cumsum produces bounded, smooth trajectory → delta
    4. Check autocorrelation signature → position vs velocity
    """
    if not episodes:
        return "unknown", 0.0

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

    ndim = all_data.shape[1]

    # Skip dead/gripper dimensions for type detection
    active_dims = [
        i for i in range(ndim)
        if not profiles[i].is_dead and profiles[i].role.role != "gripper"
    ]
    if not active_dims:
        return "unknown", 0.0

    active_data = all_data[:, active_dims]
    active_profiles = [profiles[i] for i in active_dims]

    # Check 1: Normalized actions (all values in [-1, 1])
    if np.all(np.abs(active_data) <= 1.05):
        mean_autocorr = np.mean([p.lag1_autocorrelation for p in active_profiles])
        if mean_autocorr > 0.90:
            return "normalized", 0.85
        elif mean_autocorr < 0.50:
            return "normalized", 0.80  # normalized velocity
        return "normalized", 0.75

    # Check 2: Zero-centered small values → delta/velocity
    means = np.array([p.mean for p in active_profiles])
    stds = np.array([p.std for p in active_profiles])
    is_zero_centered = np.mean(np.abs(means) < 0.5 * stds) > 0.7

    if is_zero_centered:
        # Integration test: does cumsum produce something reasonable?
        integrated_smoothness = _test_integration(episodes, active_dims)
        if integrated_smoothness > 0.8:
            return "delta", 0.85
        return "velocity", 0.70

    # Check 3: High autocorrelation → absolute position
    mean_autocorr = np.mean([p.lag1_autocorrelation for p in active_profiles])
    if mean_autocorr > 0.90:
        return "absolute_position", 0.80

    # Check 4: Fall back to range-based heuristics
    mean_range = np.mean([p.range_width for p in active_profiles])
    mean_std = np.mean([p.std for p in active_profiles])

    if mean_range < 0.1:
        return "delta", 0.60
    if mean_std > 10:
        return "absolute_position", 0.55

    return "unknown", 0.30


def _test_integration(episodes: list[np.ndarray], dims: list[int]) -> float:
    """Test if integrating (cumsum) the signal produces smooth, bounded trajectories.

    High smoothness after integration → the original signal is delta/velocity.
    """
    smoothness_scores = []
    for ep in episodes[:5]:  # test on first 5 episodes
        if ep.ndim == 1:
            ep = ep.reshape(-1, 1)
        active = ep[:, dims] if dims else ep
        integrated = np.cumsum(active, axis=0)

        # Check smoothness: ratio of low-frequency power
        for d in range(integrated.shape[1]):
            ratio = _spectral_low_power_ratio(integrated[:, d])
            smoothness_scores.append(ratio)

        # Check boundedness: integrated shouldn't diverge wildly
        range_width = np.ptp(integrated, axis=0)
        is_bounded = np.all(range_width < 100)  # reasonable workspace bounds
        if not is_bounded:
            return 0.0

    return float(np.mean(smoothness_scores)) if smoothness_scores else 0.0


def _estimate_frequency(episodes: list[np.ndarray]) -> float:
    """Estimate the dominant frequency content of the data.

    Uses the autocorrelation time to estimate how fast the signal changes.
    Returns an estimate in Hz (assumes ~30 Hz sampling if unknown).
    """
    if not episodes or len(episodes[0]) < 20:
        return 0.0

    ep = episodes[0]
    if ep.ndim == 1:
        ep = ep.reshape(-1, 1)

    # Find the lag where autocorrelation drops below 0.5 (half-life)
    autocorrs = []
    for d in range(min(ep.shape[1], 5)):  # sample a few dims
        x = ep[:, d]
        x_c = x - np.mean(x)
        var = np.sum(x_c**2)
        if var < 1e-15:
            continue
        for lag in range(1, min(100, len(x) // 2)):
            ac = np.sum(x_c[:-lag] * x_c[lag:]) / var
            if ac < 0.5:
                autocorrs.append(lag)
                break

    if not autocorrs:
        return 0.0

    # Half-life in samples → frequency estimate
    median_halflife = float(np.median(autocorrs))
    # Dominant frequency ≈ 1 / (2 * half_life) of the sampling rate
    # Without knowing the actual sampling rate, return the relative frequency
    return 1.0 / max(2.0 * median_halflife, 1.0)


def _compute_smoothness(episodes: list[np.ndarray]) -> float:
    """Compute overall trajectory smoothness (0-1).

    Uses spectral low-frequency power ratio averaged across dimensions and episodes.
    """
    ratios = []
    for ep in episodes[:10]:  # sample up to 10 episodes
        if ep.ndim == 1:
            ep = ep.reshape(-1, 1)
        if len(ep) < 16:
            continue
        for d in range(ep.shape[1]):
            ratios.append(_spectral_low_power_ratio(ep[:, d]))

    return float(np.mean(ratios)) if ratios else 0.5


def _detect_kinematic_groups(
    all_data: np.ndarray,
    profiles: list[DimensionProfile],
) -> list[list[int]]:
    """Detect groups of correlated dimensions (kinematic chains).

    Bimanual robots have two groups of ~7 dims with high intra-group
    and low inter-group correlation.
    """
    ndim = all_data.shape[1]
    if ndim < 4:
        return [list(range(ndim))]

    # Compute correlation matrix
    active_dims = [i for i in range(ndim) if not profiles[i].is_dead]
    if len(active_dims) < 4:
        return [active_dims]

    active_data = all_data[:, active_dims]
    corr = np.abs(np.corrcoef(active_data.T))
    # Replace NaN with 0
    corr = np.nan_to_num(corr, nan=0.0)

    # Simple grouping: greedily build groups of highly correlated dims
    groups: list[list[int]] = []
    assigned = set()

    for i in range(len(active_dims)):
        if i in assigned:
            continue
        group = [active_dims[i]]
        assigned.add(i)
        for j in range(i + 1, len(active_dims)):
            if j in assigned:
                continue
            # Check if dim j is correlated with the group
            mean_corr = np.mean([corr[i, j] for i in range(len(active_dims)) if active_dims[i] in group])
            if mean_corr > 0.30:
                group.append(active_dims[j])
                assigned.add(j)
        groups.append(group)

    return groups


def _detect_bimanual(
    profiles: list[DimensionProfile],
    kinematic_groups: list[list[int]],
) -> bool:
    """Detect if this is a bimanual (dual-arm) robot.

    Heuristic: >=10 active dims, and two large groups of ~equal size.
    """
    ndim = len(profiles)
    active = [p for p in profiles if not p.is_dead and p.role.role != "gripper"]

    # Bimanual: 12-14 active dims typical
    if len(active) < 10:
        return False

    # Check if there are two roughly equal-sized kinematic groups
    large_groups = [g for g in kinematic_groups if len(g) >= 4]
    if len(large_groups) >= 2:
        sizes = sorted([len(g) for g in large_groups], reverse=True)
        if sizes[0] <= 2 * sizes[1]:  # roughly balanced
            return True

    # Dimension count heuristic: 12-14 dims with no dead → likely bimanual
    if 12 <= ndim <= 16 and len(active) >= 10:
        return True

    return False


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


def profile_array(
    episodes: list[np.ndarray],
) -> ArrayProfile:
    """Profile a robotics data array to infer what it represents.

    Given raw episode data (no metadata), infers dimension roles, action type,
    quality signals, and structural properties using statistical analysis only.

    Args:
        episodes: List of per-episode arrays, each of shape (timesteps, dims).
            Variable-length episodes are supported.

    Returns:
        ArrayProfile with full inference results.
    """
    if not episodes:
        return _empty_profile()

    # Normalize to 2D arrays
    normed = []
    for ep in episodes:
        arr = np.asarray(ep, dtype=np.float64)
        if arr.ndim == 1:
            arr = arr.reshape(-1, 1)
        normed.append(arr)
    episodes = normed

    ndim = episodes[0].shape[1]
    all_data = np.concatenate(episodes, axis=0)

    # Step 1: Profile each dimension
    dim_profiles = []
    for d in range(ndim):
        profile = _profile_dimension(all_data, d, episodes)
        dim_profiles.append(profile)

    # Step 2: Infer dimension roles
    for profile in dim_profiles:
        profile.role = _infer_dimension_role(profile, dim_profiles, episodes=episodes)

    # Step 3: Detect action type
    action_type, action_confidence = _detect_action_type(episodes, dim_profiles)

    # Step 4: Array-level properties
    is_normalized = bool(np.all(np.abs(all_data) <= 1.05))
    scale = float(np.median(np.abs(all_data)))
    frequency = _estimate_frequency(episodes)
    smoothness = _compute_smoothness(episodes)

    # Step 5: Quality signals
    dead_dims = [p.index for p in dim_profiles if p.is_dead]
    clip_dims = [p.index for p in dim_profiles if p.is_clipping]
    noisy_dims = [
        p.index for p in dim_profiles
        if p.spectral_low_power_ratio < 0.40 and not p.is_dead and not p.is_bimodal
    ]
    gripper_dims = [p.index for p in dim_profiles if p.role.role == "gripper"]

    # Step 6: Structural inference
    kinematic_groups = _detect_kinematic_groups(all_data, dim_profiles)
    is_bimanual = _detect_bimanual(dim_profiles, kinematic_groups)

    # Step 7: Overall confidence
    role_confidences = [p.role.confidence for p in dim_profiles if not p.is_dead]
    confidence = float(np.mean(role_confidences)) if role_confidences else 0.0
    confidence = confidence * 0.7 + action_confidence * 0.3  # weight action type detection

    return ArrayProfile(
        dimension_profiles=dim_profiles,
        dimension_roles=[p.role for p in dim_profiles],
        action_type=action_type,
        action_type_confidence=action_confidence,
        is_normalized=is_normalized,
        scale=scale,
        estimated_frequency_hz=frequency,
        smoothness=smoothness,
        dead_dimensions=dead_dims,
        clipping_dimensions=clip_dims,
        noisy_dimensions=noisy_dims,
        gripper_dimensions=gripper_dims,
        num_dimensions=ndim,
        is_bimanual=is_bimanual,
        kinematic_groups=kinematic_groups,
        confidence=confidence,
    )


def _empty_profile() -> ArrayProfile:
    """Return a default profile for empty data."""
    return ArrayProfile(
        dimension_profiles=[],
        dimension_roles=[],
        action_type="unknown",
        action_type_confidence=0.0,
        is_normalized=False,
        scale=0.0,
        estimated_frequency_hz=0.0,
        smoothness=0.0,
        dead_dimensions=[],
        clipping_dimensions=[],
        noisy_dimensions=[],
        gripper_dimensions=[],
        num_dimensions=0,
        is_bimanual=False,
        kinematic_groups=[],
        confidence=0.0,
    )


# ---------------------------------------------------------------------------
# Convenience: profile from a DataFrame (orbit standard format)
# ---------------------------------------------------------------------------


def profile_from_dataframe(
    df: "pd.DataFrame",
    column: str = "action",
) -> ArrayProfile:
    """Profile a column from an orbit-standard DataFrame.

    Args:
        df: DataFrame with episode_index and the target column.
        column: Column name containing array data (default: "action").

    Returns:
        ArrayProfile for the specified column.
    """
    import pandas as pd

    if column not in df.columns or df[column].isna().all():
        return _empty_profile()

    episodes = []
    for _, group in df.groupby("episode_index"):
        vals = group[column].tolist()
        # Convert list-of-lists or list-of-arrays to 2D array
        try:
            ep_array = np.array([np.asarray(v, dtype=np.float64) for v in vals])
            if ep_array.ndim == 1:
                ep_array = ep_array.reshape(-1, 1)
            episodes.append(ep_array)
        except (ValueError, TypeError):
            continue

    return profile_array(episodes)


# ---------------------------------------------------------------------------
# Compare two ArrayProfiles (e.g., action vs state)
# ---------------------------------------------------------------------------


def identify_action_vs_state(
    profiles: dict[str, ArrayProfile],
) -> dict[str, str]:
    """Given multiple array profiles, identify which is action and which is state.

    Heuristic: actions have lower autocorrelation (deltas) or are more
    smooth/directional. States have higher autocorrelation and broader range.

    Args:
        profiles: Dict mapping array names to their profiles.

    Returns:
        Dict mapping array names to roles: "action", "state", "image", "unknown".
    """
    assignments: dict[str, str] = {}

    if not profiles:
        return assignments

    # Score each profile on "action-likeness"
    scores: dict[str, float] = {}
    for name, prof in profiles.items():
        if prof.num_dimensions == 0:
            assignments[name] = "unknown"
            continue

        score = 0.0
        # Actions tend to be lower-dimensional
        if prof.num_dimensions <= 14:
            score += 0.2
        # Actions with gripper dims → strong signal
        if prof.gripper_dimensions:
            score += 0.3
        # Delta/velocity actions → strong signal
        if prof.action_type in ("delta", "velocity"):
            score += 0.3
        # Smooth trajectories → action
        if prof.smoothness > 0.6:
            score += 0.1
        # Higher dimension count → more likely state (state often includes more obs)
        if prof.num_dimensions > 20:
            score -= 0.3

        scores[name] = score

    if len(scores) >= 2:
        # Assign highest score as action, next as state
        sorted_names = sorted(scores.keys(), key=lambda n: scores[n], reverse=True)
        assignments[sorted_names[0]] = "action"
        assignments[sorted_names[1]] = "state"
        for name in sorted_names[2:]:
            assignments[name] = "unknown"
    elif len(scores) == 1:
        name = list(scores.keys())[0]
        assignments[name] = "action"

    return assignments
