"""Cross-episode consistency metrics from Sakr et al. (ACM THRI 2025).

Computes 8 motion-based metrics per episode, then takes the coefficient of
variation (CV) across episodes for each metric. High CV = inconsistent
demonstrator = low predicted training success.

The paper showed that cross-episode CV of these metrics explains 70-91% of
training success variance (R²=0.702 in Study 1, R²=0.907 in Study 2).

Reference:
    Sakr et al., "Consistency Matters: Quantifying the Impact of
    Demonstrator Consistency on Robot Learning", ACM THRI 2025.
"""

from __future__ import annotations

import logging
from dataclasses import dataclass, field

import numpy as np

logger = logging.getLogger(__name__)


@dataclass
class ConsistencyReport:
    """Cross-episode consistency diagnostics.

    These are DIAGNOSTIC metrics — they tell you what's wrong with your
    data, not what your success rate will be. High CVs mean your
    demonstrations are inconsistent on that dimension. The per-episode
    scores identify which specific episodes are outliers.
    """

    # Per-metric cross-episode coefficient of variation
    path_length_cv: float  # CV of Σ‖xₜ-xₜ₋₁‖² across episodes
    joint_path_length_cv: float  # CV of Σ‖qₜ-qₜ₋₁‖² (same as path_length for action-only)
    cartesian_jerk_cv: float  # CV of Σ(jerk²) across episodes
    joint_jerk_cv: float  # CV of Σ(jerk²) in joint space
    curvature_cv: float  # CV of curvature across episodes
    effort_cv: float  # CV of Σ(action²) across episodes
    joint_limit_cv: float  # CV of joint limit proximity across episodes
    orientation_cv: float  # CV of rotation path length (when ≥6D)

    # Composite consistency score (inverse of weighted mean CV, 0-1)
    overall_consistency: float

    # Per-episode consistency scores (for curation)
    per_episode_scores: dict[int, float] = field(default_factory=dict)

    # Raw per-episode metric values (for downstream analysis)
    per_episode_metrics: dict[str, list[float]] = field(default_factory=dict)


# ---------------------------------------------------------------------------
# Per-episode metric computations
# ---------------------------------------------------------------------------


def _compute_path_length(episode: np.ndarray) -> float:
    """Mean squared step distance (normalized by episode length).

    Normalizing by T makes this metric comparable across episodes with
    different frame counts or frame rates.
    """
    if len(episode) < 2:
        return 0.0
    diffs = np.diff(episode, axis=0)  # (T-1, D)
    return float(np.sum(np.sum(diffs**2, axis=1)) / len(diffs))


def _compute_jerk(episode: np.ndarray) -> float:
    """Mean squared jerk (3rd derivative, normalized by episode length)."""
    if len(episode) < 4:
        return 0.0
    vel = np.diff(episode, axis=0)  # (T-1, D)
    acc = np.diff(vel, axis=0)  # (T-2, D)
    jerk = np.diff(acc, axis=0)  # (T-3, D)
    return float(np.sum(jerk**2) / len(jerk))


def _compute_curvature(episode: np.ndarray) -> float:
    """Mean curvature along trajectory (normalized by episode length).

    Uses κ = |vel × acc| / |vel|³, generalized to D dimensions via
    the cross-product magnitude formula: |a×b| = sqrt(|a|²|b|² - (a·b)²).
    """
    if len(episode) < 3 or episode.shape[1] < 2:
        return 0.0
    vel = np.diff(episode, axis=0)  # (T-1, D)
    acc = np.diff(vel, axis=0)  # (T-2, D)
    vel_trunc = vel[:-1]  # align with acc: (T-2, D)

    speed_sq = np.sum(vel_trunc**2, axis=1)  # (T-2,)
    speed = np.sqrt(speed_sq) + 1e-8

    # Cross product magnitude: |v||a|sin(θ) = sqrt(|v|²|a|² - (v·a)²)
    acc_sq = np.sum(acc**2, axis=1)
    dot_va = np.sum(vel_trunc * acc, axis=1)
    cross_mag_sq = speed_sq * acc_sq - dot_va**2
    # Numerical safety: clamp to zero before sqrt
    cross_mag = np.sqrt(np.maximum(cross_mag_sq, 0.0))

    curvature = cross_mag / (speed**3)
    return float(np.sum(curvature) / len(curvature)) if len(curvature) > 0 else 0.0


def _compute_effort(episode: np.ndarray) -> float:
    """Mean squared action magnitude (proxy for joint effort, normalized by T)."""
    if len(episode) == 0:
        return 0.0
    return float(np.sum(episode**2) / len(episode))


