"""Quality scoring engine for robot training datasets.

Combines coverage-based scores with action quality metrics (smoothness, jerk,
mutual information), episode quality (completion, consistency), and signal
diagnostics from the signal_diagnostics module.
"""

from __future__ import annotations

import logging
from dataclasses import dataclass, field

import numpy as np

from orbit.analyzer.coverage import FullCoverageReport
from orbit.analyzer.dataset_loader import DatasetInfo

logger = logging.getLogger(__name__)

# Default expected episodes for a task
DEFAULT_EXPECTED_EPISODES = 50

# Weights for quality score components (without visual diversity)
BASE_WEIGHTS = {
    "episode_count": 0.20,
    "position_diversity": 0.10,
    "action_diversity": 0.15,
    "episode_consistency": 0.35,
    "temporal_coverage": 0.20,
}

# Weight for visual diversity when embedding analysis is available
VISUAL_DIVERSITY_WEIGHT = 0.10

# Confidence thresholds
CONFIDENCE_HIGH_THRESHOLD = 50
CONFIDENCE_MEDIUM_THRESHOLD = 20

# Confidence widths for estimated success rate
CONFIDENCE_WIDTHS = {
    "high": 0.05,
    "medium": 0.10,
    "low": 0.15,
}

# Weights for the enhanced quality score when action/state data is available
ENHANCED_WEIGHTS = {
    "coverage": 0.35,
    "action_quality": 0.25,
    "episode_quality": 0.20,
    "signal_summary": 0.20,
}


@dataclass
class QualityReport:
    """Overall quality assessment of a dataset."""

    overall_score: float  # 0.0 to 1.0
    component_scores: dict[str, float]  # each component and its score
    ready_to_train: bool  # True if overall >= 0.65
    confidence: str  # "low", "medium", "high"
    estimated_success_rate: tuple[float, float]  # range, e.g. (0.45, 0.55)


@dataclass
class ReadyAssessment:
    """Ready-to-train assessment with actionable details."""

    verdict: str  # "YES" | "NOT YET" | "NO"
    reason: str  # one-line explanation
    blocking_issues: list[str] = field(default_factory=list)
    nice_to_haves: list[str] = field(default_factory=list)


# ---------------------------------------------------------------------------
# Ported quality metrics from Desktop/Orbit research version
# ---------------------------------------------------------------------------


def compute_action_smoothness(actions: np.ndarray) -> float:
    """Score action trajectory smoothness via jerk (3rd derivative).

    Smooth expert demonstrations have low jerk.

    Args:
        actions: Array of shape (T, action_dim) for a single episode.

    Returns:
        Score in [0, 1]. 1.0 = perfectly smooth, 0.0 = very jerky.
    """
    if len(actions) < 4:
        return 0.5

    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

    # Scale the mapping by the action range so the 0→1 mapping adapts to
    # any action space (delta-EEF ~0.001 vs raw servo ticks ~4095).
    action_range = float(np.mean(np.ptp(actions, axis=0)))
    jerk_divisor = max(action_range * 0.01, 2.0) if action_range > 200.0 else 2.0
    return float(np.clip(1.0 - normalized_jerk / jerk_divisor, 0.0, 1.0))


