"""Dataset readiness grader — problem-based assessment for training readiness.

Every penalty corresponds to a specific, detectable problem with a name,
a measurement, and a threshold with a clear rationale.

No ML models, no LLM calls, no benchmark lookup.
Thresholds from: Sakr et al. 2024, Iglewicz & Hoaglin 1993, published
robot learning baselines (ACT, Diffusion Policy, RoboMimic).
"""

from __future__ import annotations

import logging
from dataclasses import dataclass, field

from orbit.analyzer.policy_fit import POLICY_PROFILES, normalize_policy_name

logger = logging.getLogger(__name__)


# ---------------------------------------------------------------------------
# Data models
# ---------------------------------------------------------------------------


@dataclass
class ReadinessResult:
    """Dataset readiness assessment for training."""

    grade: str  # A/B/C/D/F
    score: int  # 0-100
    summary: str  # "Good data with minor timing issues"
    strengths: list[str]  # ["High consistency (0.92)", ...]
    problems: list[str]  # ["2 dead joints", ...]
    top_action: str  # "Run orbit clean to remove 3 bad episodes"
    episode_count: int
    policy_type: str


# Minimum episodes per policy for critical failure check
_MIN_EPISODES = {
    "act": 15,
    "diffusion_policy": 25,
    "dp3": 25,
    "smolvla": 10,
    "bc": 50,
    "bc_rnn": 30,
}

# Minimum viable episodes per policy — the threshold below which training
# is unlikely to work at all.  Used for critical failure checks.
# Above this, training is viable; quality depends on other factors.
_VIABLE_EPISODES = {
    "act": 20,
    "diffusion_policy": 50,
    "dp3": 50,
    "smolvla": 10,
    "bc": 50,
    "bc_rnn": 30,
}

# Target episodes per policy — the point of diminishing returns.
# ACT paper: 50 demos (ALOHA). DP paper: 200 demos (PushT, RoboMimic).
# Used for sufficiency scoring (not critical failure).
_TARGET_EPISODES = {
    "act": 50,
    "diffusion_policy": 100,
    "dp3": 100,
    "smolvla": 30,
    "bc": 100,
    "bc_rnn": 80,
}



# ---------------------------------------------------------------------------
# Main grader
# ---------------------------------------------------------------------------


def _worst_episodes(per_episode_scores: dict[int, float], n: int = 3) -> str:
    """Return comma-separated indices of the N worst-scoring episodes."""
    if not per_episode_scores:
        return "unknown"
    ranked = sorted(per_episode_scores.items(), key=lambda x: x[1])
    worst = [str(idx) for idx, _ in ranked[:n]]
    return ", ".join(worst)