def _compute_joint_limit_proximity(episode: np.ndarray) -> float:
    """How well actions stay away from their observed min/max range.

    Uses the formula from Sakr et al.: Π 4(q-q⁻)(q⁺-q)/range²
    Maximized when joints are at center of range, zero at limits.
    Normalized by episode length for cross-episode comparability.
    """
    if len(episode) < 2 or episode.shape[1] < 1:
        return 0.0
    mins = np.min(episode, axis=0)
    maxs = np.max(episode, axis=0)
    ranges = maxs - mins
    # Skip dimensions with zero range (dead joints)
    active = ranges > 1e-8
    if not np.any(active):
        return 0.0

    # Normalized position within range, per timestep
    ep_active = episode[:, active]
    mins_active = mins[active]
    ranges_active = ranges[active]
    normalized = (ep_active - mins_active) / ranges_active  # 0-1

    # Joint limit metric: 4*x*(1-x) — maximized at center (=1.0)
    proximity = 4 * normalized * (1 - normalized)  # (T, D_active)
    # Product across joints, then mean across time
    per_timestep = np.prod(proximity + 1e-8, axis=1)
    return float(np.mean(per_timestep))


def _compute_orientation_path_length(episode: np.ndarray) -> float:
    """Mean squared rotation step for orientation dimensions.

    If the episode has >= 6 dimensions, assumes the last 3 are
    orientation (Euler angles or axis-angle). Normalized by episode
    length for cross-episode comparability.

    For episodes with < 6 dims, returns 0.0 (no orientation data).
    """
    if episode.shape[1] < 6:
        return 0.0
    # Assume last 3 dims are orientation
    orient = episode[:, -3:]
    if len(orient) < 2:
        return 0.0
    diffs = np.diff(orient, axis=0)
    # Wrap angular differences to [-π, π]
    diffs = np.arctan2(np.sin(diffs), np.cos(diffs))
    return float(np.sum(diffs**2) / len(diffs))


# ---------------------------------------------------------------------------
# Coefficient of variation
# ---------------------------------------------------------------------------


def _cv(values: list[float]) -> float:
    """Coefficient of variation. Returns 0 if mean is near zero."""
    if len(values) < 2:
        return 0.0
    arr = np.array(values)
    mean = np.mean(arr)
    if abs(mean) < 1e-10:
        return 0.0
    return float(np.std(arr) / abs(mean))


# ---------------------------------------------------------------------------
# Sakr-based success prediction
# ---------------------------------------------------------------------------

# Weights derived from Sakr et al. feature importance rankings.
# Path length CV and jerk CV were the most predictive (highest R²
# contribution). These weights reflect their published ordering.
_SAKR_CV_WEIGHTS = {
    "path_length_cv": 0.20,
    "joint_path_length_cv": 0.15,
    "cartesian_jerk_cv": 0.18,
    "joint_jerk_cv": 0.15,
    "curvature_cv": 0.10,
    "effort_cv": 0.08,
    "joint_limit_cv": 0.07,
    "orientation_cv": 0.07,
}


# ---------------------------------------------------------------------------
# Per-episode consistency scoring (for curation)
# ---------------------------------------------------------------------------


def _score_episodes(
    metric_values: dict[str, list[float]],
    n_episodes: int,
) -> dict[int, float]:
    """Score each episode by how close it is to the dataset's central tendency.

    Episodes with metric values far from the median on multiple metrics
    get lower scores (they're inconsistent outliers).

    Returns:
        dict mapping episode_index → consistency_score (0-1, higher = better)
    """
    if n_episodes == 0:
        return {}

    # For each metric, compute how far each episode is from the median
    # using a robust z-score (MAD-based)
    per_episode_z = np.zeros(n_episodes)
    n_metrics = 0

    for metric_name, values in metric_values.items():
        if len(values) != n_episodes:
            continue
        arr = np.array(values)
        median = np.median(arr)
        mad = np.median(np.abs(arr - median))
        if mad < 1e-10:
            continue
        z_scores = np.abs(arr - median) / (mad * 1.4826)  # MAD → σ
        per_episode_z += z_scores
        n_metrics += 1

    if n_metrics == 0:
        return {i: 0.5 for i in range(n_episodes)}

    per_episode_z /= n_metrics

    # Convert mean z-score to a 0-1 score: z=0 → 1.0, z=3 → ~0.05
    scores = np.exp(-0.5 * per_episode_z**2)
    return {i: float(scores[i]) for i in range(n_episodes)}


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


