"""Unsupervised data curation for robot training datasets.

Two methods for scoring and selecting the best episodes:

Method A — Consistency-based (fast, for datasets <200 episodes):
    Score each episode by deviation from the dataset's central tendency
    on the Sakr consistency metrics. Outlier episodes get lower scores.

Method B — MI-based (for datasets ≥200 episodes):
    Mutual information scoring inspired by DemInf (RSS 2025). Estimates
    each episode's information contribution to the dataset using k-NN
    MI between states and actions.

Reference:
    Hejna et al., "Demonstration Information is Enough for Data Curation
    in Offline Imitation Learning", RSS 2025.
"""

from __future__ import annotations

import logging
from dataclasses import dataclass, field

import numpy as np

logger = logging.getLogger(__name__)


@dataclass
class CurationResult:
    """Per-episode quality scores for dataset curation."""

    episode_scores: dict[int, float]  # episode_index → quality score (0-1)
    ranked_episodes: list[int]  # episode indices sorted best → worst
    recommended_keep: list[int]  # episodes to keep
    recommended_remove: list[int]  # episodes to remove
    budget_curve: list[tuple[int, float]]  # [(n_episodes, estimated_quality), ...]
    method: str  # "consistency" or "mi"
    summary: str


# ---------------------------------------------------------------------------
# Method A: Consistency-based scoring
# ---------------------------------------------------------------------------


def _consistency_scores(
    action_episodes: list[np.ndarray],
    state_episodes: list[np.ndarray] | None = None,
) -> dict[int, float]:
    """Score episodes by deviation from central tendency on multiple metrics.

    Uses the same metric functions as consistency_metrics.py but scores
    each episode individually (not cross-episode CV).
    """
    from orbit.analyzer.consistency_metrics import (
        _compute_curvature,
        _compute_effort,
        _compute_jerk,
        _compute_path_length,
    )

    n = len(action_episodes)
    if n < 3:
        return {i: 0.5 for i in range(n)}

    episodes = []
    for ep in action_episodes:
        arr = np.asarray(ep, dtype=np.float64)
        if arr.ndim == 1:
            arr = arr.reshape(-1, 1)
        episodes.append(arr)

    # Compute per-episode metrics
    metrics = {
        "path_length": [_compute_path_length(ep) for ep in episodes],
        "jerk": [_compute_jerk(ep) for ep in episodes],
        "curvature": [_compute_curvature(ep) for ep in episodes],
        "effort": [_compute_effort(ep) for ep in episodes],
        "length": [float(len(ep)) for ep in episodes],
    }

    # For each metric, compute robust z-scores (MAD-based)
    total_z = np.zeros(n)
    n_active = 0

    for name, values in metrics.items():
        arr = np.array(values)
        median = np.median(arr)
        mad = np.median(np.abs(arr - median))
        if mad < 1e-10:
            continue
        z = np.abs(arr - median) / (mad * 1.4826)
        total_z += z
        n_active += 1

    if n_active == 0:
        return {i: 0.5 for i in range(n)}

    total_z /= n_active

    # Convert to 0-1 scores: z=0 → 1.0, z=3 → ~0.01
    scores = np.exp(-0.5 * total_z**2)
    return {i: float(scores[i]) for i in range(n)}


# ---------------------------------------------------------------------------
# Method B: MI-based scoring (DemInf-inspired)
# ---------------------------------------------------------------------------


def _mi_scores(
    action_episodes: list[np.ndarray],
    state_episodes: list[np.ndarray] | None = None,
) -> dict[int, float]:
    """Score episodes by mutual information contribution.

    Approximates DemInf: for each episode, estimate how much it
    contributes to the overall state-action MI of the dataset.
    Episodes with high MI contribution contain informative transitions;
    low MI episodes are redundant or noisy.
    """
    n = len(action_episodes)
    if n < 5:
        return _consistency_scores(action_episodes, state_episodes)

    # Prepare data
    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)

    states_2d = None
    if state_episodes and len(state_episodes) == n:
        states_2d = []
        for ep in state_episodes:
            arr = np.asarray(ep, dtype=np.float64)
            if arr.ndim == 1:
                arr = arr.reshape(-1, 1)
            states_2d.append(arr)

    try:
        from sklearn.feature_selection import mutual_info_regression

        # Compute per-episode MI between states and actions
        # If no states, use lagged actions as proxy states
        scores = {}
        for i in range(n):
            actions = episodes_2d[i]
            if states_2d is not None:
                states = states_2d[i]
                # Align lengths
                min_len = min(len(actions), len(states))
                actions = actions[:min_len]
                states = states[:min_len]
            else:
                # Use lagged actions as proxy for states
                if len(actions) < 3:
                    scores[i] = 0.5
                    continue
                states = actions[:-1]
                actions = actions[1:]

            if len(actions) < 5:
                scores[i] = 0.5
                continue

            # Subsample for speed
            max_samples = 500
            if len(actions) > max_samples:
                rng = np.random.RandomState(42 + i)
                idx = rng.choice(len(actions), max_samples, replace=False)
                actions = actions[idx]
                states = states[idx]

            # MI per action dim
            mi_vals = []
            for d in range(actions.shape[1]):
                mi = mutual_info_regression(
                    states, actions[:, d],
                    n_neighbors=3,
                    random_state=42,
                )
                mi_vals.append(float(np.mean(mi)))

            scores[i] = float(np.mean(mi_vals))

        # Normalize to 0-1
        if scores:
            vals = list(scores.values())
            min_mi = min(vals)
            max_mi = max(vals)
            rng_mi = max_mi - min_mi
            if rng_mi > 1e-10:
                scores = {k: (v - min_mi) / rng_mi for k, v in scores.items()}
            else:
                scores = {k: 0.5 for k in scores}

        return scores

    except ImportError:
        logger.info("sklearn not available, falling back to consistency scoring")
        return _consistency_scores(action_episodes, state_episodes)


