"""Per-episode quality ranking for robot training datasets.

Scores each episode on smoothness, completion, consistency, length, and
range utilization. Produces a ranked list with keep/review/remove verdicts
and estimates the quality improvement from removing low-quality episodes.
"""

from __future__ import annotations

import logging
from dataclasses import dataclass, field

import numpy as np

logger = logging.getLogger(__name__)


@dataclass
class EpisodeRank:
    """Quality ranking for a single episode."""

    episode_index: int
    smoothness: float  # 0-1
    completion: float  # 0-1
    consistency: float  # 0-1
    length_score: float  # 0-1
    range_utilization: float  # 0-1
    composite_score: float  # 0-1, weighted combination
    verdict: str  # "keep", "review", "remove"


@dataclass
class EpisodeRankingResult:
    """Result of per-episode quality ranking."""

    episodes: list[EpisodeRank]
    keep_count: int
    review_count: int
    remove_count: int
    remove_indices: list[int]
    estimated_improvement: float  # quality score delta if bad episodes removed
    current_quality_estimate: float  # approximate current quality
    improved_quality_estimate: float  # approximate quality after removal
    summary: str


# Weights for composite score
WEIGHTS = {
    "smoothness": 0.25,
    "completion": 0.20,
    "consistency": 0.30,
    "length_score": 0.15,
    "range_utilization": 0.10,
}


def compute_episode_ranking(
    action_episodes: list[np.ndarray],
    state_episodes: list[np.ndarray] | None = None,
    curation_scores: dict[int, float] | None = None,
) -> EpisodeRankingResult:
    """Rank episodes by quality and identify candidates for removal.

    Args:
        action_episodes: List of per-episode action arrays (T_i, action_dim).
        state_episodes: Optional list of per-episode state arrays for
            completion scoring.
        curation_scores: Optional dict mapping episode_index → quality score
            (0-1) from data_curator.py. When provided, these are used as the
            primary ranking signal (70% weight) with heuristic scores as
            secondary (30% weight).

    Returns:
        EpisodeRankingResult with ranked episodes and improvement estimate.
    """
    if not action_episodes:
        return EpisodeRankingResult(
            episodes=[], keep_count=0, review_count=0, remove_count=0,
            remove_indices=[], estimated_improvement=0.0,
            current_quality_estimate=0.0, improved_quality_estimate=0.0,
            summary="No episodes to rank.",
        )

    n_episodes = len(action_episodes)

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

    # Pre-compute cross-episode statistics
    lengths = [len(ep) for ep in episodes_2d]
    median_length = float(np.median(lengths)) if lengths else 1.0

    # Mean action sequence (for consistency scoring)
    # Resample all episodes to median length for comparison
    resampled = _resample_episodes(episodes_2d, int(median_length))
    if resampled:
        mean_trajectory = np.mean(np.array(resampled), axis=0)
    else:
        mean_trajectory = None

    # Global action range per dimension (for range utilization)
    all_actions = np.concatenate(episodes_2d, axis=0)
    global_range = np.ptp(all_actions, axis=0)
    global_range[global_range < 1e-8] = 1.0

    # Score each episode
    rankings: list[EpisodeRank] = []
    for i, ep in enumerate(episodes_2d):
        smoothness = _score_smoothness(ep)
        completion = _score_completion(
            state_episodes[i] if state_episodes and i < len(state_episodes) else ep
        )
        consistency = _score_consistency(ep, mean_trajectory, int(median_length))
        length_score = _score_length(len(ep), median_length)
        range_util = _score_range_utilization(ep, global_range)

        heuristic_score = (
            WEIGHTS["smoothness"] * smoothness
            + WEIGHTS["completion"] * completion
            + WEIGHTS["consistency"] * consistency
            + WEIGHTS["length_score"] * length_score
            + WEIGHTS["range_utilization"] * range_util
        )

        if curation_scores is not None and i in curation_scores:
            # Blend curation (primary) with heuristic (secondary)
            composite = curation_scores[i] * 0.70 + heuristic_score * 0.30
        else:
            composite = heuristic_score

        if composite >= 0.7:
            verdict = "keep"
        elif composite >= 0.4:
            verdict = "review"
        else:
            verdict = "remove"

        rankings.append(EpisodeRank(
            episode_index=i,
            smoothness=round(smoothness, 4),
            completion=round(completion, 4),
            consistency=round(consistency, 4),
            length_score=round(length_score, 4),
            range_utilization=round(range_util, 4),
            composite_score=round(composite, 4),
            verdict=verdict,
        ))

    # Sort by composite score (best first for display)
    rankings.sort(key=lambda r: r.composite_score, reverse=True)

    keep = sum(1 for r in rankings if r.verdict == "keep")
    review = sum(1 for r in rankings if r.verdict == "review")
    remove = sum(1 for r in rankings if r.verdict == "remove")
    remove_indices = sorted([r.episode_index for r in rankings if r.verdict == "remove"])

    # Estimate quality improvement from removal
    all_scores = [r.composite_score for r in rankings]
    current_avg = float(np.mean(all_scores)) if all_scores else 0.0

    kept_scores = [r.composite_score for r in rankings if r.verdict != "remove"]
    improved_avg = float(np.mean(kept_scores)) if kept_scores else current_avg

    improvement = round(improved_avg - current_avg, 4)

    # Summary
    if remove == 0:
        summary = f"All {n_episodes} episodes are good quality."
    else:
        worst = sorted(rankings, key=lambda r: r.composite_score)[:3]
        worst_str = ", ".join(f"#{r.episode_index} ({r.composite_score:.2f})" for r in worst)
        summary = (
            f"Removing {remove} low-quality episode(s) would improve "
            f"average quality from {current_avg:.2f} to {improved_avg:.2f}. "
            f"Worst: {worst_str}"
        )

    return EpisodeRankingResult(
        episodes=rankings,
        keep_count=keep,
        review_count=review,
        remove_count=remove,
        remove_indices=remove_indices,
        estimated_improvement=improvement,
        current_quality_estimate=round(current_avg, 4),
        improved_quality_estimate=round(improved_avg, 4),
        summary=summary,
    )