def compute_consistency_metrics(
    action_episodes: list[np.ndarray],
    state_episodes: list[np.ndarray] | None = None,
) -> ConsistencyReport:
    """Compute Sakr et al. consistency metrics across episodes.

    Computes 8 motion-based metrics per episode, then takes the cross-episode
    coefficient of variation (CV) for each. Uses the CVs to predict training
    success based on the paper's findings.

    Args:
        action_episodes: List of per-episode action arrays (T_i, action_dim).
        state_episodes: Optional list of per-episode state arrays. If provided,
            Cartesian metrics use state data (preferred); joint metrics still
            use action data.

    Returns:
        ConsistencyReport with CVs, predicted success, and per-episode scores.
    """
    n_episodes = len(action_episodes)

    if n_episodes < 2:
        return ConsistencyReport(
            path_length_cv=0.0,
            joint_path_length_cv=0.0,
            cartesian_jerk_cv=0.0,
            joint_jerk_cv=0.0,
            curvature_cv=0.0,
            effort_cv=0.0,
            joint_limit_cv=0.0,
            orientation_cv=0.0,
            overall_consistency=0.5,
            per_episode_scores={0: 0.5} if n_episodes == 1 else {},
        )

    # Ensure all episodes are 2D
    episodes_2d = []
    for ep in action_episodes:
        arr = np.asarray(ep, dtype=np.float64)
        if arr.ndim == 1:
            arr = arr.reshape(-1, 1)
        episodes_2d.append(arr)

    # Use state episodes for Cartesian metrics if available
    cartesian_episodes = episodes_2d
    if state_episodes is not None and len(state_episodes) == n_episodes:
        state_2d = []
        for ep in state_episodes:
            arr = np.asarray(ep, dtype=np.float64)
            if arr.ndim == 1:
                arr = arr.reshape(-1, 1)
            state_2d.append(arr)
        cartesian_episodes = state_2d

    # Compute per-episode metric values
    path_lengths = [_compute_path_length(ep) for ep in cartesian_episodes]
    joint_path_lengths = [_compute_path_length(ep) for ep in episodes_2d]
    cartesian_jerks = [_compute_jerk(ep) for ep in cartesian_episodes]
    joint_jerks = [_compute_jerk(ep) for ep in episodes_2d]
    curvatures = [_compute_curvature(ep) for ep in cartesian_episodes]
    efforts = [_compute_effort(ep) for ep in episodes_2d]
    joint_limits = [_compute_joint_limit_proximity(ep) for ep in episodes_2d]
    orientations = [_compute_orientation_path_length(ep) for ep in episodes_2d]

    # Store raw per-episode values
    per_episode_metrics = {
        "path_length": path_lengths,
        "joint_path_length": joint_path_lengths,
        "cartesian_jerk": cartesian_jerks,
        "joint_jerk": joint_jerks,
        "curvature": curvatures,
        "effort": efforts,
        "joint_limit": joint_limits,
        "orientation": orientations,
    }

    # Compute cross-episode CV for each metric
    path_length_cv = _cv(path_lengths)
    joint_path_length_cv = _cv(joint_path_lengths)
    cartesian_jerk_cv = _cv(cartesian_jerks)
    joint_jerk_cv = _cv(joint_jerks)
    curvature_cv = _cv(curvatures)
    effort_cv = _cv(efforts)
    joint_limit_cv = _cv(joint_limits)
    orientation_cv = _cv(orientations)

    cvs = {
        "path_length_cv": path_length_cv,
        "joint_path_length_cv": joint_path_length_cv,
        "cartesian_jerk_cv": cartesian_jerk_cv,
        "joint_jerk_cv": joint_jerk_cv,
        "curvature_cv": curvature_cv,
        "effort_cv": effort_cv,
        "joint_limit_cv": joint_limit_cv,
        "orientation_cv": orientation_cv,
    }

    # Overall consistency: inverse of weighted mean CV, mapped to 0-1
    weighted_cv = sum(
        _SAKR_CV_WEIGHTS.get(k, 0.0) * v for k, v in cvs.items()
    )
    overall_consistency = float(np.clip(1.0 / (1.0 + weighted_cv), 0.0, 1.0))

    # Per-episode consistency scores
    per_episode_scores = _score_episodes(per_episode_metrics, n_episodes)

    return ConsistencyReport(
        path_length_cv=round(path_length_cv, 4),
        joint_path_length_cv=round(joint_path_length_cv, 4),
        cartesian_jerk_cv=round(cartesian_jerk_cv, 4),
        joint_jerk_cv=round(joint_jerk_cv, 4),
        curvature_cv=round(curvature_cv, 4),
        effort_cv=round(effort_cv, 4),
        joint_limit_cv=round(joint_limit_cv, 4),
        orientation_cv=round(orientation_cv, 4),
        overall_consistency=round(overall_consistency, 4),
        per_episode_scores=per_episode_scores,
        per_episode_metrics=per_episode_metrics,
    )