def compute_spectral_smoothness(actions: np.ndarray) -> float:
    """Score action smoothness via frequency-domain analysis.

    Computes the fraction of total spectral power in low frequencies.
    Smooth demonstrations concentrate power in low-frequency components.

    Args:
        actions: Array of shape (T, action_dim) for a single episode.

    Returns:
        Score in [0, 1]. 1.0 = all power in low frequencies.
    """
    if len(actions) < 16:
        return 0.5

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

    low_power_ratios = []
    for dim in range(actions.shape[1]):
        col = actions[:, dim]
        fft_coeffs = np.fft.rfft(col)
        power = np.abs(fft_coeffs) ** 2
        total_power = np.sum(power)
        if total_power < 1e-10:
            low_power_ratios.append(1.0)  # constant signal = perfectly smooth
            continue
        # Low frequency = bottom 25% of frequency bins
        n_low = max(1, len(power) // 4)
        low_ratio = np.sum(power[:n_low]) / total_power
        low_power_ratios.append(float(low_ratio))

    return float(np.clip(np.mean(low_power_ratios), 0.0, 1.0))


def compute_action_efficiency(actions: np.ndarray) -> float:
    """Score how direct/purposeful a trajectory is.

    Computes ratio of straight-line displacement to total path length.
    Expert demos take direct paths (ratio near 1.0), mediocre demos
    hesitate and backtrack (ratio near 0.0).

    Args:
        actions: Array of shape (T, action_dim) for a single episode.

    Returns:
        Score in [0, 1]. 1.0 = perfectly direct, 0.0 = all backtracking.
    """
    if len(actions) < 3:
        return 0.5

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

    # Total path length (sum of step-to-step distances)
    diffs = np.diff(actions, axis=0)
    step_lengths = np.sqrt(np.sum(diffs ** 2, axis=1))
    total_path = np.sum(step_lengths)

    if total_path < 1e-10:
        return 0.5  # no movement

    # Straight-line displacement (start to end)
    displacement = np.sqrt(np.sum((actions[-1] - actions[0]) ** 2))

    return float(np.clip(displacement / total_path, 0.0, 1.0))


def compute_completion_convergence(
    action_episodes: list[np.ndarray],
    state_episodes: list[np.ndarray] | None = None,
) -> float:
    """Score how well episodes converge to a tight final state.

    Uses state data if available (preferred), otherwise action data.
    Compares variance of last 20% vs first 20% of each episode.

    Args:
        action_episodes: List of per-episode action arrays.
        state_episodes: Optional list of per-episode state arrays.

    Returns:
        Score in [0, 1]. 1.0 = all episodes converge tightly.
    """
    episodes = state_episodes if state_episodes else action_episodes
    if not episodes:
        return 0.5

    scores = []
    for data in episodes:
        if len(data) < 10:
            continue
        if data.ndim == 1:
            data = data.reshape(-1, 1)
        split = max(2, len(data) // 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
        scores.append(1.0 / (1.0 + ratio))

    if not scores:
        return 0.5

    return float(np.clip(np.mean(scores), 0.0, 1.0))


def compute_episode_completion(states: np.ndarray) -> float:
    """Score whether an episode converges toward a goal state.

    Compares state variance in the last 20% vs first 20% of the
    trajectory. Decreasing variance suggests goal convergence.

    Args:
        states: Array of shape (T, state_dim) for a single episode.

    Returns:
        Score in [0, 1]. Higher = more convergent.
    """
    T = len(states)
    if T < 10:
        return 0.5

    split = max(2, T // 5)
    early_var = np.mean(np.var(states[:split], axis=0)) + 1e-8
    late_var = np.mean(np.var(states[-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 compute_observation_consistency(states_list: list[np.ndarray]) -> float:
    """Score consistency of initial states across episodes.

    Low variance of initial states means episodes start from similar
    positions, indicating controlled data collection.

    Args:
        states_list: List of state arrays, one per episode.

    Returns:
        Score in [0, 1]. Higher = more consistent initial states.
    """
    if len(states_list) < 2:
        return 0.5

    initial_states = []
    for states in states_list:
        if len(states) > 0:
            initial_states.append(states[0])

    if len(initial_states) < 2:
        return 0.5

    initial_arr = np.array(initial_states)
    per_dim_var = np.var(initial_arr, axis=0)
    mean_var = np.mean(per_dim_var)

    # Map variance to [0, 1]: low variance = high consistency
    # Sigmoid mapping: var=0 → 1.0, var=1.0 → ~0.27
    return float(np.clip(1.0 / (1.0 + mean_var), 0.0, 1.0))


def compute_temporal_autocorrelation(actions: np.ndarray) -> float:
    """Score temporal autocorrelation of actions.

    High autocorrelation indicates smooth, deliberate teleop control
    (actions at time t are correlated with t-1).

    Args:
        actions: Array of shape (T, action_dim) for a single episode.

    Returns:
        Score in [0, 1]. Higher = smoother control signal.
    """
    if len(actions) < 3:
        return 0.5

    correlations = []
    for dim in range(actions.shape[1] if actions.ndim == 2 else 1):
        col = actions[:, dim] if actions.ndim == 2 else actions
        col = col - np.mean(col)
        var = np.var(col)
        if var < 1e-10:
            correlations.append(0.0)
            continue
        # Lag-1 autocorrelation
        autocorr = np.mean(col[:-1] * col[1:]) / var
        correlations.append(float(np.clip(autocorr, -1.0, 1.0)))

    # Average autocorrelation, mapped to [0, 1]
    mean_corr = np.mean(correlations) if correlations else 0.0
    # Map from [-1, 1] to [0, 1]: high positive correlation is good
    return float(np.clip((mean_corr + 1.0) / 2.0, 0.0, 1.0))


def compute_action_state_mi(
    actions: np.ndarray,
    states: np.ndarray,
    max_samples: int = 10000,
) -> float:
    """Estimate mutual information between actions and states.

    Uses the KSG k-nearest-neighbor estimator (Kraskov et al. 2004)
    via sklearn, computing MI per action dimension then averaging.
    Falls back to a correlation proxy if sklearn is unavailable.

    Args:
        actions: Array of shape (N, action_dim).
        states: Array of shape (N, state_dim).
        max_samples: Max timesteps for efficiency.

    Returns:
        Score in [0, 1]. Higher = more mutual information.
    """
    if len(actions) < 5 or len(states) < 5:
        return 0.5

    n = min(len(actions), len(states), max_samples)
    a = actions[:n]
    s = states[:n]

    if a.ndim == 1:
        a = a.reshape(-1, 1)
    if s.ndim == 1:
        s = s.reshape(-1, 1)

    try:
        from sklearn.feature_selection import mutual_info_regression

        # Subsample if too many timesteps
        if n > max_samples:
            rng = np.random.RandomState(42)
            idx = rng.choice(n, max_samples, replace=False)
            a = a[idx]
            s = s[idx]

        # Compute MI per action dimension against all state dimensions
        mi_values = []
        for dim in range(a.shape[1]):
            mi = mutual_info_regression(
                s, a[:, dim],
                n_neighbors=5,
                random_state=42,
            )
            mi_values.append(float(np.mean(mi)))

        # Average MI across action dims, normalize to [0, 1]
        # MI is in nats; sigmoid mapping: mi=0→0, mi=1→0.63, mi=2→0.86
        mean_mi = float(np.mean(mi_values))
        score = 1.0 - np.exp(-mean_mi)
        return float(np.clip(score, 0.0, 1.0))

    except ImportError:
        # Fallback to correlation proxy if sklearn unavailable
        a_std = np.std(a, axis=0, keepdims=True)
        s_std = np.std(s, axis=0, keepdims=True)
        a_std[a_std < 1e-8] = 1.0
        s_std[s_std < 1e-8] = 1.0
        a_norm = (a - np.mean(a, axis=0, keepdims=True)) / a_std
        s_norm = (s - np.mean(s, axis=0, keepdims=True)) / s_std
        cross_corr = np.abs(a_norm.T @ s_norm) / n
        mi_proxy = float(np.mean(cross_corr))
        return float(np.clip(mi_proxy, 0.0, 1.0))


# ---------------------------------------------------------------------------
# Lightweight signal summary from action data
# ---------------------------------------------------------------------------


def _compute_signal_summary_from_actions(
    action_episodes: list[np.ndarray],
    zero_vel_threshold: float = 0.01,
    dead_ratio: float = 0.80,
) -> float:
    """Compute a signal health score [0, 1] directly from action episodes.

    This avoids requiring the full signal_diagnostics module as a dependency
    while still penalizing dead joints, high zero-velocity, and truncation.

    Returns:
        1.0 = perfectly healthy signals, 0.0 = severely degraded.
    """
    if not action_episodes:
        return 0.5

    all_actions = np.concatenate(action_episodes, axis=0)
    if all_actions.ndim == 1:
        all_actions = all_actions.reshape(-1, 1)

    action_dim = all_actions.shape[1]

    # 1. Dead joint penalty: each dead joint costs score proportionally
    dead_count = 0
    for d in range(action_dim):
        col = all_actions[:, d]
        zvr = float(np.mean(np.abs(col) < zero_vel_threshold))
        if zvr > dead_ratio:
            dead_count += 1

    # Dead joints are severe — each one is a big hit
    dead_penalty = min(1.0, (dead_count / max(1, action_dim)) * 3.0)

    # 2. Truncation penalty: fraction of episodes much shorter than median
    lengths = [len(ep) for ep in action_episodes]
    if len(lengths) >= 2:
        median_len = float(np.median(lengths))
        if median_len > 0:
            truncated = sum(1 for length in lengths if length < 0.5 * median_len)
            trunc_ratio = truncated / len(lengths)
        else:
            trunc_ratio = 0.0
    else:
        trunc_ratio = 0.0

    # 3. Global zero velocity penalty
    global_zvr = float(np.mean(np.abs(all_actions) < zero_vel_threshold))
    zvr_penalty = max(0.0, (global_zvr - 0.3) / 0.7)  # Penalize above 30%

    # Combine: dead joints are the most severe penalty
    score = 1.0 - dead_penalty * 0.6 - trunc_ratio * 0.25 - zvr_penalty * 0.15
    return float(np.clip(score, 0.0, 1.0))


# ---------------------------------------------------------------------------
# Main quality scoring
# ---------------------------------------------------------------------------


def compute_quality_score(
    info: DatasetInfo,
    coverage: FullCoverageReport,
    expected_episodes: int = DEFAULT_EXPECTED_EPISODES,
    visual_diversity: float | None = None,
    action_episodes: list[np.ndarray] | None = None,
    state_episodes: list[np.ndarray] | None = None,
    signal_diagnostics_summary: float | None = None,
    consistency: "ConsistencyReport | None" = None,
) -> QualityReport:
    """Compute an overall quality score and ready-to-train signal.

    Args:
        info: Dataset metadata.
        coverage: Full coverage analysis report.
        expected_episodes: Expected number of episodes for the task.
        visual_diversity: Optional visual diversity score (0.0-1.0).
        action_episodes: Optional list of per-episode action arrays for
            enhanced scoring.
        state_episodes: Optional list of per-episode state arrays for
            enhanced scoring.
        signal_diagnostics_summary: Optional summary score from signal
            diagnostics (0.0-1.0).
        consistency: Optional ConsistencyReport from consistency_metrics.py.
            Stored as a diagnostic component; does not affect the overall score.

    Returns:
        QualityReport with overall score, component scores, and confidence.
    """
    # Base component scores (always computed)
    episode_count_score = min(1.0, info.num_episodes / expected_episodes)

    component_scores: dict[str, float] = {
        "episode_count": episode_count_score,
        "position_diversity": coverage.position_diversity,
        "action_diversity": coverage.action_diversity,
        "episode_consistency": coverage.episode_consistency,
        "temporal_coverage": coverage.temporal_coverage,
    }

    # Enhanced metrics when action/state data is available
    use_enhanced = action_episodes is not None and len(action_episodes) > 0

    if use_enhanced:
        # Action quality: smoothness + autocorrelation + MI
        jerk_scores = [
            compute_action_smoothness(ep) for ep in action_episodes if len(ep) >= 4
        ]
        spectral_scores = [
            compute_spectral_smoothness(ep) for ep in action_episodes if len(ep) >= 16
        ]
        autocorr_scores = [
            compute_temporal_autocorrelation(ep) for ep in action_episodes if len(ep) >= 3
        ]

        jerk_smoothness = float(np.mean(jerk_scores)) if jerk_scores else 0.5
        spectral_smoothness = float(np.mean(spectral_scores)) if spectral_scores else None
        if spectral_smoothness is not None:
            action_smoothness = jerk_smoothness * 0.6 + spectral_smoothness * 0.4
        else:
            action_smoothness = jerk_smoothness
        action_autocorrelation = float(np.mean(autocorr_scores)) if autocorr_scores else 0.5

        mi_score = 0.5
        if state_episodes is not None and len(state_episodes) > 0:
            all_actions = np.concatenate(action_episodes, axis=0)
            all_states = np.concatenate(state_episodes, axis=0)
            mi_score = compute_action_state_mi(all_actions, all_states)

        action_quality = action_smoothness * 0.4 + action_autocorrelation * 0.3 + mi_score * 0.3
        component_scores["action_quality"] = action_quality

        # Episode quality: completion + observation consistency
        completion_scores = []
        if state_episodes is not None:
            for states in state_episodes:
                if len(states) >= 10:
                    completion_scores.append(compute_episode_completion(states))

        episode_completion = float(np.mean(completion_scores)) if completion_scores else 0.5

        obs_consistency = 0.5
        if state_episodes is not None:
            obs_consistency = compute_observation_consistency(state_episodes)

        episode_quality = episode_completion * 0.5 + obs_consistency * 0.5
        component_scores["episode_quality"] = episode_quality

        # Signal diagnostics summary — compute from actions if not provided
        if signal_diagnostics_summary is not None:
            sig_summary = signal_diagnostics_summary
        else:
            # Lightweight signal health check from action data
            sig_summary = _compute_signal_summary_from_actions(action_episodes)
        component_scores["signal_summary"] = sig_summary

        # Enhanced weighted combination
        # Coverage score = weighted average of base coverage metrics
        total_weight = sum(BASE_WEIGHTS.values())
        coverage_score = (
            episode_count_score * (BASE_WEIGHTS["episode_count"] / total_weight)
            + coverage.position_diversity * (BASE_WEIGHTS["position_diversity"] / total_weight)
            + coverage.action_diversity * (BASE_WEIGHTS["action_diversity"] / total_weight)
            + coverage.episode_consistency * (BASE_WEIGHTS["episode_consistency"] / total_weight)
            + coverage.temporal_coverage * (BASE_WEIGHTS["temporal_coverage"] / total_weight)
        )

        overall = (
            ENHANCED_WEIGHTS["coverage"] * coverage_score
            + ENHANCED_WEIGHTS["action_quality"] * action_quality
            + ENHANCED_WEIGHTS["episode_quality"] * episode_quality
            + ENHANCED_WEIGHTS["signal_summary"] * sig_summary
        )

        # Consistency is a diagnostic, not a scoring input.
        # Store it as a component for display but don't change the score.
        if consistency is not None:
            component_scores["consistency"] = consistency.overall_consistency
    else:
        # Standard scoring (backward compatible)
        if visual_diversity is not None:
            component_scores["visual_diversity"] = visual_diversity
            scale = 1.0 - VISUAL_DIVERSITY_WEIGHT
            weights = {k: v * scale for k, v in BASE_WEIGHTS.items()}
            weights["visual_diversity"] = VISUAL_DIVERSITY_WEIGHT
        else:
            weights = dict(BASE_WEIGHTS)

        overall = sum(
            weights[key] * component_scores[key] for key in weights if key in component_scores
        )

    # Confidence level
    if info.num_episodes > CONFIDENCE_HIGH_THRESHOLD:
        confidence = "high"
    elif info.num_episodes >= CONFIDENCE_MEDIUM_THRESHOLD:
        confidence = "medium"
    else:
        confidence = "low"

    # Estimated success rate — no pretense of prediction accuracy.
    # The quality score is a data quality metric, not a success predictor.
    # Map it to a broad range that's honest about uncertainty.
    center = overall
    width = CONFIDENCE_WIDTHS[confidence]
    estimated_success_rate = (
        max(0.0, center - width),
        min(1.0, center + width),
    )

    return QualityReport(
        overall_score=overall,
        component_scores=component_scores,
        ready_to_train=overall >= 0.65,
        confidence=confidence,
        estimated_success_rate=estimated_success_rate,
    )


def ready_to_train_assessment(quality: QualityReport) -> ReadyAssessment:
    """Generate a ready-to-train assessment with verdict and details.

    Args:
        quality: The quality report to assess.

    Returns:
        ReadyAssessment with verdict, reason, and issues.
    """
    blocking_issues: list[str] = []
    nice_to_haves: list[str] = []

    # Check for critically low components
    for name, score in quality.component_scores.items():
        label = name.replace("_", " ").title()
        if score < 0.3:
            blocking_issues.append(f"{label} is critically low ({score:.2f})")
        elif score < 0.5:
            nice_to_haves.append(f"Improve {label} (currently {score:.2f})")

    has_critical_component = any(s < 0.3 for s in quality.component_scores.values())

    if quality.overall_score >= 0.65 and not has_critical_component:
        verdict = "YES"
        reason = (
            f"Quality score {quality.overall_score:.2f} meets threshold "
            f"with no critically low components"
        )
    elif quality.overall_score >= 0.45:
        verdict = "NOT YET"
        if has_critical_component:
            reason = (
                f"Quality score {quality.overall_score:.2f} — "
                f"some components are critically low and need attention"
            )
        else:
            reason = (
                f"Quality score {quality.overall_score:.2f} — "
                f"close to threshold but needs improvement"
            )
    else:
        verdict = "NO"
        reason = (
            f"Quality score {quality.overall_score:.2f} is significantly below the 0.65 threshold"
        )

    return ReadyAssessment(
        verdict=verdict,
        reason=reason,
        blocking_issues=blocking_issues,
        nice_to_haves=nice_to_haves,
    )