class DatasetReadinessGrader:
    """Grades dataset readiness for training based on detected problems.

    Problem catalog approach: every penalty maps to a named, measurable problem.
    No magic numbers — each threshold is from published research or statistics.

    Reference thresholds:
    - Sakr et al. 2024: consistency predicting success at 70-91% accuracy
    - Iglewicz & Hoaglin 1993: Modified Z-Score 3.5 threshold for outliers
    - RoboMimic: expert vs novice demonstrator impact
    - LeRobot community: episode count best practices
    """

    def grade(
        self,
        quality_score: float,
        episode_consistency: float,
        signal_health: dict,
        policy_fit_score: float,
        episode_ranking: dict,
        action_divergence: float,
        episode_count: int,
        policy_type: str,
        temporal_alignment: dict | None = None,
        consistency: "ConsistencyReport | None" = None,
        outlier_report: "DatasetOutlierReport | None" = None,
        vision_quality: "VisionQualityResult | None" = None,
        proxy_result: "ProxyTrainingResult | None" = None,
    ) -> ReadinessResult:
        """Grade dataset readiness for training.

        Args:
            quality_score: 0-100 from quality.py (overall_score * 100).
            episode_consistency: 0-1 from coverage.py.
            signal_health: Dict with 'dead_dimensions', 'blockers', 'warnings',
                and optionally 'clipping_joints' (list of joint indices with
                clipping rate > threshold).
            policy_fit_score: 0-1 from policy_fit.py.
            episode_ranking: Dict with 'total', 'remove_count', 'review_count'.
            action_divergence: 0-1 from action_divergence.py.
            episode_count: Total number of episodes.
            policy_type: Canonical policy name.
            temporal_alignment: Optional dict with 'misaligned_pct' (0-1).
            consistency: Optional ConsistencyReport from consistency_metrics.py.
            outlier_report: Optional DatasetOutlierReport from outlier_detection.py.
            vision_quality: Optional VisionQualityResult from vision_quality.py.
            proxy_result: Optional ProxyTrainingResult from proxy_trainer.py.

        Returns:
            ReadinessResult with grade, score, and actionable feedback.
        """
        policy_type = normalize_policy_name(policy_type)
        strengths: list[str] = []
        problems: list[str] = []

        # --- Critical failures (automatic F) --- always checked first
        dead_dims = signal_health.get("dead_dimensions", [])
        total_joints = signal_health.get("total_joints", 1)
        dead_ratio = len(dead_dims) / max(total_joints, 1)

        # Critical failure only when the MAJORITY of the action space is
        # non-functional.  Many real datasets have padded/unused action
        # dimensions (e.g., 7-DOF action space where 3 dims are constant
        # padding).  3/7 dead = 43% — that's padded data, not a hardware
        # failure.  But 5/7 dead = 71% — the robot isn't working.
        if dead_ratio > 0.60:
            return self._critical_failure(
                "Hardware problem: too many dead joints "
                f"({len(dead_dims)} of {total_joints})",
                episode_count, policy_type,
            )

        min_eps = _MIN_EPISODES.get(policy_type, 15)
        if episode_count < min_eps:
            return self._critical_failure(
                f"Not enough demonstrations ({episode_count}) — "
                f"{policy_type} needs at least {min_eps}",
                episode_count, policy_type,
            )

        misaligned_pct = 0.0
        if temporal_alignment:
            misaligned_pct = temporal_alignment.get("misaligned_pct", 0.0)
        if misaligned_pct > 0.50:
            return self._critical_failure(
                f"Severe timing issues on {misaligned_pct:.0%} of episodes",
                episode_count, policy_type,
            )

        total_eps = episode_ranking.get("total", 1)
        remove_count = episode_ranking.get("remove_count", 0)
        remove_pct = remove_count / max(total_eps, 1)
        if remove_pct > 0.50:
            return self._critical_failure(
                f"Majority of demonstrations are unusable "
                f"({remove_count} of {total_eps} flagged for removal)",
                episode_count, policy_type,
            )

        # --- Problem-based penalty scoring ---
        base = 100
        multimodal_policy = policy_type in ("diffusion_policy", "dp3")

        # ============================================================
        # STRUCTURAL PROBLEM: Dead joints
        # Measurement: n_dead / n_total ratio
        # Penalty: ratio * 40 (e.g., 3/14 dead = -8.6 points)
        # ============================================================
        if dead_dims:
            penalty = int(round(dead_ratio * 40))
            penalty = max(penalty, len(dead_dims))  # at least 1 per dead joint
            base -= penalty
            problems.append(
                f"{len(dead_dims)} dead {'joint' if len(dead_dims) == 1 else 'joints'} — wastes model capacity"
            )

        # ============================================================
        # STRUCTURAL PROBLEM: Clipping joints
        # Measurement: n_clipping joints hitting actuator limits >10% of time
        # Penalty: min(n_clipping * 2, 8)
        # ============================================================
        clipping_joints = signal_health.get("clipping_joints", [])
        n_clipping = len(clipping_joints)
        if n_clipping > 0:
            clip_penalty = min(8, 2 + n_clipping)
            base -= clip_penalty
            problems.append(f"{n_clipping} {'joint' if n_clipping == 1 else 'joints'} clipping")
        else:
            if not dead_dims:
                strengths.append("Clean signal health — no dead joints or clipping")

        # ============================================================
        # STRUCTURAL PROBLEM: Episode outliers (Modified Z-Score detection)
        # Measurement: outlier_fraction from outlier_detection.py
        # Penalty: 0 if <10%, 8 if 10-25%, 15 if >25%
        # ============================================================
        if outlier_report is not None:
            of = outlier_report.outlier_fraction
            if of > 0.25:
                base -= 15
                problems.append(
                    f"{outlier_report.outlier_count} outlier episodes ({of:.0%}) — "
                    f"dataset is {outlier_report.dataset_health}"
                )
            elif of > 0.10:
                base -= 8
                problems.append(
                    f"{outlier_report.outlier_count} outlier episodes ({of:.0%}) — "
                    f"dataset is {outlier_report.dataset_health}"
                )
            elif of <= 0.05:
                strengths.append(
                    f"Very few outlier episodes ({outlier_report.outlier_count}/{outlier_report.total_episodes})"
                )

        # ============================================================
        # STRUCTURAL PROBLEM: Episode quality (flagged for removal)
        # Measurement: remove_pct from episode ranker
        # Penalty: 15 if >20%, 8 if >10%
        # ============================================================
        review_count = episode_ranking.get("review_count", 0)
        review_pct = review_count / max(total_eps, 1)

        if remove_pct > 0.20:
            base -= 15
            problems.append(
                f"{remove_count} episodes ({remove_pct:.0%}) flagged for removal"
            )
        elif remove_pct > 0.10:
            base -= 8
            problems.append(
                f"{remove_count} episodes ({remove_pct:.0%}) flagged for removal"
            )

        # Review flags: only penalize when many reviews AND actual removals
        if review_pct > 0.50 and remove_pct > 0.05:
            base -= 5
            problems.append(
                f"{review_count} episodes ({review_pct:.0%}) need review"
            )

        if remove_pct <= 0.05 and review_pct <= 0.10:
            strengths.append("Episode quality is high — very few flagged episodes")

        # ============================================================
        # STRUCTURAL PROBLEM: Insufficient episodes for policy
        # Measurement: actual vs viable/target thresholds
        # Penalty: 0 if >= viable, small if between viable and target
        # The viable threshold is what you need to train at all;
        # the target is where diminishing returns kick in.
        # ============================================================
        viable_eps = _VIABLE_EPISODES.get(policy_type, 30)
        target_eps = _TARGET_EPISODES.get(policy_type, 100)
        if episode_count < viable_eps:
            base -= 15
            problems.append(
                f"Only {episode_count} episodes — {policy_type} needs at least {viable_eps}"
            )
        elif episode_count < target_eps:
            # Small penalty: viable but below ideal
            ratio = episode_count / max(target_eps, 1)
            penalty = int(5 * (1.0 - ratio))
            base -= penalty
            if penalty >= 3:
                problems.append(
                    f"{episode_count} episodes — more data would help "
                    f"(target: {target_eps}+ for {policy_type})"
                )
            else:
                strengths.append(f"Sufficient episodes ({episode_count}) for {policy_type}")
        else:
            strengths.append(f"Sufficient episodes ({episode_count}) for {policy_type}")

        # ============================================================
        # CONSISTENCY PROBLEM: High cross-episode CVs (Sakr metrics)
        # Measurement: count of CVs > 0.50 across 8 metrics
        # Penalty: 5 per elevated metric, capped at 15
        # Note: halved for multimodal policies (DP, DP3)
        # Note: softened when action divergence is low — high CVs with
        #   low divergence means diverse-but-consistent demonstrations
        #   (good coverage), not conflicting strategies.
        # ============================================================
        if consistency is not None:
            if consistency.overall_consistency >= 0.65:
                strengths.append(
                    f"Good cross-episode consistency ({consistency.overall_consistency:.2f})"
                )
            elif consistency.overall_consistency < 0.40:
                base -= 5
                problems.append(
                    f"Low cross-episode consistency ({consistency.overall_consistency:.2f}) "
                    f"— demonstrations are too variable"
                )

            # Count how many CVs are elevated (CV > 0.60 — raised from 0.50
            # to reduce false positives on real-world data where some
            # trajectory variation is natural and healthy)
            all_cvs = [
                consistency.cartesian_jerk_cv,
                consistency.path_length_cv,
                consistency.joint_path_length_cv,
                consistency.joint_jerk_cv,
                consistency.curvature_cv,
                consistency.effort_cv,
                consistency.joint_limit_cv,
                consistency.orientation_cv,
            ]
            high_cv_count = sum(1 for cv in all_cvs if cv > 0.60)
            # Only penalize if majority of CVs are elevated
            if high_cv_count > len(all_cvs) // 2:
                cv_penalty = min(10, high_cv_count * 3)
                if multimodal_policy:
                    cv_penalty = cv_penalty // 2
                if action_divergence < 0.10:
                    cv_penalty = cv_penalty // 2
                base -= cv_penalty
                problems.append(
                    f"{high_cv_count}/8 consistency metrics are elevated (CV > 0.60)"
                )
            elif high_cv_count >= 3:
                cv_penalty = min(6, high_cv_count * 2)
                if multimodal_policy:
                    cv_penalty = cv_penalty // 2
                if action_divergence < 0.10:
                    cv_penalty = cv_penalty // 2
                base -= cv_penalty

            # Specific diagnostics
            if consistency.cartesian_jerk_cv > 0.5:
                worst = _worst_episodes(consistency.per_episode_scores, 3)
                problems.append(
                    f"Jerk varies across episodes (CV={consistency.cartesian_jerk_cv:.2f}) "
                    f"— some demos are much jerkier. Worst: episodes {worst}"
                )
            if consistency.path_length_cv > 0.5:
                worst = _worst_episodes(consistency.per_episode_scores, 3)
                problems.append(
                    f"Path length varies across episodes (CV={consistency.path_length_cv:.2f}) "
                    f"— demos take very different routes. Check episodes {worst}"
                )
            if consistency.effort_cv > 0.5:
                problems.append(
                    f"Effort varies across episodes (CV={consistency.effort_cv:.2f}) "
                    f"— force/torque usage is inconsistent"
                )

        # ============================================================
        # CONSISTENCY PROBLEM: Episode consistency (coverage-based)
        # Measurement: episode_consistency from coverage.py
        # Only applied when no dead joints (to avoid double-counting)
        # ============================================================
        if dead_dims and episode_consistency < 0.65:
            if episode_consistency < 0.30:
                base -= 6
                problems.append(
                    f"Low consistency ({episode_consistency:.2f}) — "
                    "partly caused by dead joints"
                )
        elif episode_consistency < 0.30:
            base -= 10
            problems.append(f"Very low consistency ({episode_consistency:.2f})")
        elif episode_consistency < 0.45:
            base -= 5
            problems.append(f"Low consistency ({episode_consistency:.2f})")
        elif episode_consistency < 0.55:
            base -= 2
        else:
            strengths.append(f"Good consistency ({episode_consistency:.2f})")

        # ============================================================
        # COVERAGE PROBLEM: Action divergence (same state → different actions)
        # Measurement: divergence score from action_divergence.py
        #
        # Continuous piecewise-linear penalty.  Belkhale et al. (NeurIPS
        # 2023) identified action divergence as the primary quality axis
        # for imitation learning.  The penalty scales linearly from 0 at
        # div=0.05 to 20 at div=0.35, then caps.
        #
        # Calibration from 27 datasets across 7 platforms:
        #   div < 0.10: clean (koch_1=0.09, insertion_scripted=0.095)
        #   div 0.10-0.13: clean boundary (pusht=0.123, tape=0.100)
        #   div 0.13-0.20: moderate (battery=0.161, cups=0.142)
        #   div > 0.20: high (fork=0.22, cable_routing=0.25)
        #   div > 0.35: severe (droid=0.44, xarm_lift=0.46)
        #
        # Halved for multimodal policies (DP, DP3) which tolerate
        # divergence by design (Chi et al. 2023 Diffusion Policy).
        # ============================================================
        if action_divergence <= 0.05:
            div_penalty = 0
        elif action_divergence <= 0.12:
            # Gentle: 0→3 over [0.05, 0.12].  Clean data (div < 0.12)
            # gets a small penalty — not enough to change grade.
            div_penalty = int(round(3 * (action_divergence - 0.05) / 0.07))
        elif action_divergence <= 0.20:
            # Steep: 3→10 over [0.12, 0.20].  This is the critical range
            # where moderate divergence costs enough to distinguish A from B.
            div_penalty = 3 + int(round(7 * (action_divergence - 0.12) / 0.08))
        elif action_divergence <= 0.35:
            # Gradual: 10→15 over [0.20, 0.35].  Already-high divergence
            # has diminishing marginal impact.
            div_penalty = 10 + int(round(5 * (action_divergence - 0.20) / 0.15))
        else:
            div_penalty = 15

        if multimodal_policy:
            div_penalty = div_penalty // 2

        if div_penalty > 0:
            base -= div_penalty
            if div_penalty >= 10:
                problems.append(
                    f"High action divergence ({action_divergence:.2f}) — "
                    "demos contradict each other in similar states"
                )
            elif div_penalty >= 4:
                problems.append(
                    f"Moderate action divergence ({action_divergence:.2f})"
                )
        elif action_divergence <= 0.05:
            strengths.append(
                f"Low action divergence ({action_divergence:.2f}) — consistent demonstrations"
            )

        # ============================================================
        # STRUCTURAL PROBLEM: Policy fit
        # Measurement: policy_fit_score from policy_fit.py
        # Penalty: proportional for <0.70, cliff at <0.50
        # ============================================================
        if policy_fit_score < 0.50:
            base -= 12
            problems.append(
                f"Poor policy fit ({policy_fit_score:.2f}) — "
                "consider a different policy architecture"
            )
        elif policy_fit_score < 0.70:
            penalty = int(8 * (0.70 - policy_fit_score) / 0.20)
            base -= penalty
            if penalty >= 4:
                problems.append(f"Moderate policy fit ({policy_fit_score:.2f})")
        else:
            strengths.append(f"Good policy fit ({policy_fit_score:.2f})")

        # ============================================================
        # VISION PROBLEMS (from Gemini, only if images available)
        # Measurement: structured VLM assessment
        # Penalties: task_visible→-5, incomplete→-10, inconsistent→-8, quality→-3
        # ============================================================
        if vision_quality is not None and vision_quality.total_penalty != 0:
            base += vision_quality.total_penalty  # total_penalty is negative
            for name, pts in vision_quality.grade_adjustments.items():
                problems.append(f"Vision: {name} ({pts:+d} pts)")

        # ============================================================
        # PROXY TRAINING SIGNAL (only if --proxy was used)
        # GO + val_loss < 0.1: boost +5 (capped at A)
        # GO + val_loss 0.1-0.5: no change
        # CAUTION: no change
        # NO-GO: penalty of 10
        # ============================================================
        if proxy_result is not None:
            if proxy_result.go_no_go == "GO" and proxy_result.validation_loss_final < 0.1:
                base = min(100, base + 5)
                strengths.append(
                    f"Proxy training signal: GO (val_loss={proxy_result.validation_loss_final:.4f})"
                )
            elif proxy_result.go_no_go == "GO":
                strengths.append(
                    f"Proxy training signal: GO (val_loss={proxy_result.validation_loss_final:.4f})"
                )
            elif proxy_result.go_no_go == "NO-GO":
                base -= 10
                problems.append(
                    f"Proxy training signal: NO-GO — {proxy_result.reason}"
                )
        else:
            # When proxy not run, suggest for B/C grades (checked after score computed)
            pass

        # ============================================================
        # CONJUNCTIVE CAPS: each component must independently meet a
        # floor.  Without these, moderate problems compensate each
        # other in the additive model (e.g. 3 dead joints + low
        # consistency + flagged episodes still scores B).
        #
        # Approach from medical competency assessment (Meyers 2018)
        # and FICO scoring: conjunctive minimums, not multiplicative.
        # ============================================================
        # Conjunctive soft caps: prevent good grades when a single
        # dimension is severely bad, using extra additive penalties
        # rather than hard cliffs (which create discontinuities).
        if action_divergence > 0.20 and not multimodal_policy:
            # Extra penalty for unimodal policies (BC, BC-RNN, ACT) when
            # divergence exceeds the moderate/high boundary.  Threshold
            # lowered from 0.30 to 0.20 (not 0.15) because centroid
            # opposition detection now correctly boosts scores into the
            # base penalty range — a lower threshold double-counts.
            # Gentle slope (25) avoids stacking too hard with base penalty.
            extra = int(min(10, (action_divergence - 0.20) * 25))
            base -= extra
            if extra >= 3 and "high divergence" not in " ".join(problems).lower():
                problems.append(
                    f"High action divergence for unimodal policy ({action_divergence:.2f}) — "
                    "consider Diffusion Policy or filtering demos"
                )
        if outlier_report is not None and outlier_report.outlier_fraction > 0.20:
            extra = int(min(8, (outlier_report.outlier_fraction - 0.20) * 40))
            base -= extra
        if policy_fit_score < 0.40:
            base -= 5  # extra penalty for very poor fit

        # ============================================================
        # Diminishing returns on stacked penalties.
        #
        # Max raw penalty sum can exceed 100 (audit found ~198 possible).
        # Real-world impact of the Nth problem is less than the 1st.
        # Apply compression: first 30 points at full weight, remainder
        # at 60% effectiveness.  This prevents catastrophic stacking
        # while preserving the ordering of individual penalties.
        #
        # Floor at 15: any dataset with real demonstrations is at least
        # marginally useful.  F grades (< 50) still exist but F(0)
        # requires a critical failure, not just penalty accumulation.
        # ============================================================
        raw_penalty = 100 - base
        if raw_penalty > 30:
            excess = raw_penalty - 30
            raw_penalty = 30 + int(excess * 0.6)
        score = max(15, min(100, 100 - raw_penalty))

        # Map to grade — calibrated against 82 datasets with known outcomes.
        # A = will train successfully with standard settings.
        # B = should train well, minor issues unlikely to block.
        # C = trainable but may need cleaning or adjustment.
        # D = significant issues, training likely to underperform.
        # F = critical problems, training will fail or be useless.
        if score >= 85:
            grade = "A"
            summary = "Ready to train — expect strong results"
        elif score >= 72:
            grade = "B"
            summary = "Good data — minor issues, should train well"
        elif score >= 58:
            grade = "C"
            summary = "Usable but has problems — run orbit clean first"
        elif score >= 40:
            grade = "D"
            summary = "Significant issues — collect more/better data"
        else:
            grade = "F"
            summary = "Critical problems — address before training"

        # Add proxy suggestion for B/C when not run
        if proxy_result is None and grade in ("B", "C"):
            problems.append("Run with --proxy for a more reliable assessment")

        # Determine top action
        top_action = self._pick_top_action(
            problems, dead_dims, remove_count, episode_count,
            policy_type, action_divergence, episode_consistency,
        )

        return ReadinessResult(
            grade=grade,
            score=score,
            summary=summary,
            strengths=strengths,
            problems=problems,
            top_action=top_action,
            episode_count=episode_count,
            policy_type=policy_type,
        )

    @staticmethod
    def _critical_failure(
        reason: str,
        episode_count: int,
        policy_type: str,
    ) -> ReadinessResult:
        return ReadinessResult(
            grade="F",
            score=0,
            summary="Critical problems — address before training",
            strengths=[],
            problems=[reason],
            top_action=reason,
            episode_count=episode_count,
            policy_type=policy_type,
        )

    @staticmethod
    def _pick_top_action(
        problems: list[str],
        dead_dims: list,
        remove_count: int,
        episode_count: int,
        policy_type: str,
        action_divergence: float,
        consistency: float,
    ) -> str:
        # Priority: dead joints > bad episodes > divergence > more data > consistency
        if dead_dims:
            return (
                f"Fix dead joints {dead_dims} — "
                "re-check hardware and re-record affected episodes"
            )
        # Only recommend cleaning when removals are significant
        remove_threshold = max(2, int(episode_count * 0.05))
        if remove_count >= remove_threshold:
            return f"Run `orbit clean` to remove {remove_count} bad episode(s)"
        if action_divergence > 0.20:
            return (
                "High action divergence — consider filtering to a single strategy "
                "or switching to Diffusion Policy"
            )
        ideal = _TARGET_EPISODES.get(policy_type, 100)
        if episode_count < ideal * 0.5:
            more = ideal - episode_count
            return f"Collect ~{more} more demonstrations to reach the sweet spot for {policy_type}"
        if consistency < 0.65:
            return "Improve demonstration consistency — practice the task motion before recording"
        if problems:
            return problems[0]
        return "Data looks good — run `orbit suggest` to get a training command"
