"""Distribution-free outlier detection for robot demonstration data.

For each episode, compute a quality vector of per-episode metrics:
  - smoothness (inverse of normalized jerk)
  - path_length (normalized)
  - completion_ratio (1.0 if episode reaches final state similar to other episodes)
  - consistency_score (how similar this episode's trajectory is to the median trajectory)
  - action_magnitude_cv (within-episode action consistency)

Then flag episodes as outliers using Modified Z-Score (Iglewicz & Hoaglin):
  - MZS = 0.6745 * (x - median) / MAD
  - |MZS| > 3.5 → outlier

This works regardless of scale, robot type, or task because it compares
each episode to the OTHER episodes in the same dataset, not to fixed thresholds.
"""

from __future__ import annotations

import logging
from dataclasses import dataclass, field

import numpy as np

logger = logging.getLogger(__name__)

# Modified Z-Score threshold (Iglewicz & Hoaglin 1993)
MZS_THRESHOLD = 3.5

# Constant for converting MAD to standard deviation equivalent
MAD_SCALE = 0.6745


@dataclass
class EpisodeOutlierResult:
    episode_index: int
    is_outlier: bool
    outlier_reasons: list[str]  # e.g., ["high jerk (MZS=4.2)", "short duration (MZS=-3.8)"]
    quality_vector: dict[str, float]
    modified_z_scores: dict[str, float]


@dataclass
class DatasetOutlierReport:
    total_episodes: int
    outlier_count: int
    outlier_fraction: float
    episode_results: list[EpisodeOutlierResult]
    dataset_health: str  # "healthy" (<10% outliers), "noisy" (10-25%), "problematic" (>25%)
    metric_summaries: dict[str, dict] = field(default_factory=dict)


def _compute_jerk(actions: np.ndarray) -> float:
    """Compute mean squared jerk (third derivative) of action trajectory."""
    if len(actions) < 4:
        return 0.0
    vel = np.diff(actions, axis=0)
    acc = np.diff(vel, axis=0)
    jerk = np.diff(acc, axis=0)
    return float(np.mean(np.sum(jerk ** 2, axis=-1)))


def _compute_path_length(actions: np.ndarray) -> float:
    """Compute total path length in action space."""
    if len(actions) < 2:
        return 0.0
    diffs = np.diff(actions, axis=0)
    return float(np.sum(np.linalg.norm(diffs, axis=-1)))


def _compute_action_magnitude_cv(actions: np.ndarray) -> float:
    """Compute CV of action magnitudes within a single episode."""
    if len(actions) < 2:
        return 0.0
    magnitudes = np.linalg.norm(actions, axis=-1)
    mean_mag = np.mean(magnitudes)
    if mean_mag < 1e-10:
        return 0.0
    return float(np.std(magnitudes) / mean_mag)


def _compute_smoothness(actions: np.ndarray) -> float:
    """Inverse of normalized jerk — higher is smoother."""
    jerk = _compute_jerk(actions)
    # Normalize by path length squared to make scale-invariant
    pl = _compute_path_length(actions)
    if pl < 1e-10:
        return 0.0
    normalized_jerk = jerk / (pl ** 2 + 1e-10)
    return 1.0 / (1.0 + normalized_jerk)


def _compute_completion_ratio(actions: np.ndarray, all_final_states: np.ndarray) -> float:
    """How close the episode's final state is to the median final state.

    Returns 1.0 if the final action is close to the median, lower otherwise.
    """
    if len(actions) == 0 or len(all_final_states) == 0:
        return 0.0
    final = actions[-1]
    median_final = np.median(all_final_states, axis=0)
    dist = np.linalg.norm(final - median_final)
    # Normalize by median absolute deviation of final states
    mad = np.median(np.linalg.norm(all_final_states - median_final, axis=-1))
    if mad < 1e-10:
        return 1.0
    # Sigmoid-like mapping: distance of 0 → 1.0, distance >> MAD → 0.0
    return float(1.0 / (1.0 + dist / (mad + 1e-10)))


def _modified_z_scores(values: np.ndarray) -> np.ndarray:
    """Compute Modified Z-Scores using MAD (Median Absolute Deviation).

    MZS = 0.6745 * (x - median) / MAD

    The constant 0.6745 is the 0.75th quantile of the standard normal
    distribution, making MAD a consistent estimator of standard deviation.
    """
    median = np.median(values)
    mad = np.median(np.abs(values - median))
    # When MAD is tiny relative to the median, the metric is near-constant
    # across episodes and should not flag outliers.  Use a relative
    # threshold so that e.g. smoothness ≈ 0.9999 ± 0.00001 doesn't explode.
    min_mad = max(1e-10, abs(median) * 1e-4)
    if mad < min_mad:
        return np.zeros_like(values)
    return MAD_SCALE * (values - median) / mad