def _score_smoothness(actions: np.ndarray) -> float:
    """Inverse of mean absolute jerk, normalized to 0-1."""
    if len(actions) < 4:
        return 0.5

    # Dead episode check: if the action range is near-zero, the episode
    # has no meaningful motion — this is NOT smooth, it's broken.
    action_range = np.ptp(actions, axis=0)
    if np.max(action_range) < 0.01:
        return 0.0

    velocity = np.diff(actions, axis=0)
    acceleration = np.diff(velocity, axis=0)
    jerk = np.diff(acceleration, axis=0)

    action_scale = np.std(actions) + 1e-8
    normalized_jerk = np.mean(np.abs(jerk)) / action_scale

    return float(np.clip(1.0 - normalized_jerk / 2.0, 0.0, 1.0))


def _score_completion(data: np.ndarray) -> float:
    """Does the trajectory converge? Compare variance of last 20% vs first 20%."""
    if data.ndim == 1:
        data = data.reshape(-1, 1)

    T = len(data)
    if T < 10:
        return 0.5

    split = max(2, T // 5)
    early_var = np.mean(np.var(data[:split], axis=0)) + 1e-8
    late_var = np.mean(np.var(data[-split:], axis=0)) + 1e-8

    ratio = late_var / early_var
    score = 1.0 / (1.0 + ratio)
    return float(np.clip(score, 0.0, 1.0))


def _score_consistency(
    actions: np.ndarray,
    mean_trajectory: np.ndarray | None,
    target_length: int,
) -> float:
    """Cosine similarity of this episode's resampled actions to the mean."""
    if mean_trajectory is None:
        return 0.5

    resampled = _resample_single(actions, target_length)
    if resampled is None:
        return 0.5

    a = resampled.flatten()
    b = mean_trajectory.flatten()

    norm_a = np.linalg.norm(a)
    norm_b = np.linalg.norm(b)
    if norm_a < 1e-8:
        # Zero episode — not consistent with anything, it's dead data
        return 0.0
    if norm_b < 1e-8:
        return 0.5

    cosine = float(np.dot(a, b) / (norm_a * norm_b))
    # Map from [-1, 1] to [0, 1]
    return float(np.clip((cosine + 1.0) / 2.0, 0.0, 1.0))


def _score_length(length: int, median_length: float) -> float:
    """Gaussian penalty centered on median length."""
    if median_length < 1:
        return 0.5

    # Gaussian with sigma = 0.4 * median
    sigma = max(1.0, 0.4 * median_length)
    deviation = abs(length - median_length)
    score = np.exp(-0.5 * (deviation / sigma) ** 2)
    return float(score)


def _score_range_utilization(actions: np.ndarray, global_range: np.ndarray) -> float:
    """Fraction of action space range this episode covers."""
    if len(actions) < 2:
        return 0.0

    ep_range = np.ptp(actions, axis=0)
    # Per-dimension utilization
    utilization = ep_range / global_range
    return float(np.clip(np.mean(utilization), 0.0, 1.0))


def _resample_single(actions: np.ndarray, target_length: int) -> np.ndarray | None:
    """Resample a single episode to target_length via linear interpolation."""
    if len(actions) < 2 or target_length < 2:
        return None

    if actions.ndim == 1:
        actions = actions.reshape(-1, 1)

    T, D = actions.shape
    old_indices = np.linspace(0, T - 1, T)
    new_indices = np.linspace(0, T - 1, target_length)

    resampled = np.zeros((target_length, D))
    for d in range(D):
        resampled[:, d] = np.interp(new_indices, old_indices, actions[:, d])

    return resampled


def _resample_episodes(
    episodes: list[np.ndarray],
    target_length: int,
) -> list[np.ndarray]:
    """Resample all episodes to the same length."""
    result = []
    for ep in episodes:
        r = _resample_single(ep, target_length)
        if r is not None:
            result.append(r)
    return result