# ---------------------------------------------------------------------------
# Budget curve: estimated quality at each subset size
# ---------------------------------------------------------------------------


def _compute_budget_curve(
    ranked_indices: list[int],
    episode_scores: dict[int, float],
    n_episodes: int,
) -> list[tuple[int, float]]:
    """Compute (n_kept, avg_quality) for various budget levels."""
    if not ranked_indices:
        return []

    curve = []
    running_sum = 0.0
    for k, idx in enumerate(ranked_indices, 1):
        running_sum += episode_scores.get(idx, 0.0)
        avg_quality = running_sum / k
        # Sample at powers of 2 and key fractions
        if k in (1, 5, 10, 25, 50, 100) or k == n_episodes or k % max(1, n_episodes // 10) == 0:
            curve.append((k, round(avg_quality, 4)))

    return curve


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


def curate_episodes(
    action_episodes: list[np.ndarray],
    state_episodes: list[np.ndarray] | None = None,
    budget: int | None = None,
    method: str = "auto",
) -> CurationResult:
    """Select the best episodes from a dataset.

    Args:
        action_episodes: List of per-episode action arrays.
        state_episodes: Optional state arrays.
        budget: If set, keep only the top N episodes.
        method: "auto", "consistency", or "mi".

    Returns:
        CurationResult with ranked episodes and curation recommendations.
    """
    n = len(action_episodes)

    if n == 0:
        return CurationResult(
            episode_scores={},
            ranked_episodes=[],
            recommended_keep=[],
            recommended_remove=[],
            budget_curve=[],
            method="none",
            summary="No episodes to curate.",
        )

    # Select method
    if method == "auto":
        method = "mi" if n >= 200 else "consistency"

    if method == "mi":
        scores = _mi_scores(action_episodes, state_episodes)
        method_label = "mi"
    else:
        scores = _consistency_scores(action_episodes, state_episodes)
        method_label = "consistency"

    # Rank episodes by score (best first)
    ranked = sorted(scores.keys(), key=lambda i: scores.get(i, 0.0), reverse=True)

    # Determine keep/remove
    if budget is not None and budget < n:
        keep = ranked[:budget]
        remove = ranked[budget:]
    else:
        # Default: remove bottom 5% if they're significantly worse
        if n >= 20:
            cutoff_idx = max(1, int(n * 0.95))
            threshold = scores.get(ranked[cutoff_idx - 1], 0.0) if cutoff_idx <= len(ranked) else 0.0
            keep = [i for i in ranked if scores.get(i, 0.0) >= threshold * 0.5]
            remove = [i for i in ranked if i not in set(keep)]
        else:
            keep = list(ranked)
            remove = []

    budget_curve = _compute_budget_curve(ranked, scores, n)

    # Summary
    if budget is not None and budget < n:
        avg_kept = np.mean([scores[i] for i in keep]) if keep else 0
        avg_all = np.mean(list(scores.values()))
        summary = (
            f"Selected {len(keep)} of {n} episodes. "
            f"Avg quality: {avg_kept:.2f} (vs {avg_all:.2f} for all). "
            f"Method: {method_label}."
        )
    elif remove:
        summary = (
            f"Identified {len(remove)} low-quality episodes for removal. "
            f"Method: {method_label}."
        )
    else:
        summary = f"All {n} episodes scored. No removals recommended. Method: {method_label}."

    return CurationResult(
        episode_scores=scores,
        ranked_episodes=ranked,
        recommended_keep=sorted(keep),
        recommended_remove=sorted(remove),
        budget_curve=budget_curve,
        method=method_label,
        summary=summary,
    )