def _episode_metrics(
    actions: np.ndarray,
    all_final_states: np.ndarray,
) -> dict[str, float]:
    """Compute the quality vector for a single episode.

    path_length is normalized by episode length so that episodes with
    different frame counts are comparable (prevents shorter episodes
    from being flagged as outliers purely due to lower cumulative distance).
    """
    n = max(1, len(actions) - 1)  # number of steps
    return {
        "smoothness": _compute_smoothness(actions),
        "path_length": _compute_path_length(actions) / n,
        "duration": float(len(actions)),
        "completion_ratio": _compute_completion_ratio(actions, all_final_states),
        "action_magnitude_cv": _compute_action_magnitude_cv(actions),
    }


def detect_outliers(episode_data: dict[int, dict]) -> DatasetOutlierReport:
    """Detect outlier episodes using Modified Z-Score on per-episode quality metrics.

    Args:
        episode_data: Dict mapping episode_index -> {
            "actions": np.ndarray (T, D),
            "states": np.ndarray (T, D) or None,
            "length": int
        }

    Returns:
        DatasetOutlierReport with per-episode results and dataset health classification.
    """
    if not episode_data:
        return DatasetOutlierReport(
            total_episodes=0,
            outlier_count=0,
            outlier_fraction=0.0,
            episode_results=[],
            dataset_health="healthy",
        )

    indices = sorted(episode_data.keys())
    n_episodes = len(indices)

    # Need at least 5 episodes for meaningful outlier detection
    if n_episodes < 5:
        results = [
            EpisodeOutlierResult(
                episode_index=idx,
                is_outlier=False,
                outlier_reasons=[],
                quality_vector={},
                modified_z_scores={},
            )
            for idx in indices
        ]
        return DatasetOutlierReport(
            total_episodes=n_episodes,
            outlier_count=0,
            outlier_fraction=0.0,
            episode_results=results,
            dataset_health="healthy",
        )

    # Collect final states for completion ratio
    all_final_states = []
    for idx in indices:
        actions = np.asarray(episode_data[idx]["actions"], dtype=np.float32)
        if actions.ndim == 1:
            actions = actions.reshape(-1, 1)
        if len(actions) > 0:
            all_final_states.append(actions[-1])
    all_final_states = np.array(all_final_states) if all_final_states else np.array([])

    # Compute per-episode metrics
    metric_names = ["smoothness", "path_length", "duration", "completion_ratio", "action_magnitude_cv"]
    per_episode_metrics: dict[int, dict[str, float]] = {}
    for idx in indices:
        actions = np.asarray(episode_data[idx]["actions"], dtype=np.float32)
        if actions.ndim == 1:
            actions = actions.reshape(-1, 1)
        per_episode_metrics[idx] = _episode_metrics(actions, all_final_states)

    # Build metric arrays and compute MZS for each metric
    metric_arrays: dict[str, np.ndarray] = {}
    mzs_arrays: dict[str, np.ndarray] = {}
    for name in metric_names:
        values = np.array([per_episode_metrics[idx][name] for idx in indices])
        metric_arrays[name] = values
        mzs_arrays[name] = _modified_z_scores(values)

    # Compute metric summaries
    metric_summaries: dict[str, dict] = {}
    for name in metric_names:
        vals = metric_arrays[name]
        median = float(np.median(vals))
        mad = float(np.median(np.abs(vals - median)))
        metric_summaries[name] = {
            "median": median,
            "mad": mad,
            "min": float(np.min(vals)),
            "max": float(np.max(vals)),
        }

    # Flag outliers: any metric with |MZS| > threshold
    episode_results: list[EpisodeOutlierResult] = []
    outlier_count = 0

    for i, idx in enumerate(indices):
        reasons: list[str] = []
        mzs_dict: dict[str, float] = {}
        for name in metric_names:
            mzs = float(mzs_arrays[name][i])
            mzs_dict[name] = round(mzs, 2)
            if abs(mzs) > MZS_THRESHOLD:
                direction = "high" if mzs > 0 else "low"
                reasons.append(f"{direction} {name} (MZS={mzs:.1f})")

        is_outlier = len(reasons) > 0
        if is_outlier:
            outlier_count += 1

        episode_results.append(EpisodeOutlierResult(
            episode_index=idx,
            is_outlier=is_outlier,
            outlier_reasons=reasons,
            quality_vector=per_episode_metrics[idx],
            modified_z_scores=mzs_dict,
        ))

    outlier_fraction = outlier_count / n_episodes

    if outlier_fraction > 0.25:
        health = "problematic"
    elif outlier_fraction > 0.10:
        health = "noisy"
    else:
        health = "healthy"

    return DatasetOutlierReport(
        total_episodes=n_episodes,
        outlier_count=outlier_count,
        outlier_fraction=round(outlier_fraction, 4),
        episode_results=episode_results,
        dataset_health=health,
        metric_summaries=metric_summaries,
    )
