"""ORBIT CLI — Data strategy copilot for robot policy training."""

from __future__ import annotations

import json
import logging
import sys

import click

logger = logging.getLogger(__name__)


def _json_default(obj):
    """Handle numpy types for JSON serialization."""
    import numpy as np

    if isinstance(obj, (np.bool_,)):
        return bool(obj)
    if isinstance(obj, (np.integer,)):
        return int(obj)
    if isinstance(obj, (np.floating,)):
        return float(obj)
    if isinstance(obj, np.ndarray):
        return obj.tolist()
    raise TypeError(f"Object of type {type(obj).__name__} is not JSON serializable")


class _OrbitGroup(click.Group):
    """Custom group that catches unexpected errors and shows clean messages."""

    def invoke(self, ctx):
        try:
            return super().invoke(ctx)
        except click.exceptions.Exit:
            raise
        except click.exceptions.Abort:
            raise
        except click.ClickException:
            raise
        except SystemExit:
            raise
        except KeyboardInterrupt:
            raise SystemExit(130)
        except Exception as e:
            from orbit.utils.display import console

            console.print(f"\n[bold red]Error:[/bold red] {e}")
            console.print("[dim]Run with ORBIT_DEBUG=1 for full traceback.[/dim]")
            import os

            if os.environ.get("ORBIT_DEBUG"):
                raise
            raise SystemExit(1) from None


@click.group(cls=_OrbitGroup)
@click.version_option(version=None, package_name="orbit-robotics", prog_name="orbit")
@click.option("--quiet", "-q", is_flag=True, help="Minimal output (grade + exit code only).")
@click.pass_context
def main(ctx, quiet):
    """Data strategy copilot for robot policy training."""
    ctx.ensure_object(dict)
    ctx.obj["quiet"] = quiet


@main.command()
@click.argument("dataset_repo")
@click.option("--episodes", type=int, default=None, help="Limit analysis to N episodes.")
@click.option("--no-cache", is_flag=True, help="Skip cache.")
@click.option("--json", "as_json", is_flag=True, help="Output as JSON.")
@click.option("--skip-embeddings", is_flag=True, help="Skip SigLIP embedding analysis.")
@click.option("--skip-coverage", is_flag=True, help="Skip coverage analysis.")
@click.option(
    "--device",
    type=click.Choice(["cpu", "cuda", "mps"]),
    default=None,
    help="Device for embedding computation.",
)
@click.option("--max-frames", type=int, default=4, help="Frames per episode for embeddings.")
@click.option("--skip-signal-diagnostics", is_flag=True, help="Skip signal diagnostics.")
@click.option("--vlm", is_flag=True, help="Enable optional VLM-powered AI assessment.")
@click.option(
    "--policy",
    default="auto",
    type=click.Choice(["auto", "act", "diffusion_policy", "dp3", "bc", "bc_rnn", "smolvla"]),
    help="Policy type for fit analysis.",
)
@click.option("--full", is_flag=True, help="Analyze all episodes (no sampling).")
@click.option("--ai", is_flag=True, help="AI quality judge — verifies A grades with Gemini (~$0.001, ~6s, requires GOOGLE_API_KEY).")
@click.option("--deep", is_flag=True, help="Run AI-powered deep analysis (requires GOOGLE_API_KEY).")
@click.option("--proxy", is_flag=True, help="Run proxy BC training for go/no-go signal (requires torch, ~2-3 min).")
@click.option("--skip-proxy", is_flag=True, help="Skip automatic proxy training for small datasets.")
@click.option(
    "--format",
    "source_format",
    default=None,
    type=click.Choice(["auto", "lerobot", "lerobot-local", "rlds", "hdf5", "rosbag", "directory"]),
    help="Override format auto-detection.",
)
@click.option("--action-topic", default=None, help="ROS topic for actions (rosbag only).")
@click.option("--state-topic", default=None, help="ROS topic for states (rosbag only).")
@click.option("--action-key", default=None, help="Column name for action data (e.g. 'end_pose', 'joint_command').")
@click.option("--state-key", default=None, help="Column name for state data (e.g. 'qpos', 'joint_positions').")
@click.option("--episode-gap", default=None, type=float, help="Seconds gap for episode segmentation (rosbag only).")
@click.option("--fps", default=None, type=float, help="Override FPS when not in metadata.")
@click.option("--detail", is_flag=True, help="Show full detailed report (default is compact summary).")
@click.option("--push", is_flag=True, help="Push results to your ORBIT dashboard after analysis.")
def analyze(
    dataset_repo: str,
    episodes: int | None,
    no_cache: bool,
    as_json: bool,
    skip_embeddings: bool,
    skip_coverage: bool,
    device: str | None,
    max_frames: int,
    skip_signal_diagnostics: bool,
    vlm: bool,
    policy: str,
    full: bool,
    ai: bool,
    deep: bool,
    proxy: bool,
    skip_proxy: bool,
    source_format: str | None,
    action_topic: str | None,
    state_topic: str | None,
    action_key: str | None,
    state_key: str | None,
    episode_gap: float | None,
    fps: float | None,
    detail: bool = False,
    push: bool = False,
):
    """Analyze robot training data.

    DATASET_REPO can be a HuggingFace repo ID (e.g. lerobot/pusht), a local
    directory, or a file path (.hdf5, .bag, .mcap).
    """
    import numpy as np

    from orbit.analyzer.coverage import compute_coverage
    from orbit.analyzer.dataset_loader import load_dataset_info, load_episode_data
    from orbit.analyzer.policy_fit import PolicyFitAnalyzer
    from orbit.analyzer.quality import compute_quality_score, ready_to_train_assessment
    from orbit.analyzer.recommendations import generate_recommendations
    from orbit.analyzer.success_predictor import DatasetReadinessGrader
    from orbit.analyzer.task_context import TaskContextAnalyzer
    from orbit.utils.display import console

    # Build loader kwargs for format-specific options
    loader_kwargs: dict = {}
    if action_topic:
        loader_kwargs["action_topic"] = action_topic
    if state_topic:
        loader_kwargs["state_topic"] = state_topic
    if action_key:
        loader_kwargs["action_key"] = action_key
    if state_key:
        loader_kwargs["state_key"] = state_key
    if episode_gap is not None:
        loader_kwargs["episode_gap"] = episode_gap
    if fps is not None:
        loader_kwargs["fps"] = fps

    # Validate dataset_repo
    if not dataset_repo or not dataset_repo.strip():
        if as_json:
            click.echo(json.dumps({"error": "Dataset source must not be empty"}))
        else:
            console.print("\n[bold red]Error:[/bold red] Dataset source must not be empty")
        sys.exit(1)

    try:
        info = load_dataset_info(dataset_repo, format=source_format, **loader_kwargs)
    except (ValueError, ConnectionError) as e:
        if as_json:
            click.echo(json.dumps({"error": str(e)}))
        else:
            console.print(f"\n[bold red]Error:[/bold red] {e}")
        sys.exit(1)

    # Load episode data for coverage analysis
    coverage = None
    episode_data = None
    if not skip_coverage:
        try:
            ep_indices = list(range(min(episodes, info.num_episodes))) if episodes else None
            # Analyze all episodes for small datasets (≤100), sample for larger ones
            if full:
                max_ep = info.num_episodes
            elif episodes:
                max_ep = episodes
            elif info.num_episodes <= 100:
                max_ep = info.num_episodes
            else:
                max_ep = 100
            episode_data = load_episode_data(
                dataset_repo, episode_indices=ep_indices, max_episodes=max_ep,
                format=source_format, **loader_kwargs,
            )
            if not as_json and episode_data is not None and not episode_data.empty:
                n_loaded = episode_data["episode_index"].nunique()
                if n_loaded < info.num_episodes:
                    if max_ep and max_ep < info.num_episodes:
                        console.print(
                            f"[dim]Analyzing {n_loaded}/{info.num_episodes} episodes "
                            f"(sampling for speed, use --full for all)[/dim]"
                        )
                    else:
                        # Requested all but some failed to load
                        console.print(
                            f"[dim]Loaded {n_loaded}/{info.num_episodes} episodes[/dim]"
                        )
        except (ValueError, ConnectionError):
            episode_data = None

        if episode_data is not None and not episode_data.empty:
            coverage = compute_coverage(info, episode_data)

    is_estimated = coverage is None
    if is_estimated:
        coverage = _synthetic_coverage(info)

    # Signal diagnostics (optional)
    signal_diag = None
    if not skip_signal_diagnostics and episode_data is not None and not episode_data.empty:
        try:
            from orbit.analyzer.coverage import _extract_array
            from orbit.analyzer.signal_diagnostics import compute_signal_diagnostics

            action_arr = (
                _extract_array(episode_data["action"]) if "action" in episode_data.columns else None
            )
            if action_arr is not None and action_arr.size > 0:
                # Group by episode
                ep_groups = episode_data.groupby("episode_index")
                action_episodes = []
                for _, ep_df in ep_groups:
                    ep_actions = _extract_array(ep_df["action"])
                    if ep_actions is not None and ep_actions.size > 0:
                        action_episodes.append(ep_actions)

                if action_episodes:
                    signal_diag = compute_signal_diagnostics(action_episodes)
                    # Fix episode count in passing/warning messages to reflect
                    # total dataset size, not just the sampled episodes.
                    n_sampled = len(action_episodes)
                    n_total = info.num_episodes
                    if n_sampled != n_total:
                        signal_diag.passing = [
                            msg.replace(f"({n_sampled})", f"({n_total})")
                            for msg in signal_diag.passing
                        ]
                        signal_diag.warnings = [
                            msg.replace(f"({n_sampled})", f"({n_total})")
                            for msg in signal_diag.warnings
                        ]
        except Exception as e:
            if not as_json:
                console.print(f"\n[dim]Signal diagnostics skipped — {e}[/dim]")

    # Embedding analysis (optional)
    embedding_analysis = None
    visual_diversity: float | None = None

    if not skip_embeddings:
        try:
            from orbit.analyzer.embeddings import (
                analyze_embeddings,
                compute_frame_embeddings,
            )

            max_ep = episodes if episodes else 100

            if as_json:
                import io
                from contextlib import redirect_stderr, redirect_stdout

                buf = io.StringIO()
                with redirect_stdout(buf), redirect_stderr(buf):
                    emb_result = compute_frame_embeddings(
                        dataset_repo,
                        info,
                        max_frames_per_episode=max_frames,
                        max_episodes=max_ep,
                        device=device,
                        no_cache=no_cache,
                    )
            else:
                emb_result = compute_frame_embeddings(
                    dataset_repo,
                    info,
                    max_frames_per_episode=max_frames,
                    max_episodes=max_ep,
                    device=device,
                    no_cache=no_cache,
                )
            embedding_analysis = analyze_embeddings(emb_result)
            visual_diversity = embedding_analysis.diversity_score
        except ImportError:
            if not as_json:
                console.print(
                    "\n[dim]Visual analysis skipped — install torch and "
                    "transformers for embedding analysis[/dim]"
                )
        except Exception as e:
            if not as_json:
                console.print(f"\n[dim]Visual analysis skipped — {e}[/dim]")

    # Build action episodes list for new analyzers
    action_episodes = []
    state_episodes = []
    if episode_data is not None and not episode_data.empty:
        try:
            from orbit.analyzer.coverage import _extract_array

            ep_groups = episode_data.groupby("episode_index")
            for _, ep_df in ep_groups:
                if "action" in ep_df.columns:
                    ep_actions = _extract_array(ep_df["action"])
                    if ep_actions is not None and ep_actions.size > 0:
                        action_episodes.append(ep_actions)
                if "state" in ep_df.columns:
                    ep_states = _extract_array(ep_df["state"])
                    if ep_states is not None and ep_states.size > 0:
                        state_episodes.append(ep_states)
        except Exception as e:
            if not as_json:
                console.print(f"\n[dim]Episode extraction skipped — {e}[/dim]")

    # Warn when action data could not be found
    if episode_data is not None and not episode_data.empty and not action_episodes:
        available = [c for c in episode_data.columns
                     if c not in ("episode_index", "frame_index", "timestamp")]
        if not as_json:
            console.print(
                "\n[bold yellow]Warning:[/bold yellow] Could not find action data. "
                "Divergence, policy fit, and readiness grade may be inaccurate."
            )
            if available:
                console.print(
                    f"[dim]  Available columns: {', '.join(available[:15])}[/dim]"
                )
            console.print(
                "[dim]  Specify the action column with: --action-key <column_name>[/dim]"
            )

    # --- Array Intelligence: profile action/state arrays ---
    array_profiles = {}
    if action_episodes:
        try:
            from orbit.analyzer.array_intelligence import profile_array

            array_profiles["action"] = profile_array(action_episodes)
            if not as_json:
                ap = array_profiles["action"]
                console.print(
                    f"[dim]Array profile: {ap.action_type} "
                    f"({ap.num_dimensions}D, "
                    f"{'bimanual' if ap.is_bimanual else 'single-arm'}"
                    f"{', gripper=' + str(ap.gripper_dimensions) if ap.gripper_dimensions else ''}"
                    f"{', dead=' + str(ap.dead_dimensions) if ap.dead_dimensions else ''}"
                    f")[/dim]"
                )
            if state_episodes:
                array_profiles["state"] = profile_array(state_episodes)
        except Exception as e:
            if not as_json:
                console.print(f"\n[dim]Array profiling skipped — {e}[/dim]")

    # --- Phase 2: Consistency metrics (Sakr et al.) ---
    consistency_report = None
    scaling_advice = None
    if action_episodes and len(action_episodes) >= 2:
        try:
            from orbit.analyzer.consistency_metrics import compute_consistency_metrics

            consistency_report = compute_consistency_metrics(
                action_episodes,
                state_episodes=state_episodes if state_episodes else None,
            )
        except Exception as e:
            if not as_json:
                console.print(f"\n[dim]Consistency metrics skipped — {e}[/dim]")

        try:
            from orbit.analyzer.scaling_advisor import compute_scaling_advice

            scaling_advice = compute_scaling_advice(
                episode_count=info.num_episodes,
                action_episodes=action_episodes,
                state_episodes=state_episodes if state_episodes else None,
            )
        except Exception as e:
            if not as_json:
                console.print(f"\n[dim]Scaling advice skipped — {e}[/dim]")

    # Compute quality and recommendations
    quality = compute_quality_score(
        info,
        coverage,
        visual_diversity=visual_diversity,
        action_episodes=action_episodes if action_episodes else None,
        state_episodes=state_episodes if state_episodes else None,
        signal_diagnostics_summary=(
            max(0.0, 1.0 - len(signal_diag.blockers) * 0.15
                - min(len(signal_diag.warnings), 6) * 0.05)
            if signal_diag
            else None
        ),
        consistency=consistency_report,
    )
    assessment = ready_to_train_assessment(quality)
    recs = generate_recommendations(info, coverage, quality)

    # --- NEW: Full prediction pipeline ---

    # Task Context Analysis
    task_context = None
    if action_episodes:
        try:
            analyzer = TaskContextAnalyzer()
            task_context = analyzer.analyze(
                action_episodes,
                metadata={
                    "robot_type": info.robot_type,
                    "task_description": dataset_repo,
                },
            )
        except Exception as e:
            if not as_json:
                console.print(f"\n[dim]Task context analysis skipped — {e}[/dim]")

    # Action Divergence — computed early so policy fit can use it.
    # Measures same-state → different-action conflicts, which is the
    # correct multimodality signal for policy compatibility (as opposed
    # to trajectory_multimodality which measures trajectory diversity,
    # which is actually beneficial for training).
    action_divergence_result = None
    if action_episodes:
        try:
            from orbit.analyzer.action_divergence import compute_action_divergence

            action_divergence_result = compute_action_divergence(
                action_episodes,
                state_episodes=state_episodes if state_episodes else None,
            )
        except Exception as e:
            if not as_json:
                console.print(f"\n[dim]Action divergence analysis skipped — {e}[/dim]")

    # Policy Fit Analysis
    policy_fit_result = None
    if action_episodes and task_context:
        try:
            pf_analyzer = PolicyFitAnalyzer()
            # Use action_divergence (same-state conflict) for multimodality,
            # NOT trajectory_multimodality (trajectory diversity).
            ad_multimodality = (
                action_divergence_result.score
                if action_divergence_result and action_divergence_result.score is not None
                else task_context.trajectory_multimodality
            )
            policy_fit_result = pf_analyzer.analyze(
                action_episodes,
                policy_type=policy,
                episode_count=info.num_episodes,
                metadata={"is_bimanual": task_context.is_bimanual},
                task_complexity=task_context.task_complexity_score,
                consistency_score=coverage.episode_consistency,
                multimodality_score=ad_multimodality,
            )
        except Exception as e:
            if not as_json:
                console.print(f"\n[dim]Policy fit analysis skipped — {e}[/dim]")

    # Community Comparison
    community_comp = None
    try:
        from orbit.analyzer.community_comparison import compute_community_comparison
        from orbit.analyzer.task_inference import infer_task_profile as _infer_tp

        _tp = _infer_tp(dataset_name=dataset_repo)
        community_comp = compute_community_comparison(
            dataset_name=dataset_repo,
            task_type=_tp.task_type,
            robot_type=info.robot_type,
            num_episodes=info.num_episodes,
            coverage_score=quality.overall_score,
            policy_type=policy_fit_result.policy_type if policy_fit_result else None,
        )
    except Exception as e:
        if not as_json:
            console.print(f"\n[dim]Community comparison skipped — {e}[/dim]")

    # --- NEW: Advanced analysis modules ---
    episode_ranking_result = None
    temporal_alignment_result = None

    if action_episodes:
        # Episode Ranking
        try:
            from orbit.analyzer.episode_ranker import compute_episode_ranking

            episode_ranking_result = compute_episode_ranking(
                action_episodes,
                state_episodes=state_episodes if state_episodes else None,
            )
        except Exception as e:
            if not as_json:
                console.print(f"\n[dim]Episode ranking skipped — {e}[/dim]")

        # Temporal Alignment
        try:
            from orbit.analyzer.temporal_alignment import compute_temporal_alignment

            temporal_alignment_result = compute_temporal_alignment(
                action_episodes,
                state_episodes=state_episodes if state_episodes else None,
            )
        except Exception as e:
            if not as_json:
                console.print(f"\n[dim]Temporal alignment check skipped — {e}[/dim]")

    # AI Assessment (VLM predictor — optional, only with --vlm flag)
    vlm_prediction = None
    if vlm:
        try:
            from orbit.analyzer.vlm_predictor import predict_success_rate

            vlm_quality_metrics = {}
            if not skip_signal_diagnostics and signal_diag is not None:
                vlm_quality_metrics["action_smoothness"] = (
                    float(np.mean([jd.smoothness for jd in signal_diag.joint_diagnostics]))
                    if signal_diag.joint_diagnostics
                    else 0.5
                )

            vlm_dataset_info = {
                "num_episodes": info.num_episodes,
                "task_description": dataset_repo,
                "robot_type": info.robot_type or "unknown",
                "policy": policy_fit_result.policy_type if policy_fit_result else "unknown",
            }

            vlm_prediction = predict_success_rate(
                coverage_report=coverage,
                signal_diagnostics=signal_diag,
                quality_metrics=vlm_quality_metrics,
                dataset_info=vlm_dataset_info,
            )
        except Exception as e:
            if not as_json:
                console.print(f"\n[dim]AI assessment skipped — {e}[/dim]")

    # --- Outlier detection (adaptive, distribution-free) ---
    outlier_report = None
    if action_episodes:
        try:
            from orbit.analyzer.outlier_detection import detect_outliers

            episode_data_for_outliers = {}
            for i, ep_actions in enumerate(action_episodes):
                ep_actions_arr = np.asarray(ep_actions, dtype=np.float32)
                if ep_actions_arr.ndim == 1:
                    ep_actions_arr = ep_actions_arr.reshape(-1, 1)
                ep_entry: dict = {"actions": ep_actions_arr, "length": len(ep_actions_arr)}
                if state_episodes and i < len(state_episodes):
                    ep_entry["states"] = np.asarray(state_episodes[i], dtype=np.float32)
                else:
                    ep_entry["states"] = None
                episode_data_for_outliers[i] = ep_entry
            outlier_report = detect_outliers(episode_data_for_outliers)
            if not as_json and outlier_report.outlier_count > 0:
                console.print(
                    f"\n[dim]Outlier detection: {outlier_report.outlier_count}/"
                    f"{outlier_report.total_episodes} episodes flagged "
                    f"({outlier_report.dataset_health})[/dim]"
                )
        except Exception as e:
            if not as_json:
                console.print(f"\n[dim]Outlier detection skipped — {e}[/dim]")

    # Proxy training — auto-run for small datasets, or on explicit --proxy
    # Two-pass architecture:
    #   1. MLP baseline (always) — comparable loss_ratio for meta-learner feature #21
    #   2. Policy-matched (optional) — actionable GO/NO-GO for user display
    proxy_result = None
    proxy_mlp_result = None
    should_proxy = (proxy or (info.num_episodes <= 200 and not skip_proxy)) and action_episodes
    if should_proxy:
        try:
            from orbit.analyzer.proxy_trainer import run_proxy_training

            resolved_policy = (
                policy_fit_result.policy_type if policy_fit_result else policy
            )
            arch_labels = {
                "mlp": "BC-MLP",
                "act": "ACT-mini",
                "diffusion": "DP-mini",
                "rnn": "LSTM",
            }
            from orbit.analyzer.proxy_architectures import get_proxy_architecture
            policy_arch = get_proxy_architecture(resolved_policy or "bc")

            # Pass 1: MLP baseline (always, for meta-learner feature)
            if not as_json:
                console.print(
                    "\n[bold]Running proxy training (BC-MLP baseline)...[/bold] (~2 min)"
                )
            proxy_mlp_result = run_proxy_training(
                action_episodes,
                state_episodes=state_episodes if state_episodes else None,
                policy_type="bc",  # always MLP for comparable baseline
            )
            # Compute loss ratio: val_loss_final / val_loss_initial
            mlp_loss_ratio = 1.0
            if proxy_mlp_result.validation_loss_curve and proxy_mlp_result.validation_loss_curve[0] > 1e-8:
                mlp_loss_ratio = proxy_mlp_result.validation_loss_final / proxy_mlp_result.validation_loss_curve[0]
            mlp_loss_ratio = min(mlp_loss_ratio, 1.0)  # cap at 1.0 (no learning)

            # Pass 2: Policy-matched architecture (for user-facing GO/NO-GO)
            if policy_arch != "mlp":
                arch_label = arch_labels.get(policy_arch, policy_arch.upper())
                est_time = "~5-10 min" if policy_arch in ("act", "diffusion") else "~2-3 min"
                if not as_json:
                    console.print(
                        f"  [bold]Running policy-matched proxy ({arch_label})...[/bold] ({est_time})"
                    )
                proxy_result = run_proxy_training(
                    action_episodes,
                    state_episodes=state_episodes if state_episodes else None,
                    policy_type=resolved_policy,
                )
            else:
                # MLP is both baseline and policy-matched
                proxy_result = proxy_mlp_result

            if not as_json:
                signal = {
                    "GO": "[bold green]GO[/bold green]",
                    "CAUTION": "[bold yellow]CAUTION[/bold yellow]",
                    "NO-GO": "[bold red]NO-GO[/bold red]",
                }.get(proxy_result.go_no_go, proxy_result.go_no_go)
                arch_label = arch_labels.get(proxy_result.architecture, proxy_result.architecture.upper())
                console.print(f"  Proxy check: {signal} (val_loss={proxy_result.validation_loss_final:.4f})")
                console.print(f"  {proxy_result.reason}")
                console.print(f"  MLP baseline loss ratio: {mlp_loss_ratio:.3f} (lower = more learnable)")
                if proxy_result.worst_episodes:
                    worst_str = ", ".join(f"#{e}" for e in proxy_result.worst_episodes[:3])
                    console.print(
                        f"  Worst episodes: {worst_str}"
                    )
                console.print(
                    f"  ({proxy_result.training_time_seconds:.0f}s, "
                    f"{proxy_result.steps_run} steps, {arch_label})"
                )
        except ImportError:
            if not as_json:
                console.print("\n[dim]Proxy training skipped — torch not installed[/dim]")
        except Exception as e:
            if not as_json:
                console.print(f"\n[dim]Proxy training failed — {e}[/dim]")

    # Dataset Readiness Grade (replaces success prediction)
    readiness_result = None
    if policy_fit_result:
        try:
            grader = DatasetReadinessGrader()
            sig_health = {
                "dead_dimensions": signal_diag.dead_dimensions if signal_diag else [],
                "blockers": signal_diag.blockers if signal_diag else [],
                "warnings": signal_diag.warnings if signal_diag else [],
                "total_joints": (
                    len(signal_diag.joint_diagnostics) if signal_diag else
                    (task_context.action_dim if task_context else 6)
                ),
                "clipping_joints": (
                    [jd.joint_index for jd in signal_diag.joint_diagnostics if jd.is_clipping]
                    if signal_diag else []
                ),
            }
            er_dict = {
                "total": len(episode_ranking_result.episodes) if episode_ranking_result else info.num_episodes,
                "remove_count": episode_ranking_result.remove_count if episode_ranking_result else 0,
                "review_count": episode_ranking_result.review_count if episode_ranking_result else 0,
            }
            ta_dict = None
            if temporal_alignment_result and not temporal_alignment_result.aligned:
                # Weight temporal penalty by action dimensionality.
                # Cross-correlation lag is unreliable on low-dim data (2D
                # simulation).  High-dim real robots (6D+) get full weight.
                action_dim = task_context.action_dim if task_context else 6
                dim_weight = min(1.0, action_dim / 6.0)

                raw_pct = 0.0
                if temporal_alignment_result.missing_frames > 0:
                    raw_pct = temporal_alignment_result.missing_frames / max(1, info.num_frames)
                if abs(temporal_alignment_result.estimated_lag) >= 2:
                    raw_pct = max(raw_pct, 0.15)

                weighted_pct = raw_pct * dim_weight
                if weighted_pct > 0.02:
                    ta_dict = {"misaligned_pct": min(0.6, weighted_pct)}

            ad_score = 0.0
            if action_divergence_result and action_divergence_result.score is not None:
                ad_score = action_divergence_result.score

            readiness_result = grader.grade(
                quality_score=quality.overall_score * 100,
                episode_consistency=coverage.episode_consistency,
                signal_health=sig_health,
                policy_fit_score=policy_fit_result.fit_score,
                episode_ranking=er_dict,
                action_divergence=ad_score,
                episode_count=info.num_episodes,
                policy_type=policy_fit_result.policy_type,
                temporal_alignment=ta_dict,
                consistency=consistency_report,
                outlier_report=outlier_report,
                proxy_result=proxy_result,
            )
        except Exception as e:
            if not as_json:
                console.print(f"\n[dim]Readiness grading skipped — {e}[/dim]")

    # --- AI quality judge (opt-in, only for A grades) ---
    ai_judge_result = None
    if ai and readiness_result and readiness_result.grade == "A":
        try:
            from orbit.analyzer.ai_quality_judge import check_grade

            n_clip = len(sig_health.get("clipping_joints", []))
            elevated_cvs = 0
            if consistency_report:
                for cv in [
                    consistency_report.cartesian_jerk_cv,
                    consistency_report.path_length_cv,
                    consistency_report.joint_path_length_cv,
                    consistency_report.joint_jerk_cv,
                    consistency_report.curvature_cv,
                    consistency_report.effort_cv,
                    consistency_report.joint_limit_cv,
                    consistency_report.orientation_cv,
                ]:
                    if cv > 0.50:
                        elevated_cvs += 1

            ai_judge_result = check_grade(
                divergence=ad_score,
                n_clipping=n_clip,
                policy_fit=policy_fit_result.fit_score,
                consistency=(
                    consistency_report.overall_consistency
                    if consistency_report else 0.5
                ),
                elevated_cvs=elevated_cvs,
            )

            if not ai_judge_result.error:
                if not as_json:
                    verdict = (
                        "[green]Keep A[/green]"
                        if ai_judge_result.keep_grade
                        else "[yellow]→ B[/yellow]"
                    )
                    console.print(f"\n  AI judge: {verdict}")
                    console.print(f"  [dim]{ai_judge_result.reasoning}[/dim]")

                if not ai_judge_result.keep_grade:
                    readiness_result = type(readiness_result)(
                        grade="B",
                        score=min(readiness_result.score, 87),
                        summary="Good data — minor issues, should train well",
                        strengths=readiness_result.strengths,
                        problems=readiness_result.problems + [
                            f"AI judge: {ai_judge_result.reasoning}"
                        ],
                        top_action=readiness_result.top_action,
                        episode_count=readiness_result.episode_count,
                        policy_type=readiness_result.policy_type,
                    )
        except Exception as e:
            if not as_json:
                console.print(f"\n[dim]AI judge skipped — {e}[/dim]")

    # --- Gemini auto-diagnosis for grade C or below ---
    gemini_diagnosis = None
    if (
        readiness_result
        and readiness_result.grade in ("C", "D", "F")
        and not skip_embeddings  # don't auto-invoke expensive APIs in quick mode
    ):
        try:
            from orbit.analyzer.gemini_analyst import GeminiDataAnalyst

            analyst = GeminiDataAnalyst()
            if analyst.available:
                if not as_json:
                    console.print("\n[bold]Running AI diagnosis (Gemini)...[/bold]")
                # Build analysis summary dict for the analyst
                _analysis_for_gemini = {
                    "repo_id": info.repo_id,
                    "num_episodes": info.num_episodes,
                    "num_frames": info.num_frames,
                    "quality_score": quality.overall_score,
                    "grade": readiness_result.grade,
                    "problems": readiness_result.problems,
                }
                if action_divergence_result and action_divergence_result.score is not None:
                    _analysis_for_gemini["action_divergence"] = {
                        "score": action_divergence_result.score,
                        "severity": action_divergence_result.severity,
                    }
                if consistency_report:
                    _analysis_for_gemini["consistency_diagnostics"] = {
                        "overall_consistency": consistency_report.overall_consistency,
                        "path_length_cv": consistency_report.path_length_cv,
                        "cartesian_jerk_cv": consistency_report.cartesian_jerk_cv,
                        "effort_cv": consistency_report.effort_cv,
                    }
                if episode_ranking_result:
                    _analysis_for_gemini["episode_ranking"] = {
                        "remove_count": episode_ranking_result.remove_count,
                        "review_count": episode_ranking_result.review_count,
                        "remove_indices": episode_ranking_result.remove_indices[:10],
                    }
                if signal_diag:
                    _analysis_for_gemini["signal_diagnostics"] = {
                        "dead_dimensions": signal_diag.dead_dimensions,
                        "n_joints": len(signal_diag.joint_diagnostics),
                    }
                gemini_diagnosis = analyst.diagnose_problems(
                    _analysis_for_gemini,
                    array_profiles=array_profiles if array_profiles else None,
                )
                if gemini_diagnosis and not as_json:
                    console.print(f"  [bold]{gemini_diagnosis.summary}[/bold]")
        except Exception as e:
            if not as_json:
                console.print(f"\n[dim]Gemini diagnosis skipped — {e}[/dim]")

    # --- Meta-learner success prediction ---
    meta_prediction = None
    try:
        from orbit.analyzer.meta_learner import (
            MODEL_PATH,
            MetaFeatures,
            extract_features,
            predict_success,
        )

        if MODEL_PATH.exists() and readiness_result is not None:
            # Build a partial analysis dict from available data
            partial_analysis = {
                "repo_id": dataset_repo,
                "num_episodes": info.num_episodes,
                "num_frames": info.num_frames,
                "camera_names": info.camera_names,
                "quality_score": quality.overall_score,
                "coverage": {
                    "episode_consistency": coverage.episode_consistency,
                },
            }
            if task_context:
                partial_analysis["task_context"] = {
                    "action_dim": task_context.action_dim,
                    "is_bimanual": task_context.is_bimanual,
                }
            if policy_fit_result:
                partial_analysis["policy_fit"] = {
                    "fit_score": policy_fit_result.fit_score,
                }
            if signal_diag:
                partial_analysis["signal_diagnostics"] = {
                    "dead_dimensions": signal_diag.dead_dimensions,
                    "joint_diagnostics": [
                        {"is_clipping": jd.is_clipping}
                        for jd in signal_diag.joint_diagnostics
                    ],
                }
            if action_divergence_result and action_divergence_result.score is not None:
                partial_analysis["action_divergence"] = {
                    "score": action_divergence_result.score,
                }
            if episode_ranking_result and episode_ranking_result.episodes:
                partial_analysis["episode_ranking"] = {
                    "keep_count": episode_ranking_result.keep_count,
                    "review_count": episode_ranking_result.review_count,
                    "remove_count": episode_ranking_result.remove_count,
                }
            if consistency_report:
                partial_analysis["consistency_diagnostics"] = {
                    "overall_consistency": consistency_report.overall_consistency,
                    "metrics": {
                        "path_length_cv": consistency_report.path_length_cv,
                        "cartesian_jerk_cv": consistency_report.cartesian_jerk_cv,
                        "effort_cv": consistency_report.effort_cv,
                        "joint_path_length_cv": consistency_report.joint_path_length_cv,
                        "joint_jerk_cv": consistency_report.joint_jerk_cv,
                        "curvature_cv": consistency_report.curvature_cv,
                        "joint_limit_cv": consistency_report.joint_limit_cv,
                        "orientation_cv": consistency_report.orientation_cv,
                    },
                }

            features = extract_features(partial_analysis)
            policy_for_pred = (
                policy_fit_result.policy_type if policy_fit_result else "diffusion_policy"
            )
            meta_prediction = predict_success(features, policy_for_pred)

            # Display is handled by _print_compact_report or _print_report
            # Only show here in --detail mode (full report doesn't include it)
            if not as_json and detail and meta_prediction.model_available:
                pct = int(meta_prediction.predicted_success_rate * 100)
                lo = int(meta_prediction.confidence_low * 100)
                hi = int(meta_prediction.confidence_high * 100)
                conf_color = {
                    "HIGH": "green", "MEDIUM": "yellow", "LOW": "dim"
                }.get(meta_prediction.confidence_level, "dim")
                console.print(
                    f"\n  Predicted success: [bold]{pct}%[/bold] "
                    f"([{conf_color}]{lo}%-{hi}%[/{conf_color}]) "
                    f"[dim][{meta_prediction.confidence_level}][/dim]"
                )
                if meta_prediction.closest_match:
                    console.print(
                        f"  [dim]Closest match: {meta_prediction.closest_match}[/dim]"
                    )
    except Exception as e:
        if not as_json:
            import logging as _logging
            _logging.getLogger(__name__).debug("Meta-learner prediction skipped: %s", e)

    if as_json:
        report = _full_report_to_dict(
            info,
            coverage,
            quality,
            assessment,
            recs,
            embedding_analysis=embedding_analysis,
            signal_diagnostics=signal_diag,
        )
        report["coverage_estimated"] = is_estimated
        if meta_prediction is not None and meta_prediction.model_available:
            report["success_prediction"] = {
                "predicted_success_rate": meta_prediction.predicted_success_rate,
                "confidence_low": meta_prediction.confidence_low,
                "confidence_high": meta_prediction.confidence_high,
                "confidence_width": meta_prediction.confidence_width,
                "confidence_level": meta_prediction.confidence_level,
                "n_training_samples": meta_prediction.n_training_samples,
                "closest_match": meta_prediction.closest_match,
            }
        if readiness_result is not None:
            report["readiness"] = {
                "grade": readiness_result.grade,
                "score": readiness_result.score,
                "summary": readiness_result.summary,
                "strengths": readiness_result.strengths,
                "problems": readiness_result.problems,
                "top_action": readiness_result.top_action,
            }
        if gemini_diagnosis is not None:
            report["ai_diagnosis"] = {
                "summary": gemini_diagnosis.summary,
                "issues": gemini_diagnosis.issues,
                "overall_assessment": gemini_diagnosis.overall_assessment,
            }
        if task_context is not None:
            report["task_context"] = {
                "task_complexity_score": task_context.task_complexity_score,
                "dimension_complexity": task_context.dimension_complexity,
                "temporal_complexity": task_context.temporal_complexity,
                "coordination_complexity": task_context.coordination_complexity,
                "object_interaction_complexity": task_context.object_interaction_complexity,
                "trajectory_multimodality": task_context.trajectory_multimodality,
                "is_bimanual": task_context.is_bimanual,
                "action_dim": task_context.action_dim,
                "estimated_task_type": task_context.estimated_task_type,
            }
        if policy_fit_result is not None:
            report["policy_fit"] = {
                "policy_type": policy_fit_result.policy_type,
                "fit_score": policy_fit_result.fit_score,
                "episode_sufficiency": policy_fit_result.episode_sufficiency,
                "action_space_compatibility": policy_fit_result.action_space_compatibility,
                "consistency_compatibility": policy_fit_result.consistency_compatibility,
                "specific_warnings": policy_fit_result.specific_warnings,
                "recommended_policies": policy_fit_result.recommended_policies,
            }
        if community_comp is not None:
            from orbit.analyzer.community_comparison import comparison_to_dict

            report["community_comparison"] = comparison_to_dict(community_comp)
        if vlm_prediction is not None:
            report["ai_assessment"] = {
                "predicted_success_rate": vlm_prediction.predicted_success_rate,
                "confidence": vlm_prediction.confidence,
                "grade": vlm_prediction.grade,
                "reasoning": vlm_prediction.reasoning,
                "similar_datasets": vlm_prediction.similar_datasets,
                "key_strengths": vlm_prediction.key_strengths,
                "key_weaknesses": vlm_prediction.key_weaknesses,
                "improvement_estimate": vlm_prediction.improvement_estimate,
                "model_used": vlm_prediction.model_used,
                "tokens_used": vlm_prediction.tokens_used,
                "cost_estimate": vlm_prediction.cost_estimate,
            }
        if action_divergence_result is not None and action_divergence_result.score is not None:
            report["action_divergence"] = {
                "score": action_divergence_result.score,
                "severity": action_divergence_result.severity,
                "num_clusters": action_divergence_result.num_clusters,
                "estimated_success_penalty": action_divergence_result.estimated_success_penalty,
                "recommendation": action_divergence_result.recommendation,
            }
        if episode_ranking_result is not None and episode_ranking_result.episodes:
            report["episode_ranking"] = {
                "keep_count": episode_ranking_result.keep_count,
                "review_count": episode_ranking_result.review_count,
                "remove_count": episode_ranking_result.remove_count,
                "remove_indices": episode_ranking_result.remove_indices,
                "estimated_improvement": episode_ranking_result.estimated_improvement,
                "current_quality": episode_ranking_result.current_quality_estimate,
                "improved_quality": episode_ranking_result.improved_quality_estimate,
                "summary": episode_ranking_result.summary,
            }
        if temporal_alignment_result is not None:
            report["temporal_alignment"] = {
                "aligned": temporal_alignment_result.aligned,
                "estimated_lag": temporal_alignment_result.estimated_lag,
                "fps_jitter": temporal_alignment_result.fps_jitter,
                "missing_frames": temporal_alignment_result.missing_frames,
                "constant_dimensions": temporal_alignment_result.constant_dimensions,
                "severity": temporal_alignment_result.severity,
                "recommendation": temporal_alignment_result.recommendation,
            }
        if consistency_report is not None:
            # Identify worst episodes for actionable output
            worst_eps = []
            if consistency_report.per_episode_scores:
                ranked = sorted(
                    consistency_report.per_episode_scores.items(),
                    key=lambda x: x[1],
                )
                worst_eps = [idx for idx, _ in ranked[:5]]

            report["consistency_diagnostics"] = {
                "overall_consistency": consistency_report.overall_consistency,
                "metrics": {
                    "path_length_cv": consistency_report.path_length_cv,
                    "joint_path_length_cv": consistency_report.joint_path_length_cv,
                    "cartesian_jerk_cv": consistency_report.cartesian_jerk_cv,
                    "joint_jerk_cv": consistency_report.joint_jerk_cv,
                    "curvature_cv": consistency_report.curvature_cv,
                    "effort_cv": consistency_report.effort_cv,
                    "joint_limit_cv": consistency_report.joint_limit_cv,
                    "orientation_cv": consistency_report.orientation_cv,
                },
                "most_inconsistent_episodes": worst_eps,
            }
        if scaling_advice is not None:
            report["scaling_advice"] = {
                "estimated_diversity": scaling_advice.estimated_diversity,
                "demos_per_setting": scaling_advice.demos_per_setting,
                "bottleneck": scaling_advice.bottleneck,
                "expected_success_at_current": scaling_advice.expected_success_at_current,
                "expected_success_at_recommended": scaling_advice.expected_success_at_recommended,
                "advice": scaling_advice.advice,
                "details": scaling_advice.details,
            }
        if proxy_result is not None:
            # Policy-matched proxy (user-facing GO/NO-GO)
            report["proxy_training"] = {
                "go_no_go": proxy_result.go_no_go,
                "reason": proxy_result.reason,
                "architecture": proxy_result.architecture,
                "validation_loss_final": proxy_result.validation_loss_final,
                "training_loss_final": proxy_result.training_loss_final,
                "overfitting_ratio": proxy_result.overfitting_ratio,
                "convergence_step": proxy_result.convergence_step,
                "training_time_seconds": proxy_result.training_time_seconds,
                "steps_run": proxy_result.steps_run,
                "training_loss_curve": proxy_result.training_loss_curve,
                "validation_loss_curve": proxy_result.validation_loss_curve,
                "worst_episodes": proxy_result.worst_episodes,
                "per_episode_errors": proxy_result.per_episode_errors,
            }
        if proxy_mlp_result is not None:
            # MLP baseline (comparable across datasets, for meta-learner)
            _mlp_lr = 1.0
            if proxy_mlp_result.validation_loss_curve and proxy_mlp_result.validation_loss_curve[0] > 1e-8:
                _mlp_lr = proxy_mlp_result.validation_loss_final / proxy_mlp_result.validation_loss_curve[0]
            report["proxy_training_mlp_baseline"] = {
                "loss_ratio": round(min(_mlp_lr, 1.0), 4),
                "validation_loss_final": proxy_mlp_result.validation_loss_final,
                "validation_loss_initial": proxy_mlp_result.validation_loss_curve[0] if proxy_mlp_result.validation_loss_curve else None,
                "training_time_seconds": proxy_mlp_result.training_time_seconds,
                "steps_run": proxy_mlp_result.steps_run,
            }
        if array_profiles:
            report["array_intelligence"] = {}
            for name, ap in array_profiles.items():
                report["array_intelligence"][name] = {
                    "action_type": ap.action_type,
                    "action_type_confidence": round(ap.action_type_confidence, 3),
                    "is_normalized": ap.is_normalized,
                    "is_bimanual": ap.is_bimanual,
                    "smoothness": round(ap.smoothness, 3),
                    "num_dimensions": ap.num_dimensions,
                    "dead_dimensions": ap.dead_dimensions,
                    "gripper_dimensions": ap.gripper_dimensions,
                    "clipping_dimensions": ap.clipping_dimensions,
                    "noisy_dimensions": ap.noisy_dimensions,
                    "dimension_roles": [
                        {"role": r.role, "confidence": round(r.confidence, 2)}
                        for r in ap.dimension_roles
                    ],
                    "confidence": round(ap.confidence, 3),
                }
        click.echo(json.dumps(report, indent=2, default=_json_default))
        return

    _print_report(
        info,
        coverage,
        quality,
        assessment,
        recs,
        embedding_analysis=embedding_analysis,
        is_estimated=is_estimated,
        signal_diagnostics=signal_diag,
        vlm_prediction=vlm_prediction,
        readiness_result=readiness_result,
        task_context=task_context,
        policy_fit=policy_fit_result,
        community_comparison=community_comp,
        action_divergence=action_divergence_result,
        episode_ranking=episode_ranking_result,
        temporal_alignment=temporal_alignment_result,
        scaling_advice=scaling_advice,
        consistency_report=consistency_report,
        meta_prediction=meta_prediction,
        detail=detail,
        proxy_result=proxy_result,
        gemini_diagnosis=gemini_diagnosis,
    )

    # Deep AI analysis (optional, requires API key)
    if deep and not as_json:
        from rich.markdown import Markdown

        from orbit.analyzer.deep_analysis import run_deep_analysis

        # Build dicts from structured results for the LLM
        _cov_dict = None
        if coverage:
            _cov_dict = {
                "position_diversity": coverage.position_diversity,
                "action_diversity": coverage.action_diversity,
                "episode_consistency": coverage.episode_consistency,
                "temporal_coverage": coverage.temporal_coverage,
            }
        _sig_dict = None
        if signal_diag:
            _sig_dict = {
                "dead_dimensions": signal_diag.dead_dimensions,
                "blockers": signal_diag.blockers,
                "warnings": signal_diag.warnings,
                "n_joints": len(signal_diag.joint_diagnostics),
                "dead_count": len(signal_diag.dead_dimensions),
                "episodes_to_remove": signal_diag.episodes_to_remove,
                "episodes_to_review": signal_diag.episodes_to_review,
            }
        _rank_dict = None
        if episode_ranking_result and episode_ranking_result.episodes:
            _rank_dict = {
                "keep_count": episode_ranking_result.keep_count,
                "review_count": episode_ranking_result.review_count,
                "remove_count": episode_ranking_result.remove_count,
                "remove_indices": episode_ranking_result.remove_indices,
                "current_quality": episode_ranking_result.current_quality_estimate,
                "improved_quality": episode_ranking_result.improved_quality_estimate,
            }
        _div_dict = None
        if action_divergence_result and action_divergence_result.score is not None:
            _div_dict = {
                "score": action_divergence_result.score,
                "severity": action_divergence_result.severity,
                "estimated_success_penalty": action_divergence_result.estimated_success_penalty,
            }
        _temp_dict = None
        if temporal_alignment_result:
            _temp_dict = {
                "aligned": temporal_alignment_result.aligned,
                "estimated_lag": temporal_alignment_result.estimated_lag,
                "severity": temporal_alignment_result.severity,
                "fps_jitter": temporal_alignment_result.fps_jitter,
                "missing_frames": temporal_alignment_result.missing_frames,
            }
        _pf_dict = None
        if policy_fit_result:
            _pf_dict = {
                "policy_type": policy_fit_result.policy_type,
                "fit_score": policy_fit_result.fit_score,
                "episode_sufficiency": policy_fit_result.episode_sufficiency,
                "consistency_compatibility": policy_fit_result.consistency_compatibility,
                "specific_warnings": policy_fit_result.specific_warnings,
            }

        console.print()
        with console.status("[dim]Running AI analysis...[/dim]"):
            deep_result = run_deep_analysis(
                dataset_id=dataset_repo,
                episode_count=info.num_episodes,
                action_dim=task_context.action_dim if task_context else 0,
                fps=info.fps,
                robot_type=info.robot_type or "unknown",
                policy_type=policy_fit_result.policy_type if policy_fit_result else policy,
                is_simulation=getattr(info, "is_simulation", False),
                readiness_grade=readiness_result.grade if readiness_result else "?",
                readiness_score=readiness_result.score if readiness_result else 0,
                coverage_results=_cov_dict,
                signal_results=_sig_dict,
                episode_ranking=_rank_dict,
                divergence_results=_div_dict,
                temporal_results=_temp_dict,
                policy_fit_results=_pf_dict,
                quality_score=round(quality.overall_score * 100, 1),
                recommendations=recs,
                no_cache=no_cache,
            )

        if deep_result:
            console.print("[bold]AI DEEP ANALYSIS[/bold]")
            console.print("\u2500" * 40)
            console.print(Markdown(deep_result))
        else:
            console.print(
                "[dim]\u26a0 AI analysis unavailable — "
                "set GOOGLE_API_KEY environment variable for deep analysis[/dim]"
            )

    # Push results only when explicitly requested
    if push:
        _maybe_push_results(dataset_repo, info, coverage, quality, recs, embedding_analysis,
                            signal_diag, readiness_result, task_context, policy_fit_result,
                            community_comp, vlm_prediction, assessment)


def _maybe_push_results(dataset_repo, info, coverage, quality, recs, embedding_analysis,
                        signal_diag, readiness_result, task_context, policy_fit_result,
                        community_comp, vlm_prediction, assessment):
    """If the user is logged in, offer to push analysis results to the dashboard.

    When a previous analysis exists, fetches it and prints a before/after diff
    so the user can see what changed.
    """
    from orbit.utils.display import console

    creds = _load_credentials()
    if not creds:
        console.print(
            "[dim]Not logged in — run `orbit login --token <token>` to enable --push[/dim]"
        )
        return

    import httpx

    report = _full_report_to_dict(
        info, coverage, quality, assessment, recs,
        embedding_analysis=embedding_analysis,
        signal_diagnostics=signal_diag,
    )

    # --- Fetch previous analysis for before/after diff ---
    previous = None
    try:
        prev_resp = httpx.get(
            f"{creds['api_url']}/analysis/latest",
            headers=_api_headers(creds["token"]),
            params={"dataset_repo": dataset_repo},
            timeout=15,
        )
        if prev_resp.status_code == 200:
            previous = prev_resp.json()
    except httpx.RequestError as e:
        import logging as _log
        _log.getLogger(__name__).debug("No previous analysis available: %s", e)

    payload = {"dataset_repo": dataset_repo}
    for key in (
        "quality_score", "estimated_success_rate", "coverage_results",
        "recommendations", "embedding_clusters", "embedding_pca",
        "phase_detection", "signal_diagnostics", "task_context",
        "policy_fit", "readiness", "community_comparison",
        "ai_assessment", "llm_assessment",
    ):
        if key in report:
            payload[key] = report[key]

    try:
        resp = httpx.post(
            f"{creds['api_url']}/analysis/upload",
            headers=_api_headers(creds["token"]),
            json=payload,
            timeout=30,
        )
        if resp.status_code == 201:
            result = resp.json()
            console.print(f"[bold green]\u2713[/bold green] Results synced. View at: [bold]{result['url']}[/bold]")

            # --- Print before/after diff if previous analysis exists ---
            if previous:
                console.print()
                console.print("[bold]CHANGES SINCE LAST ANALYSIS[/bold]")
                console.print("\u2500" * 40)

                diff_items = []
                # Quality score diff
                prev_quality = previous.get("quality_score")
                curr_quality = report.get("quality_score")
                if prev_quality is not None and curr_quality is not None:
                    delta = curr_quality - prev_quality
                    arrow = "\u2191" if delta > 0 else ("\u2193" if delta < 0 else "\u2192")
                    color = "green" if delta > 0 else ("red" if delta < 0 else "dim")
                    diff_items.append(
                        f"  Quality Score:  {prev_quality:.2f} [{color}]{arrow} {curr_quality:.2f}[/{color}]"
                        f"  ([{color}]{delta:+.2f}[/{color}])"
                    )

                # Episode count diff
                prev_eps = previous.get("coverage_results", {}).get("num_episodes")
                curr_eps = info.num_episodes
                if prev_eps is not None:
                    ep_delta = curr_eps - prev_eps
                    if ep_delta != 0:
                        arrow = "\u2191" if ep_delta > 0 else "\u2193"
                        color = "green" if ep_delta > 0 else "yellow"
                        diff_items.append(
                            f"  Episodes:       {prev_eps} [{color}]{arrow} {curr_eps}[/{color}]"
                            f"  ([{color}]{ep_delta:+d}[/{color}])"
                        )

                # Coverage diff
                prev_cov = previous.get("coverage_results", {})
                cov_keys = [
                    ("position_diversity", "Position Div"),
                    ("action_diversity", "Action Div"),
                    ("episode_consistency", "Consistency"),
                    ("temporal_coverage", "Temporal Cov"),
                ]
                for cov_key, cov_label in cov_keys:
                    prev_val = prev_cov.get(cov_key)
                    curr_val = getattr(coverage, cov_key, None)
                    if prev_val is not None and curr_val is not None:
                        delta = curr_val - prev_val
                        if abs(delta) >= 0.01:
                            arrow = "\u2191" if delta > 0 else "\u2193"
                            color = "green" if delta > 0 else "red"
                            diff_items.append(
                                f"  {cov_label + ':':<16}{prev_val:.2f} [{color}]{arrow} {curr_val:.2f}[/{color}]"
                                f"  ([{color}]{delta:+.2f}[/{color}])"
                            )

                # Readiness grade diff
                prev_grade = previous.get("readiness", {}).get("grade")
                curr_grade = readiness_result.grade if readiness_result else None
                prev_score = previous.get("readiness", {}).get("score")
                curr_score = readiness_result.score if readiness_result else None
                if prev_score is not None and curr_score is not None:
                    delta = curr_score - prev_score
                    arrow = "\u2191" if delta > 0 else ("\u2193" if delta < 0 else "\u2192")
                    color = "green" if delta > 0 else ("red" if delta < 0 else "dim")
                    diff_items.append(
                        f"  Readiness:      {prev_grade} ({prev_score}) [{color}]{arrow} {curr_grade} ({curr_score})[/{color}]"
                        f"  ([{color}]{delta:+d}[/{color}])"
                    )

                if diff_items:
                    for item in diff_items:
                        console.print(item)
                else:
                    console.print("  [dim]No significant changes detected.[/dim]")
                console.print()
        else:
            console.print(f"[dim]Push failed ({resp.status_code})[/dim]")
    except httpx.RequestError:
        console.print("[dim]Push failed \u2014 could not reach API[/dim]")


def _synthetic_coverage(info):
    """Generate synthetic coverage scores based on dataset metadata."""
    from orbit.analyzer.coverage import FullCoverageReport

    ep_factor = min(1.0, info.num_episodes / 50)
    return FullCoverageReport(
        position_diversity=min(1.0, 0.3 + ep_factor * 0.5),
        action_diversity=min(1.0, 0.2 + ep_factor * 0.5),
        episode_consistency=min(1.0, 0.5 + ep_factor * 0.4),
        temporal_coverage=min(1.0, 0.2 + ep_factor * 0.4),
    )


def _coverage_status(score: float) -> tuple[str, str]:
    """Return (status, label) for a coverage score."""
    if score >= 0.65:
        return "ok", "good"
    elif score >= 0.45:
        return "warn", "moderate"
    else:
        return "bad", "needs improvement"


def _score_color(score: float) -> str:
    if score >= 0.65:
        return "green"
    elif score >= 0.45:
        return "yellow"
    else:
        return "red"



def _priority_style(priority: str) -> str:
    if priority == "HIGH":
        return "bold red"
    elif priority == "MEDIUM":
        return "bold yellow"
    else:
        return "dim"


def _info_to_dict(info) -> dict:
    """Convert DatasetInfo to a JSON-serializable dict."""
    from orbit.utils.display import format_duration

    return {
        "repo_id": info.repo_id,
        "num_episodes": info.num_episodes,
        "num_frames": info.num_frames,
        "fps": info.fps,
        "robot_type": info.robot_type,
        "camera_names": info.camera_names,
        "video_keys": info.video_keys,
        "shapes": info.shapes,
        "duration_seconds": info.duration_seconds,
        "duration_human": format_duration(info.duration_seconds),
    }


def _full_report_to_dict(
    info,
    coverage,
    quality,
    assessment,
    recs,
    embedding_analysis=None,
    signal_diagnostics=None,
) -> dict:
    """Convert the full analysis report to a JSON-serializable dict."""
    data = _info_to_dict(info)
    data["coverage"] = {
        "position_diversity": coverage.position_diversity,
        "action_diversity": coverage.action_diversity,
        "episode_consistency": coverage.episode_consistency,
        "temporal_coverage": coverage.temporal_coverage,
    }
    data["quality_score"] = quality.overall_score
    data["component_scores"] = quality.component_scores
    data["ready_to_train"] = quality.ready_to_train
    data["ready_to_train_verdict"] = assessment.verdict
    data["ready_to_train_reason"] = assessment.reason
    data["confidence"] = quality.confidence
    low, high = quality.estimated_success_rate
    data["estimated_success_rate"] = {
        "low": round(low, 2),
        "high": round(high, 2),
        "display": f"{low:.0%}-{high:.0%}",
    }
    data["recommendations"] = [
        {
            "priority": r.priority,
            "category": r.category,
            "description": r.description,
            "expected_impact": r.expected_impact,
            "effort": r.effort,
        }
        for r in recs
    ]
    data["blocking_issues"] = assessment.blocking_issues
    data["nice_to_haves"] = assessment.nice_to_haves

    if embedding_analysis is not None:
        data["embedding_analysis"] = {
            "n_clusters": embedding_analysis.n_clusters,
            "cluster_sizes": embedding_analysis.cluster_sizes,
            "visual_diversity": embedding_analysis.diversity_score,
            "outlier_episodes": embedding_analysis.outlier_indices,
            "insights": embedding_analysis.insights,
        }
        data["visual_diversity"] = embedding_analysis.diversity_score

    if signal_diagnostics is not None:
        data["signal_diagnostics"] = _signal_diag_to_dict(signal_diagnostics)

    return data


def _signal_diag_to_dict(sd) -> dict:
    """Convert a SignalDiagnosticsReport to a JSON-serializable dict."""
    return {
        "joint_diagnostics": [
            {
                "joint_index": jd.joint_index,
                "joint_name": jd.joint_name,
                "range_used": list(jd.range_used),
                "range_limit": list(jd.range_limit) if jd.range_limit else None,
                "range_utilization": round(jd.range_utilization, 4),
                "zero_velocity_ratio": round(jd.zero_velocity_ratio, 4),
                "smoothness": round(jd.smoothness, 4),
                "directional_bias": round(jd.directional_bias, 4),
                "is_dead": jd.is_dead,
                "is_clipping": jd.is_clipping,
                "issues": jd.issues,
            }
            for jd in sd.joint_diagnostics
        ],
        "gripper_analysis": {
            "gripper_dim": sd.gripper_analysis.gripper_dim,
            "open_ratio": round(sd.gripper_analysis.open_ratio, 4),
            "close_ratio": round(sd.gripper_analysis.close_ratio, 4),
            "transitions_per_episode": round(sd.gripper_analysis.transitions_per_episode, 2),
            "has_sufficient_grasps": sd.gripper_analysis.has_sufficient_grasps,
            "issues": sd.gripper_analysis.issues,
        },
        "episode_diagnostics": [
            {
                "episode_index": ed.episode_index,
                "quality_score": round(ed.quality_score, 4),
                "length": ed.length,
                "is_truncated": ed.is_truncated,
                "jerk_spike_count": ed.jerk_spike_count,
                "recommendation": ed.recommendation,
                "issues": ed.issues,
            }
            for ed in sd.episode_diagnostics
        ],
        "global_zero_velocity_ratio": round(sd.global_zero_velocity_ratio, 4),
        "dead_dimensions": sd.dead_dimensions,
        "episodes_to_remove": sd.episodes_to_remove,
        "episodes_to_review": sd.episodes_to_review,
        "blockers": sd.blockers,
        "warnings": sd.warnings,
        "passing": sd.passing,
        "prescription": sd.prescription,
    }


def _print_compact_report(
    info,
    readiness_result,
    meta_prediction,
    policy_fit,
    episode_ranking,
    action_divergence,
    signal_diagnostics,
    proxy_result=None,
    gemini_diagnosis=None,
) -> None:
    """Print a concise summary — grade, prediction, issues, next step.

    This is the default output. Users run --detail for the full report.
    """
    from rich.panel import Panel

    from orbit.utils.display import console

    policy_name = policy_fit.policy_type.upper() if policy_fit else "AUTO"

    # Header
    console.print()
    console.print(Panel(
        f"[bold]ORBIT Analysis: {info.repo_id}[/bold]\n"
        f"{info.num_episodes} episodes \u00b7 {policy_name}",
        border_style="cyan",
    ))

    # Grade
    if readiness_result:
        grade = readiness_result.grade
        score = readiness_result.score
        color = {"A": "green", "B": "green", "C": "yellow", "D": "yellow", "F": "red"}.get(grade, "white")
        console.print(f"  Grade: [{color}][bold]{grade}[/bold][/{color}] ({score}/100) \u2014 {readiness_result.summary}")

    # Meta-learner prediction
    if meta_prediction and meta_prediction.model_available:
        pct = int(meta_prediction.predicted_success_rate * 100)
        lo = int(meta_prediction.confidence_low * 100)
        hi = int(meta_prediction.confidence_high * 100)
        conf_color = {"HIGH": "green", "MEDIUM": "yellow", "LOW": "dim"}.get(
            meta_prediction.confidence_level, "dim"
        )
        caveat = ""
        if meta_prediction.confidence_level == "LOW":
            caveat = " [dim](wide interval — treat as directional)[/dim]"
        console.print(
            f"  Predicted success: [bold]{pct}%[/bold] "
            f"([{conf_color}]{lo}%-{hi}%[/{conf_color}]){caveat}"
        )

    # Issues (max 3)
    if readiness_result and readiness_result.problems:
        n_blocking = sum(1 for p in readiness_result.problems if "dead" in p.lower() or "critical" in p.lower())
        n_total = len(readiness_result.problems)
        console.print(f"\n  {n_total} issue{'s' if n_total != 1 else ''} found:")
        for problem in readiness_result.problems[:3]:
            console.print(f"    [yellow]\u26a0[/yellow] {problem}")
        if n_total > 3:
            console.print(f"    [dim]... and {n_total - 3} more (run with --detail to see all)[/dim]")

    # Proxy training result
    if proxy_result is not None:
        signal = {
            "GO": "[bold green]GO[/bold green]",
            "CAUTION": "[bold yellow]CAUTION[/bold yellow]",
            "NO-GO": "[bold red]NO-GO[/bold red]",
        }.get(proxy_result.go_no_go, proxy_result.go_no_go)
        worst_str = ""
        if proxy_result.worst_episodes:
            worst_str = ", worst episodes: " + ", ".join(f"#{e}" for e in proxy_result.worst_episodes[:3])
        console.print(
            f"\n  Proxy check: {signal} "
            f"(val_loss={proxy_result.validation_loss_final:.4f}{worst_str})"
        )

    # AI Diagnosis (Gemini)
    if gemini_diagnosis is not None:
        console.print(f"\n  [bold]AI Diagnosis:[/bold] {gemini_diagnosis.summary}")
        if gemini_diagnosis.issues:
            top_fix = gemini_diagnosis.issues[0].get("fix", "")
            if top_fix:
                console.print(f"    Fix: {top_fix}")

    # Next step
    if readiness_result:
        console.print(f"\n  [bold]Next:[/bold] {readiness_result.top_action}")

    # Footer with hint
    hints = []
    if proxy_result is None:
        hints.append("--proxy for training signal")
    hints.append("--detail for full diagnostics")
    console.print(f"\n  [dim]Run with {', '.join(hints)}[/dim]")
    console.print()


def _print_report(
    info,
    coverage,
    quality,
    assessment,
    recs,
    embedding_analysis=None,
    is_estimated=False,
    signal_diagnostics=None,
    vlm_prediction=None,
    readiness_result=None,
    task_context=None,
    policy_fit=None,
    community_comparison=None,
    action_divergence=None,
    episode_ranking=None,
    temporal_alignment=None,
    scaling_advice=None,
    consistency_report=None,
    meta_prediction=None,
    detail: bool = False,
    proxy_result=None,
    gemini_diagnosis=None,
) -> None:
    """Print dataset report using Rich.

    Default (compact): grade, prediction, top issues, next step.
    --detail: full narrative report with all diagnostics.
    """
    from rich.panel import Panel
    from rich.table import Table

    from orbit.utils.display import console

    # ── COMPACT MODE (default, no --detail) ────────────────────────────
    if not detail:
        _print_compact_report(
            info, readiness_result, meta_prediction, policy_fit,
            episode_ranking, action_divergence, signal_diagnostics,
            proxy_result=proxy_result,
            gemini_diagnosis=gemini_diagnosis,
        )
        return

    # --- Helper: progress bar string ---
    def _bar(score: float, width: int = 25) -> str:
        filled = int(score * width)
        return "\u2593" * filled + "\u2591" * (width - filled)

    def _grade(score: float) -> tuple[str, str]:
        if score >= 0.85:
            return "A", "green"
        if score >= 0.72:
            return "B", "green"
        if score >= 0.58:
            return "C", "yellow"
        if score >= 0.40:
            return "D", "yellow"
        return "F", "red"

    def _percentile_label(pct_str: str) -> str:
        return pct_str if pct_str else ""

    # --- HEADER BOX ---
    action_dim = task_context.action_dim if task_context else "?"
    robot_name = info.robot_type or "robot"
    policy_name = policy_fit.policy_type if policy_fit else "unknown"
    bimanual_str = "bimanual" if (task_context and task_context.is_bimanual) else ""
    header_line2 = f"{info.num_episodes} episodes \u00b7 {action_dim} action dims"
    if bimanual_str:
        header_line2 += f" \u00b7 {robot_name.upper()} {bimanual_str}"
    else:
        header_line2 += f" \u00b7 {robot_name.upper()}"
    header_line2 += f" \u00b7 {policy_name.upper()}"

    console.print()
    console.print(Panel(
        f"[bold]ORBIT Analysis: {info.repo_id}[/bold]\n{header_line2}",
        border_style="cyan",
        padding=(0, 2),
    ))

    if is_estimated:
        console.print(
            "[bold yellow]WARNING:[/bold yellow] Coverage analysis failed \u2014 "
            "showing estimated scores. Run with --full for accurate results.\n"
        )

    # --- DATASET READINESS (headline) ---
    if readiness_result is not None:
        rr = readiness_result
        grade_color = {"A": "green", "B": "green", "C": "yellow", "D": "yellow", "F": "red"}.get(rr.grade, "dim")
        console.print(
            f"[bold]Dataset Readiness:[/bold] "
            f"[bold {grade_color}]{rr.grade}[/bold {grade_color}] "
            f"(score: {rr.score}/100)"
        )
        console.print(f"{rr.summary}")
        console.print()
        for s in rr.strengths:
            console.print(f"  [green]\u2713[/green] {s}")
        for p in rr.problems:
            console.print(f"  [red]\u2717[/red] {p}")
        if rr.top_action:
            console.print()
            console.print(f"[bold]Top action:[/bold] {rr.top_action}")

        # Next-grade guidance: tell the user what it takes to reach the next grade
        if rr.grade in ("B", "C", "D", "F") and rr.problems:
            next_grade = {"B": "A", "C": "B", "D": "C", "F": "D"}.get(rr.grade, "?")
            next_threshold = {"B": 90, "C": 75, "D": 60, "F": 40}.get(rr.grade, 0)
            gap = next_threshold - rr.score
            if gap > 0 and rr.problems:
                top_fix = rr.problems[0]
                console.print()
                console.print(
                    f"  [dim]\u2192 To reach [bold]{next_grade}[/bold] "
                    f"(need {next_threshold}+, you're at {rr.score}): "
                    f"fix \"{top_fix}\"[/dim]"
                )
        console.print()

    # --- YOUR DATA AT A GLANCE ---
    console.print("[bold]YOUR DATA AT A GLANCE[/bold]")
    console.print("\u2500" * 40)

    # Episode count — only show percentile when the comparison set is
    # meaningful (same-scale datasets). Comparing 50 ALOHA demos against
    # foundation model datasets with 800K episodes is misleading.
    ep_pct_str = ""
    if community_comparison is not None:
        pct = community_comparison.your_episode_percentile
        # Suppress misleading percentiles when readiness says episodes are sufficient
        if readiness_result and "Sufficient episodes" in " ".join(readiness_result.strengths):
            pass  # Don't show "bottom 25%" when the grader says you have enough
        elif pct != "unknown":
            ep_pct_str = f"  ({pct})"

    # Coverage score (avg of coverage dimensions)
    cov_avg = (
        coverage.position_diversity + coverage.action_diversity +
        coverage.episode_consistency + coverage.temporal_coverage
    ) / 4.0

    # Signal health score — dead joints = 0, clipping = 0.5, healthy = 1.0
    signal_score = None
    if signal_diagnostics is not None:
        n_joints = len(signal_diagnostics.joint_diagnostics)
        if n_joints > 0:
            score_sum = sum(
                0.0 if j.is_dead else (0.5 if j.is_clipping else 1.0)
                for j in signal_diagnostics.joint_diagnostics
            )
            signal_score = score_sum / n_joints
        else:
            signal_score = 0.0

    glance_items = [
        ("Episodes", f"{info.num_episodes:<6}", None, ep_pct_str),
        ("Coverage", f"{cov_avg:.2f}  ", cov_avg, ""),
        ("Signal Health", f"{signal_score:.2f}  " if signal_score is not None else "N/A   ", signal_score, ""),
    ]

    for label, val_str, score, suffix in glance_items:
        if score is not None:
            bar_color = _score_color(score)
            console.print(
                f"  {label + ':':<16}{val_str}"
                f"[{bar_color}]{_bar(score)}[/{bar_color}]{suffix}"
            )
        else:
            console.print(f"  {label + ':':<16}{val_str}{suffix}")

    # Warn if position coverage may be inaccurate for high-dim state
    state_dim = 0
    if hasattr(info, "shapes") and info.shapes:
        for key in ("observation.state", "state"):
            if key in info.shapes:
                shape = info.shapes[key]
                state_dim = shape[0] if isinstance(shape, (list, tuple)) else shape
                break
    if state_dim > 6:
        console.print(
            f"  [dim]Note: Position coverage estimated from first 3 of {state_dim} state dims "
            "— may not reflect workspace if state ordering differs.[/dim]"
        )

    console.print()

    # --- WHAT'S WORKING ---
    working_items = []
    attention_items = []

    # From signal diagnostics
    if signal_diagnostics is not None:
        n_joints = len(signal_diagnostics.joint_diagnostics)
        n_active = sum(1 for j in signal_diagnostics.joint_diagnostics if not j.is_dead)
        n_dead = n_joints - n_active

        if n_dead == 0:
            working_items.append(f"All {n_joints} joints active \u2014 no dead servos detected")
        else:
            attention_items.append(
                (f"{n_dead} dead {'joint' if n_dead == 1 else 'joints'} detected (dimensions: {signal_diagnostics.dead_dimensions})",
                 "Re-check hardware and re-record affected episodes",
                 "HIGH")
            )

        # Gripper
        ga = signal_diagnostics.gripper_analysis
        if ga.gripper_dim is not None and ga.has_sufficient_grasps:
            working_items.append(
                f"Gripper transitions: {ga.transitions_per_episode:.1f}/episode (sufficient for grasping tasks)"
            )
        elif ga.gripper_dim is not None and not ga.has_sufficient_grasps:
            attention_items.append(
                (f"Gripper transitions low: {ga.transitions_per_episode:.1f}/episode (need \u22651.0)",
                 "Ensure each demonstration includes a full grasp-release cycle",
                 "HIGH")
            )

        # Episode consistency
        ep_lengths = [e.length for e in signal_diagnostics.episode_diagnostics]
        if ep_lengths:
            import statistics
            mean_len = statistics.mean(ep_lengths)
            if len(ep_lengths) > 1:
                stdev_len = statistics.stdev(ep_lengths)
                cv = stdev_len / mean_len if mean_len > 0 else 0
                if cv < 0.40:
                    working_items.append(f"Episode lengths consistent (CV={cv:.2f})")
                else:
                    attention_items.append(
                        (f"Episode lengths vary significantly (CV={cv:.2f})",
                         "Try to keep demonstration duration more consistent",
                         "LOW")
                    )

        # Flagged episodes — with specific reasons from episode ranker
        flagged = signal_diagnostics.episodes_to_review + signal_diagnostics.episodes_to_remove
        if flagged:
            ep_details = []
            for ed in signal_diagnostics.episode_diagnostics:
                if ed.episode_index in flagged and ed.jerk_spike_count > 0:
                    # Try to get specific reason from episode ranking
                    reason = "motion quality issues"
                    if episode_ranking is not None:
                        for er in episode_ranking.episodes:
                            if er.episode_index == ed.episode_index:
                                components = {
                                    "smoothness": er.smoothness,
                                    "completion": er.completion,
                                    "consistency": er.consistency,
                                    "length": er.length_score,
                                }
                                worst_name = min(components, key=components.get)
                                worst_val = components[worst_name]
                                reason = f"low {worst_name} ({worst_val:.2f})"
                                break
                    ep_details.append(f"Episode {ed.episode_index} ({reason})")
            if ep_details:
                ep_list = ", ".join(ep_details[:3])
                attention_items.append(
                    (ep_list,
                     f"Inspect episode {flagged[0]} for quality issues",
                     "MEDIUM")
                )

        # From passing checks — skip items already covered above
        existing_prefixes = {item.split("(")[0].strip().lower() for item in working_items}
        for p in signal_diagnostics.passing:
            prefix = p.split("(")[0].strip().lower()
            # Skip generic "All joints active" — already have detailed version
            if prefix == "all joints active":
                continue
            if prefix not in existing_prefixes and p not in [w for w, _, _ in attention_items]:
                working_items.append(p)
                existing_prefixes.add(prefix)

        # From blockers/warnings as attention items
        for b in signal_diagnostics.blockers:
            attention_items.append((b, "Fix before training", "HIGH"))
        for w in signal_diagnostics.warnings:
            # Avoid duplicating items already added
            if not any(w in item[0] for item in attention_items):
                attention_items.append((w, "", "MEDIUM"))

    # Coverage attention items — temporal diversity
    # For unimodal policies (ACT, BC), consistent patterns are a feature, not a bug.
    # Only warn for multimodal policies (diffusion, dp3) where diversity matters.
    _is_multimodal_policy = policy_fit and policy_fit.policy_type in ("diffusion_policy", "dp3")
    if coverage.temporal_coverage < 0.25 and _is_multimodal_policy:
        attention_items.append(
            (f"Temporal diversity low ({coverage.temporal_coverage:.2f}) \u2014 most episodes follow the same trajectory pattern",
             "Try varying your speed and approach angle between demonstrations",
             "MEDIUM")
        )

    # Print WHAT'S WORKING
    if working_items:
        console.print("[bold green]WHAT'S WORKING[/bold green]")
        for item in working_items:
            console.print(f"  [green]\u2713[/green] {item}")
        console.print()

    # Print WHAT NEEDS ATTENTION
    if attention_items:
        console.print("[bold yellow]WHAT NEEDS ATTENTION[/bold yellow]")
        for issue, fix, priority in attention_items:
            console.print(f"  [yellow]\u26a0[/yellow] {issue}")
            if fix:
                console.print(f"    [dim]\u2192 {fix}[/dim]")
        console.print()

    # --- COMMUNITY COMPARISON ---
    if community_comparison is not None:
        comp = community_comparison
        console.print("[bold]COMMUNITY COMPARISON[/bold]")
        console.print("\u2500" * 40)

        # Filtering note (domain-aware context)
        if comp.filtering_note:
            console.print(f"  [dim]{comp.filtering_note}[/dim]")
            console.print()

        n_similar = len(comp.similar_all)
        task_desc = task_context.estimated_task_type if task_context else "similar"
        console.print(
            f"  Your dataset vs {n_similar} similar {task_desc} datasets"
            f"{' on ' + robot_name.upper() if robot_name != 'robot' else ''}:"
        )

        your_eps = comp.your_episodes
        peer_eps = comp.peer_avg_episodes
        eps_sufficient = (
            readiness_result is not None
            and "Sufficient episodes" in " ".join(readiness_result.strengths)
        )
        if eps_sufficient or your_eps >= peer_eps:
            console.print(f"    Your episodes ({your_eps}) are sufficient for your policy")
        else:
            console.print(
                f"    Your episodes ({your_eps}) are below peer average ({int(peer_eps)})"
            )

        if comp.your_coverage is not None and comp.peer_avg_coverage is not None:
            console.print(
                f"    Your coverage ({comp.your_coverage:.2f}) is "
                f"{'above' if comp.your_coverage >= comp.peer_avg_coverage else 'below'} "
                f"average (peer avg: {comp.peer_avg_coverage:.2f})"
            )

        # Quality with peer context
        q = quality.overall_score
        if q >= 0.85:
            q_label = f"Your quality: {q:.2f} (above most peers)"
        elif q >= 0.70:
            q_label = f"Your quality: {q:.2f} (above average)"
        elif q >= 0.55:
            q_label = f"Your quality: {q:.2f} (average)"
        else:
            q_label = f"Your quality: {q:.2f} (below average)"
        console.print(
            f"    {q_label}"
        )

        # Closest successful dataset with domain + difficulty labels
        if comp.similar_successful:
            best = comp.similar_successful[0]
            domain_tag = "[SIM]" if best.is_sim else "[REAL]"
            diff_tag = f"[{best.difficulty_tier.upper()}]"
            console.print(
                f"\n  Closest successful dataset: "
                f"[bold]{best.dataset}[/bold] {domain_tag} {diff_tag}"
            )
            console.print(
                f"    {best.num_episodes} episodes, {best.success_rate:.0%} success ({best.policy})"
            )

        console.print()

    # --- COMPATIBLE DATASETS FOR MIXING ---
    if community_comparison is not None and community_comparison.similar_successful:
        console.print("[bold]COMPATIBLE DATASETS FOR MIXING[/bold]")
        console.print("\u2500" * 40)
        for i, peer in enumerate(community_comparison.similar_successful[:2]):
            stars = "\u2605" * min(3, int(peer.success_rate * 3 + 0.5))
            domain_tag = "[SIM]" if peer.is_sim else "[REAL]"
            diff_tag = f"[{peer.difficulty_tier.upper()}]"
            console.print(
                f"\n  [bold]{peer.dataset}[/bold] {stars} {domain_tag} {diff_tag}"
            )
            console.print(
                f"    {peer.num_episodes} eps \u00b7 {peer.success_rate:.0%} success \u00b7 "
                f"Strategy: {'Direct merge' if i == 0 else 'Pretrain+finetune'}"
            )
        console.print()

    # --- POLICY FIT ---
    if policy_fit is not None:
        console.print("[bold]POLICY FIT[/bold]")
        console.print("\u2500" * 40)

        if hasattr(policy_fit, 'recommended_policies') and policy_fit.recommended_policies:
            for rp in policy_fit.recommended_policies:
                if isinstance(rp, dict):
                    name = rp.get("name", rp.get("policy", "?"))
                    score = rp.get("score", rp.get("fit_score", 0))
                    note = rp.get("note", "")
                else:
                    name = str(rp)
                    score = policy_fit.fit_score
                    note = ""
                score_color = _score_color(score) if isinstance(score, (int, float)) else "dim"
                suffix = "  \u2014 " + note if note else ""
                console.print(
                    f"  {name}: [{score_color}]{score:.2f}[/{score_color}]{suffix}"
                )
        else:
            console.print(
                f"  Best fit: [bold]{policy_fit.policy_type}[/bold] "
                f"(score [{_score_color(policy_fit.fit_score)}]{policy_fit.fit_score:.2f}[/{_score_color(policy_fit.fit_score)}])"
            )
            if policy_fit.specific_warnings:
                for w in policy_fit.specific_warnings:
                    console.print(f"    [dim]\u2022 {w}[/dim]")

        console.print()

    # --- EPISODE QUALITY RANKING ---
    if episode_ranking is not None and episode_ranking.episodes:
        er = episode_ranking
        sampled = er.keep_count + er.review_count + er.remove_count
        total_in_dataset = info.num_episodes
        is_sampled = sampled < total_in_dataset
        console.print("[bold]EPISODE QUALITY RANKING[/bold]")
        if is_sampled:
            console.print(f"[dim]Based on {sampled}/{total_in_dataset} sampled episodes (use --full to check all)[/dim]")
        console.print("\u2500" * 40)
        console.print(
            f"  Keep: {er.keep_count}/{sampled} ({er.keep_count * 100 // max(1, sampled)}%)  \u2502  "
            f"Review: {er.review_count} ({er.review_count * 100 // max(1, sampled)}%)  \u2502  "
            f"Remove: {er.remove_count} ({er.remove_count * 100 // max(1, sampled)}%)"
        )
        if er.remove_count > 0:
            worst = sorted(er.episodes, key=lambda e: e.composite_score)[:3]
            worst_str = ", ".join(f"#{e.episode_index} (score: {e.composite_score:.2f})" for e in worst)
            console.print(f"  Worst episodes: {worst_str}")
            console.print(
                f"  Removing them would improve quality: "
                f"{er.current_quality_estimate:.2f} \u2192 {er.improved_quality_estimate:.2f}"
            )
        console.print()

    # --- ACTION DIVERGENCE ---
    if action_divergence is not None and action_divergence.score is not None:
        ad = action_divergence
        sev_color = {"low": "green", "moderate": "yellow", "high": "yellow", "critical": "red"}.get(ad.severity, "dim")
        console.print("[bold]ACTION DIVERGENCE[/bold]")
        console.print("\u2500" * 40)
        console.print(f"  Score: [{sev_color}]{ad.score:.2f}[/{sev_color}] ({ad.severity})")
        if ad.estimated_success_penalty > 0:
            if ad.severity == "high" or ad.severity == "critical":
                console.print(f"  Some states have contradictory actions — consider filtering or switching to Diffusion Policy")
        console.print()

    # --- TEMPORAL ALIGNMENT ---
    if temporal_alignment is not None:
        ta = temporal_alignment
        console.print("[bold]TEMPORAL ALIGNMENT[/bold]")
        console.print("\u2500" * 40)
        if ta.aligned:
            if ta.severity == "negligible":
                console.print(
                    f"  [green]\u2713[/green] State-action alignment: lag = {ta.estimated_lag} frames (negligible, <2% of episode)"
                )
            else:
                console.print(
                    f"  [green]\u2713[/green] State-action alignment: OK (lag = {ta.estimated_lag} frames)"
                )
        else:
            # Dampen severity label for low-dim data where lag detection is unreliable
            action_dim = 0
            if hasattr(info, "shapes") and info.shapes:
                for key in ("action",):
                    if key in info.shapes:
                        shape = info.shapes[key]
                        action_dim = shape[0] if isinstance(shape, (list, tuple)) else shape
            displayed_severity = ta.severity
            if action_dim > 0 and action_dim < 6 and ta.severity == "severe":
                displayed_severity = "mild (low-dim data)"
            console.print(
                f"  [yellow]\u26a0[/yellow] State-action alignment: lag = {ta.estimated_lag} frames ({displayed_severity})"
            )
        if ta.fps_jitter > 0:
            jitter_ok = ta.fps_jitter < 0.05
            check = "[green]\u2713[/green]" if jitter_ok else "[yellow]\u26a0[/yellow]"
            fps_str = ta.details.get("measured_fps", "?")
            console.print(
                f"  {check} FPS consistency: {fps_str} fps (jitter={ta.fps_jitter:.3f})"
            )
        if ta.constant_dimensions:
            console.print(
                f"  [yellow]\u26a0[/yellow] Constant state dims: {ta.constant_dimensions}"
            )
        console.print()

    # --- SCALING ADVICE ---
    # Only show when dataset has multiple environments — the ICLR 2025
    # scaling law is about multi-environment generalization, not
    # single-task imitation learning.
    if (
        scaling_advice is not None
        and scaling_advice.bottleneck != "quality"
        and scaling_advice.estimated_diversity > 1
    ):
        console.print("[bold]SCALING ADVICE[/bold]")
        console.print("\u2500" * 40)
        console.print(f"  {scaling_advice.advice}")
        console.print(
            f"  [dim]Estimated diversity: ~{scaling_advice.estimated_diversity} environments, "
            f"~{scaling_advice.demos_per_setting:.0f} demos each[/dim]"
        )
        console.print()

    # --- CONSISTENCY DIAGNOSTICS ---
    if consistency_report is not None:
        cv_items = [
            ("Path length", consistency_report.path_length_cv, "trajectory distance varies"),
            ("Jerk", consistency_report.cartesian_jerk_cv, "smoothness varies"),
            ("Curvature", consistency_report.curvature_cv, "path shape varies"),
            ("Effort", consistency_report.effort_cv, "action magnitude varies"),
            ("Joint limit", consistency_report.joint_limit_cv, "joint range usage varies"),
            ("Orientation", consistency_report.orientation_cv, "rotation patterns vary"),
        ]
        elevated = [(name, cv, desc) for name, cv, desc in cv_items if cv > 0.40]

        if elevated:
            console.print("[bold]CONSISTENCY DIAGNOSTICS[/bold]")
            console.print("\u2500" * 40)
            for name, cv, desc in elevated:
                level = "high" if cv > 0.5 else "moderate"
                console.print(f"  {name} CV: {cv:.2f} ({level}) \u2014 {desc}")
            if consistency_report.per_episode_scores:
                ranked = sorted(consistency_report.per_episode_scores.items(), key=lambda x: x[1])
                worst = [str(idx) for idx, _ in ranked[:3]]
                console.print(f"  Most inconsistent episodes: {', '.join(worst)}")
            console.print()

    # --- NEXT STEPS (prioritized, deduplicated) ---
    console.print("[bold]NEXT STEPS (prioritized)[/bold]")
    console.print("\u2500" * 40)

    # Merge recommendations from all sources and deduplicate
    next_steps = []
    seen = set()

    # From attention items (highest priority since they're specific issues)
    for issue, fix, priority in (attention_items if signal_diagnostics else []):
        if fix and fix not in seen:
            seen.add(fix)
            next_steps.append((priority, fix if fix != "Fix before training" else issue))

    # From readiness grade problems
    if readiness_result and readiness_result.problems:
        for prob in readiness_result.problems[:3]:
            if prob not in seen:
                seen.add(prob)
                next_steps.append(("HIGH", prob))

    # From standard recommendations — filter out canned advice for grade-A
    is_grade_a = readiness_result and readiness_result.grade == "A"
    is_unimodal = policy_fit and policy_fit.policy_type in ("act", "bc", "bc_rnn")
    _canned_patterns = ("nearly fails but recovers", "different strategies", "follow the same")
    for rec in recs[:5]:
        # Suppress generic advice for grade-A datasets
        if is_grade_a and any(p in rec.description for p in _canned_patterns):
            continue
        # "Try different strategies" is harmful for unimodal policies
        if is_unimodal and "different strategies" in rec.description:
            continue
        if rec.description not in seen:
            seen.add(rec.description)
            next_steps.append((rec.priority, rec.description))

    # Sort: HIGH first, then MEDIUM, then LOW
    priority_order = {"HIGH": 0, "MEDIUM": 1, "LOW": 2}
    next_steps.sort(key=lambda x: priority_order.get(x[0], 1))

    for priority, step in next_steps[:6]:
        style = _priority_style(priority)
        console.print(f"  [{style}][{priority}][/{style}] {step}")

    console.print()
    if readiness_result and readiness_result.grade in ("C", "D", "F"):
        console.print(
            f"[bold]Next steps:[/bold]\n"
            f"  1. orbit clean {info.repo_id}              # Fix data issues\n"
            f"  2. orbit suggest {info.repo_id} --policy {policy_fit.policy_type if policy_fit else 'act'}  "
            f"# Get training command\n"
        )
    elif readiness_result and readiness_result.grade == "B":
        pol = policy_fit.policy_type if policy_fit else "act"
        console.print(
            f"[dim]Unsure whether to train? Run:[/dim]\n"
            f"  orbit analyze {info.repo_id} --policy {pol} --proxy  "
            f"# 2-3 min proxy training for a real go/no-go signal\n"
        )
    else:
        console.print(
            f"[dim]Run `orbit suggest {info.repo_id}` to get a ready-to-run training command.[/dim]\n"
        )



@main.command()
@click.argument("dataset")
@click.option("--output", "-o", default=None, help="Output repo ID or local path for cleaned dataset.")
@click.option("--policy", "-p", default=None, help="Target policy (act, diffusion, bc, smolvla).")
@click.option("--aggressive", is_flag=True, help="Also remove 'review' episodes, not just 'remove'.")
@click.option("--dry-run", is_flag=True, help="Show what would be removed without modifying anything.")
@click.option("--keep-dead-joints", is_flag=True, help="Don't flag dead joint dimensions.")
@click.option("--json", "as_json", is_flag=True, help="Output as JSON.")
def clean(dataset, output, policy, aggressive, dry_run, keep_dead_joints, as_json):
    """Clean a dataset by removing bad episodes and flagging problems.

    Runs signal diagnostics and episode ranking, then produces a cleaned
    dataset with bad episodes removed and dead joints optionally dropped.

    Examples:
        orbit clean lerobot/my-dataset
        orbit clean lerobot/my-dataset -o user/my-dataset-clean --aggressive
        orbit clean lerobot/my-dataset --dry-run
        orbit clean ./local-dataset -p act
    """
    from orbit.analyzer.dataset_loader import load_dataset_info, load_episode_data
    from orbit.cleaner import (
        build_cleaning_plan,
        execute_cleaning,
        print_cleaning_result,
        print_dry_run,
    )
    from orbit.utils.display import console

    try:
        info = load_dataset_info(dataset)
    except (ValueError, ConnectionError) as e:
        console.print(f"\n[bold red]Error:[/bold red] {e}")
        sys.exit(1)

    if not as_json and info.num_episodes > 500:
        console.print(
            f"[yellow]\u26a0 Dataset has {info.num_episodes} episodes — "
            "this may take a few minutes.[/yellow]"
        )

    if not as_json:
        console.print(f"Loading all {info.num_episodes} episodes for cleaning...")

    try:
        episode_data = load_episode_data(dataset, max_episodes=None)
    except (ValueError, ConnectionError) as e:
        console.print(f"\n[bold red]Error:[/bold red] Could not load episode data: {e}")
        sys.exit(1)

    if episode_data is None or episode_data.empty:
        console.print("[bold red]Error:[/bold red] No episode data loaded.")
        sys.exit(1)

    plan = build_cleaning_plan(
        dataset_repo=dataset,
        episode_data=episode_data,
        info=info,
        aggressive=aggressive,
        keep_dead_joints=keep_dead_joints,
        policy=policy,
    )

    if as_json:
        import dataclasses
        import json as _json

        result_dict = dataclasses.asdict(plan)
        if not dry_run:
            result = execute_cleaning(plan, output=output)
            result_dict = dataclasses.asdict(result)
        click.echo(_json.dumps(result_dict, indent=2, default=str))
        return

    if dry_run:
        print_dry_run(plan)
        return

    result = execute_cleaning(plan, output=output)
    print_cleaning_result(result)


@main.command()
@click.argument("dataset")
@click.option("--policy", "-p", default="bc", help="Target policy (act, diffusion_policy, bc, smolvla).")
@click.option("--json", "as_json", is_flag=True, help="Output as JSON.")
def improve(dataset, policy, as_json):
    """Analyze, identify bad episodes, and PROVE that removing them helps.

    Runs the full analysis → episode ranking → proxy training pipeline
    twice (before and after outlier removal) to demonstrate measurable
    improvement in data learnability.

    Examples:
        orbit improve lerobot/xarm_lift_medium --policy bc
        orbit improve lerobot/droid_100 --policy act
    """
    import numpy as np

    from orbit.analyzer.dataset_loader import load_dataset_info, load_episode_data
    from orbit.analyzer.episode_ranker import compute_episode_ranking
    from orbit.analyzer.signal_diagnostics import compute_signal_diagnostics
    from orbit.utils.display import console

    # Step 1: Load dataset
    try:
        info = load_dataset_info(dataset)
    except (ValueError, ConnectionError) as e:
        console.print(f"\n[bold red]Error:[/bold red] {e}")
        sys.exit(1)

    if not as_json:
        console.print(f"\n[bold]orbit improve: {dataset}[/bold]")
        console.print(f"  Loading {info.num_episodes} episodes...")

    try:
        episode_data = load_episode_data(dataset, max_episodes=200)
    except (ValueError, ConnectionError) as e:
        console.print(f"\n[bold red]Error:[/bold red] Could not load episode data: {e}")
        sys.exit(1)

    if episode_data is None or episode_data.empty:
        console.print("[bold red]Error:[/bold red] No episode data loaded.")
        sys.exit(1)

    # Extract action episodes
    action_episodes: list[np.ndarray] = []
    state_episodes: list[np.ndarray] = []
    for ep_idx in sorted(episode_data["episode_index"].unique()):
        ep = episode_data[episode_data["episode_index"] == ep_idx]
        if "action" in ep.columns and ep["action"].notna().any():
            actions = np.stack(ep["action"].dropna().values)
            if actions.ndim >= 2:
                action_episodes.append(actions)
        if "state" in ep.columns and ep["state"].notna().any():
            states = np.stack(ep["state"].dropna().values)
            if states.ndim >= 2:
                state_episodes.append(states)

    if len(action_episodes) < 5:
        console.print("[bold red]Error:[/bold red] Need at least 5 episodes with action data.")
        sys.exit(1)

    # Step 2: Rank episodes to find outliers
    if not as_json:
        console.print("  Ranking episodes...")
    ranking = compute_episode_ranking(
        action_episodes,
        state_episodes=state_episodes if state_episodes else None,
    )

    # Use "remove" episodes, or worst "review" episodes if no removes
    remove_indices = set(ranking.remove_indices)
    if not remove_indices:
        # Aggressive mode: take the worst-scoring review episodes
        review_eps = sorted(
            [ep for ep in ranking.episodes if ep.verdict == "review"],
            key=lambda e: e.composite_score,
        )
        # Remove the worst quartile of review episodes
        n_worst = max(1, len(review_eps) // 4)
        remove_indices = {ep.episode_index for ep in review_eps[:n_worst]}

    n_remove = len(remove_indices)

    if n_remove == 0:
        if as_json:
            click.echo(json.dumps({
                "dataset": dataset,
                "policy": policy,
                "outlier_count": 0,
                "message": "No outlier episodes found — data is already clean.",
            }, indent=2))
        else:
            console.print("\n  [green]No outlier episodes found — data is already clean.[/green]")
        return

    if not as_json:
        console.print(f"  Identified [bold]{n_remove}[/bold] outlier episodes: {sorted(remove_indices)[:10]}")

    # Step 3: Proxy training BEFORE removal (all episodes)
    try:
        from orbit.analyzer.proxy_trainer import run_proxy_training
    except ImportError:
        console.print("[bold red]Error:[/bold red] Proxy training requires torch. Install with: pip install torch")
        sys.exit(1)

    if not as_json:
        console.print("\n  [bold]Running proxy training on ALL episodes (before)...[/bold]")
    proxy_before = run_proxy_training(
        action_episodes,
        state_episodes=state_episodes if state_episodes else None,
        policy_type="bc",  # always MLP for comparable baseline
    )
    loss_before = proxy_before.validation_loss_final
    ratio_before = 1.0
    if proxy_before.validation_loss_curve and proxy_before.validation_loss_curve[0] > 1e-8:
        ratio_before = loss_before / proxy_before.validation_loss_curve[0]

    # Step 4: Remove outlier episodes
    clean_actions = [a for i, a in enumerate(action_episodes) if i not in remove_indices]
    clean_states = [s for i, s in enumerate(state_episodes) if i not in remove_indices] if state_episodes else None

    if len(clean_actions) < 3:
        console.print("[bold red]Error:[/bold red] Too few episodes remaining after removal.")
        sys.exit(1)

    # Step 5: Proxy training AFTER removal
    if not as_json:
        console.print(f"  [bold]Running proxy training on {len(clean_actions)} clean episodes (after)...[/bold]")
    proxy_after = run_proxy_training(
        clean_actions,
        state_episodes=clean_states,
        policy_type="bc",  # same MLP baseline
    )
    loss_after = proxy_after.validation_loss_final
    ratio_after = 1.0
    if proxy_after.validation_loss_curve and proxy_after.validation_loss_curve[0] > 1e-8:
        ratio_after = loss_after / proxy_after.validation_loss_curve[0]

    # Step 6: Report
    improvement_pct = 0.0
    if ratio_before > 1e-8:
        improvement_pct = (1.0 - ratio_after / ratio_before) * 100

    result = {
        "dataset": dataset,
        "policy": policy,
        "total_episodes": len(action_episodes),
        "outlier_count": n_remove,
        "outlier_indices": sorted(remove_indices),
        "clean_episodes": len(clean_actions),
        "before": {
            "proxy_val_loss": round(loss_before, 6),
            "loss_ratio": round(min(ratio_before, 1.0), 4),
            "go_no_go": proxy_before.go_no_go,
        },
        "after": {
            "proxy_val_loss": round(loss_after, 6),
            "loss_ratio": round(min(ratio_after, 1.0), 4),
            "go_no_go": proxy_after.go_no_go,
        },
        "improvement_pct": round(improvement_pct, 1),
        "recommendation": (
            f"Run `orbit clean {dataset} --policy {policy}` to apply changes"
            if improvement_pct > 5 else
            "Marginal improvement — consider collecting more data instead"
        ),
    }

    if as_json:
        click.echo(json.dumps(result, indent=2))
    else:
        console.print(f"\n  {'='*50}")
        console.print(f"  [bold]IMPROVEMENT REPORT[/bold]")
        console.print(f"  {'='*50}")
        console.print(f"  Episodes: {len(action_episodes)} total, {n_remove} outliers removed")
        console.print(f"  Proxy val_loss BEFORE: {loss_before:.6f} (loss_ratio={min(ratio_before,1.0):.4f})")
        console.print(f"  Proxy val_loss AFTER:  {loss_after:.6f} (loss_ratio={min(ratio_after,1.0):.4f})")

        if improvement_pct > 0:
            console.print(f"  [bold green]Improvement: {improvement_pct:.1f}%[/bold green]")
        else:
            console.print(f"  [yellow]No improvement ({improvement_pct:.1f}%)[/yellow]")

        console.print(f"\n  {result['recommendation']}")
        console.print()


@main.command()
@click.argument("dataset")
@click.option("--policy", "-p", default=None, help="Target policy (auto-selects if not specified).")
@click.option("--gpu-memory", default=None, type=int, help="Available GPU memory in GB.")
@click.option(
    "--framework", "-f",
    default="lerobot",
    type=click.Choice(["lerobot", "openvla", "openpi", "groot", "osmo", "custom", "all"]),
    help="Training framework (default: lerobot). Use 'all' to see all options.",
)
@click.option("--json", "as_json", is_flag=True, help="Output as JSON.")
def suggest(dataset, policy, gpu_memory, framework, as_json):
    """Suggest training configuration for a dataset.

    Analyzes the dataset and recommends a policy, hyperparameters,
    and a ready-to-run training command for the chosen framework.

    \b
    Examples:
        orbit suggest lerobot/my-dataset
        orbit suggest lerobot/my-dataset --policy act --gpu-memory 16
        orbit suggest lerobot/my-dataset --framework openvla
        orbit suggest lerobot/my-dataset --framework all
    """
    from orbit.analyzer.dataset_loader import load_dataset_info, load_episode_data
    from orbit.utils.display import console

    try:
        info = load_dataset_info(dataset)
    except (ValueError, ConnectionError) as e:
        console.print(f"\n[bold red]Error:[/bold red] {e}")
        sys.exit(1)

    # Quick analysis — only need a sample for suggestion
    try:
        episode_data = load_episode_data(dataset, max_episodes=100)
    except (ValueError, ConnectionError) as e:
        console.print(f"\n[bold red]Error:[/bold red] Could not load episode data: {e}")
        sys.exit(1)

    if framework == "lerobot":
        from orbit.suggester import generate_suggestion, print_suggestion

        suggestion = generate_suggestion(
            dataset_repo=dataset,
            info=info,
            episode_data=episode_data,
            policy=policy,
            gpu_memory_gb=gpu_memory,
        )
        if as_json:
            import dataclasses
            import json as _json

            console.print(_json.dumps(dataclasses.asdict(suggestion), indent=2, default=str))
        else:
            print_suggestion(suggestion)
    else:
        from orbit.suggester import generate_multi_framework_suggestion, print_multi_framework_suggestion

        base, configs = generate_multi_framework_suggestion(
            dataset_repo=dataset,
            info=info,
            episode_data=episode_data,
            policy=policy,
            gpu_memory_gb=gpu_memory,
            framework=framework,
        )
        if as_json:
            import dataclasses
            import json as _json

            result = {
                "base": dataclasses.asdict(base),
                "frameworks": {k: dataclasses.asdict(v) for k, v in configs.items()},
            }
            console.print(_json.dumps(result, indent=2, default=str))
        else:
            print_multi_framework_suggestion(base, configs)


@main.command()
@click.argument("training_dir")
@click.option("--dataset", "-d", default=None, help="Dataset repo ID (auto-detected from config if not specified).")
@click.option("--ai", is_flag=True, help="Enable AI-powered diagnosis with data-aware reasoning.")
@click.option("--json", "as_json", is_flag=True, help="Output as JSON.")
def debug(training_dir, dataset, ai, as_json):
    """Diagnose a failed or underperforming training run.

    Reads training logs from a LeRobot output directory, identifies
    what went wrong, and connects failures to data issues.

    \b
    Examples:
        orbit debug outputs/train/my-run
        orbit debug outputs/train/my-run --ai
        orbit debug outputs/train/my-run --ai -d lerobot/my-dataset
    """
    if as_json:
        import json as _json
        from pathlib import Path

        from orbit.debugger import (
            _analyze_dynamics,
            _find_config,
            _parse_lerobot_logs,
            _parse_wandb_logs,
            _rule_based_diagnosis,
        )
        from orbit.utils.display import console

        training_path = Path(training_dir)
        if not training_path.exists():
            console.print(_json.dumps({"error": f"Training directory not found: {training_dir}"}))
            sys.exit(1)

        log_entries = _parse_lerobot_logs(training_path)
        wandb_entries = _parse_wandb_logs(training_path)
        all_entries = sorted(log_entries + wandb_entries, key=lambda x: x.get("step", 0))
        dynamics = _analyze_dynamics(all_entries)
        config = _find_config(training_path)
        findings = _rule_based_diagnosis(dynamics, config, dataset)

        result = {
            "status": dynamics["status"],
            "message": dynamics["message"],
            "total_steps": dynamics["total_steps"],
            "loss_curve": dynamics.get("loss_curve_summary", {}),
            "findings": findings,
            "dataset": dataset,
        }
        console.print(_json.dumps(result, indent=2, default=str))
        return

    from orbit.debugger import debug_training

    debug_training(training_dir, dataset_id=dataset, use_ai=ai)


@main.command()
@click.argument("source")
@click.option("--to", "output_format", default="lerobot-v3", help="Target format (default: lerobot-v3).")
@click.option("--output", "-o", required=True, help="Output directory path.")
@click.option(
    "--format",
    "source_format",
    default=None,
    type=click.Choice(["auto", "lerobot", "lerobot-local", "rlds", "hdf5", "rosbag", "directory"]),
    help="Override source format auto-detection.",
)
@click.option("--action-topic", default=None, help="ROS topic for actions (rosbag only).")
@click.option("--state-topic", default=None, help="ROS topic for states (rosbag only).")
@click.option("--action-key", default=None, help="Column name for action data (e.g. 'end_pose', 'joint_command').")
@click.option("--state-key", default=None, help="Column name for state data (e.g. 'qpos', 'joint_positions').")
@click.option("--episode-gap", default=None, type=float, help="Episode segmentation gap (rosbag only).")
@click.option("--fps", default=None, type=float, help="Override FPS for output dataset.")
def convert(
    source: str,
    output_format: str,
    output: str,
    source_format: str | None,
    action_topic: str | None,
    state_topic: str | None,
    action_key: str | None,
    state_key: str | None,
    episode_gap: float | None,
    fps: float | None,
):
    """Convert a dataset from any format to LeRobot v3.

    Examples:
        orbit convert ./recording.mcap --to lerobot-v3 -o ./my-dataset/
        orbit convert ./demo.hdf5 --to lerobot-v3 -o ./my-dataset/
        orbit convert ./rlds-dataset/ --to lerobot-v3 -o ./my-dataset/
        orbit convert lerobot/some-dataset --to lerobot-v3 -o ./converted/
    """
    from orbit.loaders.convert import convert as do_convert
    from orbit.utils.display import console

    loader_kwargs: dict = {}
    if action_topic:
        loader_kwargs["action_topic"] = action_topic
    if state_topic:
        loader_kwargs["state_topic"] = state_topic
    if action_key:
        loader_kwargs["action_key"] = action_key
    if state_key:
        loader_kwargs["state_key"] = state_key
    if episode_gap is not None:
        loader_kwargs["episode_gap"] = episode_gap

    try:
        out_path = do_convert(
            source=source,
            output_dir=output,
            output_format=output_format,
            source_format=source_format,
            fps=fps,
            **loader_kwargs,
        )
        console.print(f"\n[green]Dataset converted successfully → {out_path}[/green]")
    except (ValueError, ImportError, ConnectionError) as e:
        console.print(f"\n[bold red]Error:[/bold red] {e}")
        sys.exit(1)


@main.command()
@click.argument("task_description")
@click.option("--robot", default="so100", help="Robot type (so100, so101, koch, aloha).")
@click.option("--policy", default="act", help="Policy type (act, diffusion_policy, smolvla).")
@click.option("--save", "save_path", default=None, help="Save plan to JSON file.")
@click.option("--format", "fmt", default=None, type=click.Choice(["json"]), help="Output format.")
def plan(task_description: str, robot: str, policy: str, save_path: str | None, fmt: str | None):
    """Plan data collection strategy.

    \b
    Examples:
        orbit plan "pick up cups" --robot so100 --policy act
        orbit plan "insert peg" --robot koch --save plan.json
        orbit plan "fold towel" --robot aloha --policy diffusion_policy
    """
    from orbit.planner.playbook import generate_playbook
    from orbit.utils.display import console

    playbook = generate_playbook(task_description, robot, policy)

    if save_path:
        playbook.save(save_path)
        console.print(f"\n[green]Plan saved to {save_path}[/green]")

    if fmt == "json":
        click.echo(json.dumps(playbook.to_dict(), indent=2))
        return

    _print_playbook(playbook)


def _print_playbook(playbook) -> None:
    """Print a formatted playbook using Rich."""
    from orbit.utils.display import console, print_header, print_section

    print_header("ORBIT Data Collection Playbook")

    console.print(f"\n[bold]Task:[/bold] {playbook.task_description}")
    console.print(f"[bold]Task type:[/bold] {playbook.task_type}")
    console.print(f"[bold]Robot:[/bold] {playbook.robot_type}")
    console.print(f"[bold]Policy:[/bold] {playbook.policy_type}")
    console.print(f"[bold]Total target episodes:[/bold] {playbook.total_target_episodes}")

    # Robot-specific tips
    if playbook.robot_config:
        print_section("Robot Setup")
        rc = playbook.robot_config
        dims = rc.workspace_dims
        console.print(f"  Workspace: {dims[0]}m x {dims[1]}m x {dims[2]}m")
        console.print(f"  Camera: {rc.recommended_camera_position}")
        res = rc.recommended_resolution
        console.print(f"  FPS: {rc.recommended_fps}  |  Resolution: {res[0]}x{res[1]}")
        console.print(f"  Arms: {rc.arm_count}")
        console.print(f"  [dim]{rc.notes}[/dim]")

    # Phases
    for i, phase in enumerate(playbook.phases, 1):
        print_section(f"Phase {i}: {phase.name} ({phase.target_episodes} episodes)")
        for req_name, req_desc in phase.requirements.items():
            label = req_name.replace("_", " ").title()
            console.print(f"  [green]\u2713[/green] {label}: [bold]{req_desc}[/bold]")

    # Estimates
    print_section("Estimates")
    console.print(f"  80% success rate: {playbook.estimates.eighty_percent_success}")
    console.print(f"  95% success rate: {playbook.estimates.ninety_five_percent_success}")

    # Policy notes
    if playbook.policy_notes:
        print_section("Policy Notes")
        for note in playbook.policy_notes:
            console.print(f"  [dim]\u2022 {note}[/dim]")

    console.print()


@main.command()
@click.argument("query", required=False, default=None)
@click.option("--task", default=None, help="Filter by task type (pick_and_place, insertion, pushing, etc.).")
@click.option("--policy", "policy_filter", default=None, help="Filter by policy (act, diffusion_policy, bc, etc.).")
@click.option("--min-success", type=float, default=None, help="Only show entries with success rate >= this value.")
@click.option("--max-success", type=float, default=None, help="Only show entries with success rate <= this value.")
@click.option("--json", "as_json", is_flag=True, help="Output as JSON.")
@click.option("--top", type=int, default=None, help="Show only top N entries by success rate.")
def benchmark(
    query: str | None,
    task: str | None,
    policy_filter: str | None,
    min_success: float | None,
    max_success: float | None,
    as_json: bool,
    top: int | None,
):
    """Browse the benchmark database of 82 validated results.

    Search by keyword, filter by task type or policy, or browse all entries.

    Examples:

      orbit benchmark                     # show all 82 entries
      orbit benchmark aloha               # search for 'aloha' entries
      orbit benchmark --task insertion     # filter by task type
      orbit benchmark --policy act --top 5 # top 5 ACT results
      orbit benchmark --min-success 0.8    # only high-performing entries
    """
    from orbit.utils.display import console

    # Load ground truth
    from pathlib import Path

    gt_path = Path(__file__).parent / "analyzer" / "data" / "ground_truth_condensed.json"
    try:
        with open(gt_path) as f:
            entries = json.load(f)
    except (FileNotFoundError, json.JSONDecodeError) as e:
        console.print(f"[bold red]Error:[/bold red] Could not load benchmark data: {e}")
        sys.exit(1)

    # Apply filters
    filtered = entries

    if query:
        q = query.lower()
        filtered = [
            e for e in filtered
            if q in e.get("id", "").lower()
            or q in e.get("dataset", "").lower()
            or q in e.get("task_type", "").lower()
            or q in e.get("policy", "").lower()
            or q in e.get("features_summary", "").lower()
        ]

    if task:
        t = task.lower()
        filtered = [e for e in filtered if t in e.get("task_type", "").lower()]

    if policy_filter:
        p = policy_filter.lower()
        filtered = [e for e in filtered if p in e.get("policy", "").lower()]

    if min_success is not None:
        filtered = [e for e in filtered if e.get("success_rate", 0) >= min_success]

    if max_success is not None:
        filtered = [e for e in filtered if e.get("success_rate", 1) <= max_success]

    # Sort by success rate descending
    filtered.sort(key=lambda e: e.get("success_rate", 0), reverse=True)

    if top:
        filtered = filtered[:top]

    if not filtered:
        if as_json:
            click.echo(json.dumps([]))
        else:
            console.print("[yellow]No matching benchmark entries found.[/yellow]")
            if query:
                console.print(f"[dim]Search query: '{query}'[/dim]")
        return

    if as_json:
        click.echo(json.dumps(filtered, indent=2))
        return

    _print_benchmark_table(filtered, len(entries), query=query, task=task, policy=policy_filter)


def _print_benchmark_table(
    entries: list[dict],
    total: int,
    query: str | None = None,
    task: str | None = None,
    policy: str | None = None,
) -> None:
    """Print benchmark entries as a formatted table."""
    from pathlib import Path

    from orbit.utils.display import console, print_header

    # Header
    filter_desc = []
    if query:
        filter_desc.append(f"query='{query}'")
    if task:
        filter_desc.append(f"task={task}")
    if policy:
        filter_desc.append(f"policy={policy}")
    filter_str = f" ({', '.join(filter_desc)})" if filter_desc else ""

    print_header(f"ORBIT Benchmark Database — {len(entries)}/{total} entries{filter_str}")

    # Group by task type for readability
    task_groups: dict[str, list[dict]] = {}
    for e in entries:
        tt = e.get("task_type", "unknown")
        task_groups.setdefault(tt, []).append(e)

    for task_type, group in task_groups.items():
        console.print(f"\n  [bold]{task_type.replace('_', ' ').title()}[/bold] ({len(group)} entries)")
        console.print(f"  {'─' * 90}")

        for e in group:
            rate = e.get("success_rate", 0)
            if rate >= 0.80:
                color = "green"
            elif rate >= 0.50:
                color = "yellow"
            else:
                color = "red"

            eps = e.get("num_episodes", "?")
            eps_str = f"{eps:>6,}" if isinstance(eps, int) else f"{eps:>6}"

            console.print(
                f"  [{color}]{rate:5.0%}[/{color}]  "
                f"{e.get('id', '?'):<42} "
                f"policy={e.get('policy', '?'):<20} "
                f"eps={eps_str}"
            )

    # Summary stats
    rates = [e.get("success_rate", 0) for e in entries]
    console.print(f"\n  [dim]Success rate range: {min(rates):.0%} – {max(rates):.0%} | "
                  f"Median: {sorted(rates)[len(rates)//2]:.0%} | "
                  f"Mean: {sum(rates)/len(rates):.0%}[/dim]")

    # Available task types
    all_tasks = sorted(set(e.get("task_type", "") for e in entries))
    all_policies = sorted(set(e.get("policy", "") for e in entries))
    console.print(f"  [dim]Task types: {', '.join(all_tasks)}[/dim]")
    console.print(f"  [dim]Policies: {', '.join(all_policies)}[/dim]")
    console.print()


@main.command()
@click.argument("plan_file")
@click.option("--dataset", required=True, help="HuggingFace dataset repo ID.")
def track(plan_file: str, dataset: str):
    """Track training progress against a collection plan."""
    from orbit.analyzer.coverage import compute_coverage
    from orbit.analyzer.dataset_loader import load_dataset_info, load_episode_data
    from orbit.tracker.progress import compute_progress, load_plan_from_file
    from orbit.utils.display import console

    # Load the plan
    try:
        playbook = load_plan_from_file(plan_file)
    except FileNotFoundError:
        console.print(
            f"\n[bold red]Error:[/bold red] Plan file not found: {plan_file}\n"
            f"Run [bold]orbit plan[/bold] first to create one."
        )
        sys.exit(1)
    except (json.JSONDecodeError, KeyError, ValueError) as e:
        console.print(f"\n[bold red]Error:[/bold red] Invalid plan file: {e}")
        sys.exit(1)

    # Load dataset info
    try:
        info = load_dataset_info(dataset)
    except (ValueError, ConnectionError) as e:
        console.print(f"\n[bold red]Error:[/bold red] {e}")
        sys.exit(1)

    # Load episode data for coverage analysis
    try:
        episode_data = load_episode_data(dataset)
    except (ValueError, ConnectionError):
        episode_data = None

    # Compute coverage
    is_estimated = False
    if episode_data is not None and not episode_data.empty:
        coverage = compute_coverage(info, episode_data)
    else:
        coverage = _synthetic_coverage(info)
        is_estimated = True

    # Compute progress
    report = compute_progress(playbook, info, coverage)

    _print_tracking_report(report, playbook, dataset, is_estimated=is_estimated)


def _print_tracking_report(
    report,
    playbook,
    dataset: str,
    is_estimated: bool = False,
) -> None:
    """Print a formatted tracking report using Rich."""
    from orbit.utils.display import console, print_header, print_section

    print_header("ORBIT Progress Tracker")

    if is_estimated:
        console.print(
            "\n[bold yellow]WARNING:[/bold yellow] Coverage analysis failed — "
            "showing estimated scores based on episode count only. "
            "Run with --embeddings for accurate results.\n"
        )

    console.print(f"\n[bold]Plan:[/bold] {playbook.task_description}")
    console.print(f"[bold]Dataset:[/bold] {dataset}")

    # Progress bar
    filled = int(report.progress_percent / 100 * 16)
    bar = "\u2588" * filled + "\u2591" * (16 - filled)
    console.print(
        f"\n[bold]Episode Progress:[/bold] {bar} "
        f"{report.current_episodes}/{report.total_target_episodes} "
        f"({report.progress_percent:.0f}%)"
    )

    phase_display = "Core demonstrations" if report.phase == "core" else "Diversity expansion"
    console.print(f"[bold]Phase:[/bold] {phase_display}")

    # Coverage checklist
    print_section("Coverage Checklist")
    for item in report.coverage_checklist:
        if item.status == "complete":
            icon = "[green]\u2713[/green]"
        elif item.status == "in_progress":
            icon = "[yellow]\u25cb[/yellow]"
        else:
            icon = "[red]\u2717[/red]"

        if item.name == "Near-failure recoveries":
            console.print(
                f"  {icon} {item.name:<30} "
                f"({int(item.current)} found \u2014 target: {int(item.target)})"
            )
        else:
            console.print(
                f"  {icon} {item.name:<30} ({item.current:.2f} \u2014 target: {item.target:.2f})"
            )

    # Ready to train
    if report.ready_to_train:
        console.print("\n[bold]Ready to Train?[/bold] [bold green]YES[/bold green]")
    else:
        console.print("\n[bold]Ready to Train?[/bold] [bold yellow]NOT YET[/bold yellow]")

    # Next session priorities
    if report.next_session_priorities:
        print_section("Next Session Priorities")
        for i, priority in enumerate(report.next_session_priorities, 1):
            console.print(f"  {i}. {priority}")

    console.print()


# ---------------------------------------------------------------------------
# Credentials helpers
# ---------------------------------------------------------------------------

_CREDENTIALS_DIR = "~/.config/orbit"
_CREDENTIALS_FILE = "credentials.json"


def _credentials_path() -> str:
    import os
    return os.path.join(os.path.expanduser(_CREDENTIALS_DIR), _CREDENTIALS_FILE)


def _load_credentials() -> dict | None:
    import os
    path = _credentials_path()
    if not os.path.exists(path):
        return None
    with open(path) as f:
        return json.load(f)


def _save_credentials(token: str, api_url: str) -> None:
    import os
    dirpath = os.path.expanduser(_CREDENTIALS_DIR)
    os.makedirs(dirpath, exist_ok=True)
    with open(_credentials_path(), "w") as f:
        json.dump({"token": token, "api_url": api_url}, f, indent=2)
    os.chmod(_credentials_path(), 0o600)


def _remove_credentials() -> None:
    import os
    path = _credentials_path()
    if os.path.exists(path):
        os.remove(path)


def _api_headers(token: str) -> dict:
    return {"Authorization": f"Bearer {token}", "Content-Type": "application/json"}


# ---------------------------------------------------------------------------
# orbit auth — unified API key management
# ---------------------------------------------------------------------------


@main.command()
@click.argument("action", type=click.Choice(["set", "show", "clear"]))
@click.argument("key_name", required=False, default=None)
@click.argument("key_value", required=False, default=None)
def auth(action: str, key_name: str | None, key_value: str | None):
    """Manage API keys for ORBIT features.

    Stores keys in ~/.config/orbit/credentials.json so you don't need
    to export environment variables every session.

    \b
    Examples:
        orbit auth set google-api-key AIza...     # Store Gemini key for --ai/--deep/--vlm
        orbit auth set orbit-token eyJ...         # Store dashboard token (same as orbit login)
        orbit auth show                           # Show stored keys (masked)
        orbit auth clear google-api-key           # Remove a stored key
    """
    import os

    from orbit.utils.display import console

    creds = _load_credentials() or {}

    if action == "set":
        if not key_name or not key_value:
            console.print("[bold red]Error:[/bold red] Usage: orbit auth set <key-name> <value>")
            console.print("  Key names: google-api-key, orbit-token")
            sys.exit(1)

        key_map = {
            "google-api-key": "google_api_key",
            "google_api_key": "google_api_key",
            "orbit-token": "token",
            "token": "token",
        }
        internal_key = key_map.get(key_name)
        if not internal_key:
            console.print(f"[bold red]Error:[/bold red] Unknown key '{key_name}'")
            console.print("  Supported keys: google-api-key, orbit-token")
            sys.exit(1)

        creds[internal_key] = key_value
        if internal_key == "token" and "api_url" not in creds:
            creds["api_url"] = "https://api.orbit-robotics.com"

        # Save
        dirpath = os.path.expanduser(_CREDENTIALS_DIR)
        os.makedirs(dirpath, exist_ok=True)
        with open(_credentials_path(), "w") as f:
            json.dump(creds, f, indent=2)
        os.chmod(_credentials_path(), 0o600)

        console.print(f"[bold green]✓[/bold green] Stored {key_name}")

        # Also set as env var for current session
        if internal_key == "google_api_key":
            os.environ["GOOGLE_API_KEY"] = key_value
            console.print("[dim]  Also set GOOGLE_API_KEY for this session[/dim]")

    elif action == "show":
        if not creds:
            console.print("[dim]No stored keys. Run: orbit auth set google-api-key <key>[/dim]")
            return
        display_map = {
            "google_api_key": "google-api-key",
            "token": "orbit-token",
            "api_url": "api-url",
        }
        for k, v in creds.items():
            name = display_map.get(k, k)
            if isinstance(v, str) and len(v) > 8:
                masked = v[:4] + "..." + v[-4:]
            else:
                masked = str(v)
            console.print(f"  {name}: {masked}")

    elif action == "clear":
        if not key_name:
            console.print("[bold red]Error:[/bold red] Specify which key to clear")
            sys.exit(1)
        key_map = {"google-api-key": "google_api_key", "orbit-token": "token"}
        internal_key = key_map.get(key_name, key_name)
        if internal_key in creds:
            del creds[internal_key]
            dirpath = os.path.expanduser(_CREDENTIALS_DIR)
            os.makedirs(dirpath, exist_ok=True)
            with open(_credentials_path(), "w") as f:
                json.dump(creds, f, indent=2)
            os.chmod(_credentials_path(), 0o600)
            console.print(f"[bold green]✓[/bold green] Cleared {key_name}")
        else:
            console.print(f"[dim]Key '{key_name}' not found[/dim]")


# ---------------------------------------------------------------------------
# orbit login
# ---------------------------------------------------------------------------

@main.command()
@click.option("--token", required=True, help="CLI API token from your dashboard.")
@click.option(
    "--api-url",
    default="https://api.orbit-robotics.com",
    help="API base URL (for self-hosted or dev setups).",
)
def login(token: str, api_url: str):
    """Authenticate the CLI with your ORBIT dashboard."""
    import httpx
    from orbit.utils.display import console

    api_url = api_url.rstrip("/")
    try:
        resp = httpx.get(f"{api_url}/auth/me", headers=_api_headers(token), timeout=10)
        if resp.status_code != 200:
            console.print("[bold red]Error:[/bold red] Invalid token or server error.")
            sys.exit(1)
        user_data = resp.json()
        username = user_data.get("hf_username", "unknown")
    except httpx.RequestError as e:
        console.print(f"[bold red]Error:[/bold red] Could not reach API — {e}")
        sys.exit(1)

    _save_credentials(token, api_url)
    console.print(f"\n[bold green]✓[/bold green] Logged in as [bold]@{username}[/bold]. CLI synced with your dashboard.")


# ---------------------------------------------------------------------------
# orbit logout
# ---------------------------------------------------------------------------

@main.command()
def logout():
    """Disconnect CLI from your ORBIT dashboard."""
    from orbit.utils.display import console

    _remove_credentials()
    console.print("\n[bold]Logged out.[/bold] CLI disconnected from dashboard.")


# ---------------------------------------------------------------------------
# orbit whoami
# ---------------------------------------------------------------------------

@main.command()
def whoami():
    """Show the currently authenticated user."""
    import httpx
    from orbit.utils.display import console

    creds = _load_credentials()
    if not creds:
        console.print("[bold yellow]Not logged in.[/bold yellow] Run: orbit login --token <token>")
        sys.exit(1)

    try:
        resp = httpx.get(
            f"{creds['api_url']}/auth/me",
            headers=_api_headers(creds["token"]),
            timeout=10,
        )
        if resp.status_code != 200:
            console.print("[bold red]Error:[/bold red] Token is invalid or expired. Run: orbit login --token <new-token>")
            sys.exit(1)
        user_data = resp.json()
    except httpx.RequestError as e:
        console.print(f"[bold red]Error:[/bold red] Could not reach API — {e}")
        sys.exit(1)

    username = user_data.get("hf_username", "unknown")
    console.print(f"\n[bold]@{username}[/bold]")
    console.print(f"[dim]Dashboard: {creds['api_url'].replace('api.', 'app.')}[/dim]")


# ---------------------------------------------------------------------------
# orbit push
# ---------------------------------------------------------------------------

@main.command()
@click.argument("results_file")
@click.option("--project", "project_id", type=int, default=None, help="Target project ID.")
def push(results_file: str, project_id: int | None):
    """Push analysis results to your ORBIT dashboard."""
    import httpx
    from orbit.utils.display import console

    creds = _load_credentials()
    if not creds:
        console.print("[bold yellow]Not logged in.[/bold yellow] Run: orbit login --token <token>")
        sys.exit(1)

    try:
        with open(results_file) as f:
            data = json.load(f)
    except (OSError, json.JSONDecodeError) as e:
        console.print(f"[bold red]Error:[/bold red] Could not read results file — {e}")
        sys.exit(1)

    # Build upload payload
    payload = {
        "dataset_repo": data.get("dataset_repo", data.get("dataset", "")),
    }
    if project_id:
        payload["project_id"] = project_id

    # Map known fields
    for key in (
        "quality_score", "estimated_success_rate", "coverage_results",
        "recommendations", "embedding_clusters", "embedding_pca",
        "phase_detection", "signal_diagnostics", "task_context",
        "policy_fit", "readiness", "community_comparison",
        "signal_diagnostics_full", "ai_assessment", "llm_assessment",
    ):
        if key in data:
            payload[key] = data[key]

    try:
        resp = httpx.post(
            f"{creds['api_url']}/analysis/upload",
            headers=_api_headers(creds["token"]),
            json=payload,
            timeout=30,
        )
        if resp.status_code == 201:
            result = resp.json()
            console.print(f"\n[bold green]✓[/bold green] Results synced.")
            console.print(f"  View at: [bold]{result['url']}[/bold]")
        else:
            console.print(f"[bold red]Error:[/bold red] Upload failed ({resp.status_code}): {resp.text[:200]}")
            sys.exit(1)
    except httpx.RequestError as e:
        console.print(f"[bold red]Error:[/bold red] Could not reach API — {e}")
        sys.exit(1)


# ---------------------------------------------------------------------------
# orbit report
# ---------------------------------------------------------------------------

VALID_POLICIES = ["act", "diffusion_policy", "dp3", "bc", "bc_rnn", "smolvla", "vla", "other"]


@main.command()
@click.argument("dataset_repo")
@click.option("--policy", required=True, type=click.Choice(VALID_POLICIES), help="Policy type used.")
@click.option("--success-rate", required=True, type=float, help="Success rate (0.0-1.0).")
@click.option("--eval-trials", required=True, type=int, help="Number of evaluation trials.")
@click.option("--task", default=None, help="Task name (e.g. 'pick up cup').")
@click.option("--notes", default=None, help="Additional context (e.g. 'trained for 2000 epochs').")
@click.option("--local-only", is_flag=True, help="Save locally only, don't push to community database.")
def report(dataset_repo: str, policy: str, success_rate: float, eval_trials: int, task: str | None, notes: str | None, local_only: bool):
    """Report training results — saves locally and to ORBIT community database.

    \b
    Results are always saved to ~/.orbit/reports/ for your records.
    Unless --local-only is set, they are also submitted to the community database.

    \b
    Examples:
        orbit report lerobot/my-dataset --policy act --success-rate 0.82 --eval-trials 20
        orbit report lerobot/my-dataset --policy act --success-rate 0.82 --eval-trials 20 \\
            --notes "Trained on curated 50 episodes, removed 12 outliers" --local-only
    """
    from datetime import datetime, timezone
    from pathlib import Path

    from orbit.utils.display import console

    # Validate inputs
    if "/" not in dataset_repo or dataset_repo.endswith("/") or dataset_repo.startswith("/"):
        console.print(
            f"\n[bold red]Error:[/bold red] Invalid repo ID '{dataset_repo}'. "
            "Expected format: owner/dataset_name (e.g. lerobot/my_dataset)"
        )
        sys.exit(1)

    if success_rate < 0.0 or success_rate > 1.0:
        console.print(
            f"\n[bold red]Error:[/bold red] success-rate must be between 0.0 and 1.0, got {success_rate}"
        )
        sys.exit(1)

    if eval_trials < 1:
        console.print(
            f"\n[bold red]Error:[/bold red] eval-trials must be >= 1, got {eval_trials}"
        )
        sys.exit(1)

    # Always save locally first
    reports_dir = Path.home() / ".orbit" / "reports"
    reports_dir.mkdir(parents=True, exist_ok=True)

    timestamp = datetime.now(timezone.utc).strftime("%Y%m%d_%H%M%S")
    safe_name = dataset_repo.replace("/", "_")
    local_report = {
        "dataset_repo": dataset_repo,
        "policy": policy,
        "success_rate": success_rate,
        "eval_trials": eval_trials,
        "task": task,
        "notes": notes,
        "timestamp": datetime.now(timezone.utc).isoformat(),
    }

    report_path = reports_dir / f"{safe_name}_{policy}_{timestamp}.json"
    with open(report_path, "w") as f:
        json.dump(local_report, f, indent=2)

    pct = int(success_rate * 100)
    console.print(f"\n[bold green]\u2713[/bold green] Report saved locally: {report_path}")
    console.print(f"  Dataset:      {dataset_repo}")
    console.print(f"  Policy:       {policy.upper()}")
    console.print(f"  Success Rate: {pct}% ({eval_trials} eval trials)")
    if notes:
        console.print(f"  Notes:        {notes}")

    if local_only:
        console.print("\n[dim]Local-only mode — not submitted to community database.[/dim]")
        return

    # Submit to community database
    import httpx

    creds = _load_credentials()
    if not creds:
        console.print(
            "\n[bold yellow]Not logged in.[/bold yellow] "
            "Result saved locally but not submitted to community database.\n"
            "Run: [bold]orbit login --token <token>[/bold] to enable community sharing."
        )
        return

    payload = {
        "dataset_repo": dataset_repo,
        "policy": policy,
        "success_rate": success_rate,
        "eval_trials": eval_trials,
    }
    if task:
        payload["task_name"] = task
    if notes:
        payload["notes"] = notes

    try:
        resp = httpx.post(
            f"{creds['api_url']}/api/benchmark/community-report",
            headers=_api_headers(creds["token"]),
            json=payload,
            timeout=15,
        )
        if resp.status_code in (200, 201):
            console.print(
                "\n  [bold green]\u2713[/bold green] Result submitted to ORBIT community database."
            )
            console.print(
                f"  View it at: https://orbit-robotics.com/explore/dataset/{dataset_repo}"
            )
            console.print("\n  Thank you for contributing to the robotics community!")
        else:
            error_detail = resp.text[:200]
            console.print(f"[bold red]Error:[/bold red] Community submission failed ({resp.status_code}): {error_detail}")
            console.print("[dim]Your result was still saved locally.[/dim]")
    except httpx.RequestError as e:
        console.print(f"[bold red]Error:[/bold red] Could not reach API — {e}")
        console.print("[dim]Your result was still saved locally.[/dim]")


# ---------------------------------------------------------------------------
# Phase 2: curate + gate commands
# ---------------------------------------------------------------------------


def _load_episodes(dataset_repo: str, source_format: str | None = None):
    """Load dataset and extract action/state episodes. Returns (info, episode_data, action_episodes, state_episodes)."""
    from orbit.analyzer.dataset_loader import load_dataset_info, load_episode_data

    info = load_dataset_info(dataset_repo, format=source_format)
    # Load all episodes — callers like orbit curate need the full dataset
    episode_data = load_episode_data(dataset_repo, format=source_format, max_episodes=info.num_episodes)

    action_episodes = []
    state_episodes = []
    if episode_data is not None and not episode_data.empty:
        from orbit.analyzer.coverage import _extract_array

        for _, ep_df in episode_data.groupby("episode_index"):
            if "action" in ep_df.columns:
                ep_actions = _extract_array(ep_df["action"])
                if ep_actions is not None and ep_actions.size > 0:
                    action_episodes.append(ep_actions)
            if "state" in ep_df.columns:
                ep_states = _extract_array(ep_df["state"])
                if ep_states is not None and ep_states.size > 0:
                    state_episodes.append(ep_states)

    return info, episode_data, action_episodes, state_episodes


@main.command()
@click.argument("dataset_repo")
@click.option("--budget", "-b", type=int, default=None, help="Keep only the top N episodes.")
@click.option("--output", "-o", default=None, help="Path to save curated episode list.")
@click.option(
    "--method",
    type=click.Choice(["auto", "consistency", "mi"]),
    default="auto",
    help="Curation method.",
)
@click.option("--json", "as_json", is_flag=True, help="Output as JSON.")
@click.option(
    "--format",
    "source_format",
    default=None,
    type=click.Choice(["auto", "lerobot", "lerobot-local", "rlds", "hdf5", "rosbag", "directory"]),
    help="Override format auto-detection.",
)
def curate(dataset_repo: str, budget: int | None, output: str | None, method: str, as_json: bool, source_format: str | None):
    """Select the best episodes from a dataset.

    Ranks every episode by quality using consistency metrics and mutual
    information scoring. Optionally outputs a curated subset.

    \b
    Examples:
        orbit curate lerobot/my-dataset --budget 50
        orbit curate ./recordings/ --budget 100 -o curated_episodes.json
    """
    from orbit.utils.display import console

    if not as_json:
        console.print(f"\n[bold]Curating:[/bold] {dataset_repo}")

    try:
        info, _, action_episodes, state_episodes = _load_episodes(dataset_repo, source_format)
    except Exception as e:
        console.print(f"[bold red]Error:[/bold red] Could not load dataset — {e}")
        sys.exit(1)

    if not action_episodes:
        console.print("[bold red]Error:[/bold red] No action data found in dataset.")
        sys.exit(1)

    from orbit.analyzer.data_curator import curate_episodes

    result = curate_episodes(
        action_episodes,
        state_episodes=state_episodes if state_episodes else None,
        budget=budget,
        method=method,
    )

    if as_json:
        report = {
            "dataset": dataset_repo,
            "total_episodes": len(action_episodes),
            "method": result.method,
            "ranked_episodes": result.ranked_episodes,
            "recommended_keep": result.recommended_keep,
            "recommended_remove": result.recommended_remove,
            "budget_curve": result.budget_curve,
            "summary": result.summary,
        }
        click.echo(json.dumps(report, indent=2, default=_json_default))
    else:
        console.print(f"\n  Dataset: {dataset_repo} ({len(action_episodes)} episodes)")
        console.print(f"  Method:  {result.method}")
        console.print(f"\n  {result.summary}")

        if result.recommended_remove:
            preview = result.recommended_remove[:10]
            console.print(
                f"\n  [yellow]Remove:[/yellow] episodes {preview}"
                + (f" and {len(result.recommended_remove) - 10} more"
                   if len(result.recommended_remove) > 10 else "")
            )

        if result.budget_curve:
            console.print("\n  Budget curve:")
            for n_eps, quality in result.budget_curve:
                bar = "█" * int(quality * 20)
                console.print(f"    {n_eps:>4} episodes → quality {quality:.2f} {bar}")

    if output:
        out_data = {
            "dataset": dataset_repo,
            "method": result.method,
            "keep_episodes": result.recommended_keep,
            "remove_episodes": result.recommended_remove,
        }
        with open(output, "w") as f:
            json.dump(out_data, f, indent=2)
        console.print(f"\n  Saved curated episode list to: {output}")


@main.command()
@click.argument("dataset_repo")
@click.option("--policy", "-p", required=True, help="Target policy (act, diffusion, smolvla, bc).")
@click.option("--min-grade", default="C", type=click.Choice(["A", "B", "C", "D"]), help="Minimum passing grade.")
@click.option(
    "--format",
    "source_format",
    default=None,
    type=click.Choice(["auto", "lerobot", "lerobot-local", "rlds", "hdf5", "rosbag", "directory"]),
    help="Override format auto-detection.",
)
@click.option("--json", "as_json", is_flag=True, help="Output as JSON.")
@click.option(
    "--output-format",
    "output_fmt",
    default=None,
    type=click.Choice(["text", "json", "github"]),
    help="Output format. 'github' emits GitHub Actions annotations.",
)
def gate(dataset_repo: str, policy: str, min_grade: str, source_format: str | None, as_json: bool, output_fmt: str | None):
    """CI/CD quality gate — pass/fail for training readiness.

    Returns exit code 0 (pass) or 1 (fail). Designed for integration
    into training pipelines as a pre-training check.

    \b
    Examples:
        orbit gate lerobot/my-dataset -p act
        orbit gate ./recordings/ -p diffusion --min-grade B
        orbit gate lerobot/my-dataset -p act --output-format github
    """
    # --json is shorthand for --output-format json
    if as_json and output_fmt is None:
        output_fmt = "json"
    from orbit.utils.display import console

    if output_fmt not in ("json", "github"):
        console.print(f"\n[bold]Quality Gate:[/bold] {dataset_repo} → {policy}")

    try:
        info, episode_data, action_episodes, state_episodes = _load_episodes(
            dataset_repo, source_format
        )
    except Exception as e:
        if as_json:
            click.echo(json.dumps({"passed": False, "error": str(e)}))
        else:
            console.print(f"[bold red]FAIL:[/bold red] Could not load dataset — {e}")
        sys.exit(1)

    if not action_episodes:
        if as_json:
            click.echo(json.dumps({"passed": False, "error": "No action data found"}))
        else:
            console.print("[bold red]FAIL:[/bold red] No action data found.")
        sys.exit(1)

    # Consistency metrics
    consistency_report = None
    if len(action_episodes) >= 2:
        try:
            from orbit.analyzer.consistency_metrics import compute_consistency_metrics

            consistency_report = compute_consistency_metrics(
                action_episodes,
                state_episodes=state_episodes if state_episodes else None,
            )
        except Exception:
            pass

    # Full analysis for grading (not hardcoded placeholders)
    from orbit.analyzer.action_divergence import compute_action_divergence
    from orbit.analyzer.coverage import compute_coverage
    from orbit.analyzer.episode_ranker import compute_episode_ranking
    from orbit.analyzer.policy_fit import PolicyFitAnalyzer
    from orbit.analyzer.quality import compute_quality_score
    from orbit.analyzer.success_predictor import DatasetReadinessGrader
    from orbit.analyzer.task_context import TaskContextAnalyzer

    coverage = compute_coverage(info, episode_data)
    episode_ranking_result = compute_episode_ranking(action_episodes)

    # Quality score — real computation, not placeholder
    quality = compute_quality_score(
        info, coverage,
        action_episodes=action_episodes if action_episodes else None,
        state_episodes=state_episodes if state_episodes else None,
    )

    task_context = TaskContextAnalyzer().analyze(
        action_episodes,
        metadata={"robot_type": info.robot_type, "task_description": dataset_repo},
    )

    # Action divergence — real computation, not 0.0
    ad_result = compute_action_divergence(
        action_episodes, state_episodes=state_episodes if state_episodes else None,
    )
    ad_score = ad_result.score if ad_result and ad_result.score is not None else 0.0

    # Signal diagnostics — real values, not empty dicts
    signal_diag = None
    try:
        from orbit.analyzer.signal_diagnostics import compute_signal_diagnostics

        if action_episodes:
            signal_diag = compute_signal_diagnostics(action_episodes)
    except Exception:
        pass

    sig_health = {
        "dead_dimensions": signal_diag.dead_dimensions if signal_diag else [],
        "blockers": signal_diag.blockers if signal_diag else [],
        "warnings": signal_diag.warnings if signal_diag else [],
        "total_joints": (
            len(signal_diag.joint_diagnostics) if signal_diag else task_context.action_dim
        ),
        "clipping_joints": (
            [jd.joint_index for jd in signal_diag.joint_diagnostics if jd.is_clipping]
            if signal_diag else []
        ),
    }

    pf_analyzer = PolicyFitAnalyzer()
    policy_fit_result = pf_analyzer.analyze(
        action_episodes,
        policy_type=policy,
        episode_count=info.num_episodes,
        metadata={"is_bimanual": task_context.is_bimanual},
        task_complexity=task_context.task_complexity_score,
        consistency_score=coverage.episode_consistency,
        multimodality_score=ad_score if ad_score > 0 else task_context.trajectory_multimodality,
    )

    grader = DatasetReadinessGrader()
    readiness = grader.grade(
        quality_score=quality.overall_score * 100,
        episode_consistency=coverage.episode_consistency,
        signal_health=sig_health,
        policy_fit_score=policy_fit_result.fit_score,
        episode_ranking={
            "total": len(action_episodes),
            "remove_count": episode_ranking_result.remove_count,
            "review_count": episode_ranking_result.review_count,
        },
        action_divergence=ad_score,
        episode_count=info.num_episodes,
        policy_type=policy_fit_result.policy_type,
        consistency=consistency_report,
    )

    grade_order = {"A": 4, "B": 3, "C": 2, "D": 1, "F": 0}
    passed = grade_order.get(readiness.grade, 0) >= grade_order.get(min_grade, 2)

    if output_fmt == "json":
        result = {
            "passed": passed,
            "grade": readiness.grade,
            "score": readiness.score,
            "min_grade": min_grade,
            "summary": readiness.summary,
            "problems": readiness.problems[:5] if readiness.problems else [],
            "top_action": readiness.top_action,
            "consistency": (
                round(consistency_report.overall_consistency, 3) if consistency_report else None
            ),
            "action_divergence": round(ad_score, 4),
            "quality_score": round(quality.overall_score, 4),
        }
        click.echo(json.dumps(result, indent=2, default=_json_default))
        sys.exit(0 if passed else 1)

    if output_fmt == "github":
        # GitHub Actions annotation format
        # https://docs.github.com/en/actions/using-workflows/workflow-commands-for-github-actions
        if passed:
            click.echo(f"::notice title=ORBIT Gate PASS::Grade {readiness.grade} ({readiness.score}/100) >= {min_grade}")
        else:
            click.echo(f"::error title=ORBIT Gate FAIL::Grade {readiness.grade} ({readiness.score}/100) < {min_grade}")
        for problem in readiness.problems[:5]:
            click.echo(f"::warning title=Data Issue::{problem}")
        if readiness.top_action:
            click.echo(f"::notice title=Next Step::{readiness.top_action}")
        sys.exit(0 if passed else 1)

    if passed:
        console.print(
            f"\n  [bold green]PASS[/bold green]  Grade: {readiness.grade} "
            f"(minimum: {min_grade})"
        )
        console.print(f"  {readiness.summary}")
        if consistency_report:
            console.print(
                f"  Consistency: {consistency_report.overall_consistency:.2f}"
            )
        sys.exit(0)
    else:
        console.print(
            f"\n  [bold red]FAIL[/bold red]  Grade: {readiness.grade} "
            f"(minimum: {min_grade})"
        )
        console.print(f"  {readiness.summary}")
        if readiness.problems:
            for p in readiness.problems[:3]:
                console.print(f"  • {p}")
        console.print(f"\n  Top action: {readiness.top_action}")
        sys.exit(1)


# ---------------------------------------------------------------------------
# Phase 3: orbit coach
# ---------------------------------------------------------------------------


@main.command()
@click.argument("dataset_path")
@click.option("--target", "-t", default=50, type=int, help="Target number of episodes to collect.")
@click.option("--policy", "-p", default="act", help="Target policy (act, diffusion, smolvla).")
@click.option("--poll-interval", default=2.0, type=float, help="Seconds between directory polls.")
def coach(dataset_path: str, target: int, policy: str, poll_interval: float):
    """Real-time data collection coach.

    Watches a dataset directory as new episodes are added and gives
    live feedback on consistency, smoothness, and quality.

    \b
    Examples:
        orbit coach ./my-dataset/ --target 50 --policy act
        orbit coach lerobot/my-dataset --target 100 --poll-interval 5
    """
    from orbit.coach import watch_and_coach

    watch_and_coach(
        dataset_path=dataset_path,
        target_episodes=target,
        policy=policy,
        poll_interval=poll_interval,
    )


# ---------------------------------------------------------------------------
# Phase 3: orbit generate
# ---------------------------------------------------------------------------


@main.command()
@click.argument("task_description")
@click.option("--sim", "simulator", default="mujoco", type=click.Choice(["mujoco", "isaac"]),
              help="Simulator backend.")
@click.option("--model", "model_path", required=True, help="Path to simulator model file (.xml for MuJoCo).")
@click.option("--episodes", "-n", default=10, type=int, help="Number of demonstrations to generate.")
@click.option("--output", "-o", default="./generated/", help="Output directory.")
@click.option("--max-attempts", default=20, type=int, help="Max LLM attempts per episode.")
def generate(task_description: str, simulator: str, model_path: str, episodes: int,
             output: str, max_attempts: int):
    """Generate robot demonstrations using an LLM agent in simulation.

    Uses an LLM to write control code in a ReAct loop, executing
    in MuJoCo until the task succeeds. Records successful trajectories.

    \b
    Caveats (from FAEA research):
        - Uses privileged state info, not camera images
        - 2-8 seconds per decision (not real-time)
        - ~$0.01-$5 per successful episode
        - Best for: generating seed data for VLA fine-tuning

    \b
    Examples:
        orbit generate "pick up the red cube" --sim mujoco --model ./scene.xml -n 20
        orbit generate "push block to target" --sim mujoco --model ./env.xml -o ./data/
    """
    from orbit.generator.generate import generate_demonstrations
    from orbit.utils.display import console

    console.print(f"\n[bold]Demo Generator[/bold]")
    console.print(f"  Task: {task_description}")
    console.print(f"  Simulator: {simulator}")
    console.print(f"  Target: {episodes} episodes")
    console.print()

    result = generate_demonstrations(
        task_description=task_description,
        simulator=simulator,
        model_path=model_path,
        num_episodes=episodes,
        max_attempts_per_episode=max_attempts,
        output_dir=output,
    )

    console.print(f"\n[bold]Results:[/bold]")
    console.print(f"  Successful: {result.successful_episodes}/{episodes}")
    console.print(f"  Failed: {result.failed_episodes}")
    console.print(f"  Cost: {result.cost_estimate}")
    console.print(f"  Output: {result.output_dir}")

    if result.successful_episodes > 0:
        console.print(f"\n  Next: orbit analyze {result.output_dir}")
    console.print()


# ---------------------------------------------------------------------------
# Phase 4: orbit monitor
# ---------------------------------------------------------------------------


@main.command()
@click.argument("training_dir")
@click.option("--dataset", "-d", default=None, help="Dataset repo ID for context.")
@click.option("--poll-interval", default=10.0, type=float, help="Seconds between log checks.")
def monitor(training_dir: str, dataset: str | None, poll_interval: float):
    """Monitor a training run in real-time.

    Watches training logs and alerts on plateaus, NaN loss,
    gradient instability, and other problems.

    \b
    Examples:
        orbit monitor ./outputs/train/my-run/
        orbit monitor ./outputs/train/my-run/ -d lerobot/my-dataset
    """
    from orbit.monitor import monitor_training

    monitor_training(
        training_dir=training_dir,
        dataset=dataset,
        poll_interval=poll_interval,
    )


# ---------------------------------------------------------------------------
# Phase 4: orbit verify
# ---------------------------------------------------------------------------


@main.command()
@click.argument("training_dir")
@click.option("--dataset", "-d", default=None, help="Dataset repo ID for prediction comparison.")
@click.option("--success-rate", default=None, type=float,
              help="Actual success rate from rollout evaluation (0.0-1.0).")
@click.option("--json", "as_json", is_flag=True, help="Output as JSON.")
def verify(training_dir: str, dataset: str | None, success_rate: float | None, as_json: bool):
    """Verify training outcome against predicted quality.

    Compares predicted dataset grade against actual training
    convergence. Stores calibration data for improving future predictions.

    \b
    Examples:
        orbit verify ./outputs/train/my-run/
        orbit verify ./outputs/train/my-run/ -d lerobot/my-dataset
        orbit verify ./outputs/train/my-run/ --success-rate 0.73
    """
    from orbit.verify import print_verification, verify_training

    result = verify_training(
        training_dir=training_dir,
        dataset=dataset,
        success_rate=success_rate,
    )

    if as_json:
        import dataclasses
        import json as _json

        from orbit.utils.display import console

        result_dict = dataclasses.asdict(result)
        console.print(_json.dumps(result_dict, indent=2, default=str))
        return

    print_verification(result, dataset=dataset)


# ---------------------------------------------------------------------------
# Phase 4: orbit train (pipeline)
# ---------------------------------------------------------------------------


@main.command("train", context_settings={"ignore_unknown_options": True})
@click.option("--gate", default=None, help='Quality gate args: "<policy> [--min-grade <grade>] [dataset]".')
@click.option("--skip-gate", is_flag=True, help="Skip the quality gate.")
@click.option("--auto-clean", is_flag=True, help="Auto-fix data issues before training.")
@click.option("--dry-run", is_flag=True, help="Show what would happen without training.")
@click.argument("training_command", nargs=-1, type=click.UNPROCESSED)
def train_cmd(gate: str | None, skip_gate: bool, auto_clean: bool, dry_run: bool,
              training_command: tuple[str, ...]):
    """Run the full training pipeline: gate -> YOUR command -> verify.

    Uses passthrough design — ORBIT wraps ANY training command. Everything
    after -- is passed directly to subprocess.run().

    \b
    Examples:
        orbit train --gate "act --min-grade B" -- lerobot-train --dataset.repo_id=./data --policy.type=act
        orbit train --gate "diffusion --min-grade C" --auto-clean -- python finetune.py --config my.yaml
        orbit train --skip-gate -- uv run scripts/train.py my_config
        orbit train --dry-run --gate "act" -- lerobot-train --dataset.repo_id=./data --policy.type=act
    """
    from orbit.hooks import run_training_pipeline

    if not training_command:
        from orbit.utils.display import console

        console.print(
            "\n[bold red]Error:[/bold red] No training command provided.\n"
            "  Usage: orbit train [--gate '<policy> [--min-grade <grade>]'] -- <your training command>\n"
            "  Example: orbit train --gate 'act --min-grade B' -- lerobot-train --dataset.repo_id=./data --policy.type=act"
        )
        sys.exit(1)

    run_training_pipeline(
        gate_args=gate,
        training_command=list(training_command),
        auto_clean=auto_clean,
        skip_gate=skip_gate,
        dry_run=dry_run,
    )


# ---------------------------------------------------------------------------
# Phase 4: orbit init
# ---------------------------------------------------------------------------


@main.command("init")
@click.option("--dataset", "-d", required=True, help="Dataset path or HuggingFace repo ID.")
@click.option("--policy", "-p", default="act", help="Target policy.")
@click.option("--output", "-o", default=".", help="Output directory for generated files.")
@click.option("--min-grade", default="C", type=click.Choice(["A", "B", "C", "D"]),
              help="Minimum passing grade for quality gate.")
def init_cmd(dataset: str, policy: str, output: str, min_grade: str):
    """Initialize a training project with Makefile, scripts, and CI config.

    Generates ready-to-use project files for a robot training pipeline
    including a Makefile, train.sh script, and GitHub Actions workflow.

    \b
    Examples:
        orbit init --dataset lerobot/my-dataset --policy act
        orbit init -d lerobot/my-dataset -p act -o ./my-project/
    """
    from orbit.training.templates import init_project
    from orbit.utils.display import console

    created = init_project(
        dataset=dataset,
        policy=policy,
        output_dir=output,
        min_grade=min_grade,
    )

    console.print(f"\n[bold green]Project initialized![/bold green]")
    console.print(f"\n  Created files:")
    for f in created:
        console.print(f"    {f}")
    console.print(f"\n  Next steps:")
    console.print(f"    make check    # Run quality gate")
    console.print(f"    make train    # Gate + train")
    console.print(f"    make all      # Gate + train + verify")
    console.print()


# ---------------------------------------------------------------------------
# Phase 3: orbit assist
# ---------------------------------------------------------------------------


@main.command()
@click.option(
    "--backend",
    type=click.Choice(["gemini", "claude", "auto"]),
    default="auto",
    help="AI provider: gemini (default), claude (optional, requires pip install orbit-robotics[assist-claude]).",
)
@click.option("--no-scan", is_flag=True, help="Skip initial directory scan.")
@click.option(
    "--task",
    "-t",
    default=None,
    help="Run a single task non-interactively (exit after completion).",
)
def assist(backend, no_scan, task):
    """Interactive AI assistant for robotics data and training.

    Scans your workspace, understands your data, and helps you go from
    raw recordings to trained policies. Uses ORBIT's tools autonomously —
    it doesn't just suggest commands, it runs them.

    \b
    Examples:
        orbit assist
        orbit assist --task "analyze my dataset and tell me if it's ready for ACT"
        orbit assist --backend claude
        orbit assist --no-scan -t "suggest training commands for lerobot/pusht"
    """
    from orbit.assist import run_assist

    run_assist(backend=backend, no_scan=no_scan, task=task)


# ---------------------------------------------------------------------------
# Phase 5: Community, Polish & Ecosystem
# ---------------------------------------------------------------------------


@main.command()
@click.option("--json", "as_json", is_flag=True, help="Output as JSON.")
def doctor(as_json: bool):
    """Check environment health and dependencies.

    \b
    Examples:
        orbit doctor
        orbit doctor --json
    """
    from orbit.doctor import print_doctor_report, run_doctor

    checks = run_doctor()
    if as_json:
        import json as json_mod

        results = [
            {"category": cat, "status": st, "message": msg, "fix": fix}
            for cat, st, msg, fix in checks
        ]
        click.echo(json_mod.dumps(results, indent=2))
    else:
        print_doctor_report(checks)


@main.command()
@click.option("--robot", default=None, help="Filter by robot type (e.g. so100, aloha, panda).")
@click.option("--task", default=None, help="Filter by task description (e.g. 'pick and place').")
@click.option("--min-grade", default=None, type=click.Choice(["A", "B", "C", "D"]), help="Minimum ORBIT grade.")
@click.option("--policy", default=None, help="Filter by policy type.")
@click.option("--min-episodes", default=None, type=int, help="Minimum number of episodes.")
@click.option("--limit", default=20, type=int, help="Maximum results to show.")
def explore(robot, task, min_grade, policy, min_episodes, limit):
    """Browse and discover LeRobot datasets on HuggingFace.

    \b
    Examples:
        orbit explore --robot so100 --task "pick and place"
        orbit explore --min-grade B --policy act --min-episodes 50
        orbit explore --limit 10
    """
    from orbit.explorer import print_explore_results, search_datasets

    results = search_datasets(
        robot=robot,
        task=task,
        min_grade=min_grade,
        policy=policy,
        min_episodes=min_episodes,
        limit=limit,
    )
    print_explore_results(results)


@main.command()
@click.argument("dataset_repo")
@click.option("--push", is_flag=True, help="Push badge to the HuggingFace dataset card.")
def badge(dataset_repo: str, push: bool):
    """Generate an ORBIT quality badge for a dataset.

    \b
    Examples:
        orbit badge lerobot/my-dataset
        orbit badge lerobot/my-dataset --push
    """
    from orbit.badge import run_badge

    run_badge(dataset_repo, push=push)


@main.command()
@click.argument("dataset_repo", required=False, default=None)
@click.option("--json", "as_json", is_flag=True, help="Output as JSON.")
def quickstart(dataset_repo: str | None, as_json: bool):
    """Get started in 30 seconds. Analyzes a dataset and shows what to do next.

    If no dataset is given, analyzes a public reference dataset so you can see
    what ORBIT does before pointing it at your own data.

    \b
    Examples:
        orbit quickstart                           # Try with a public dataset
        orbit quickstart lerobot/my-dataset        # Analyze your own data
    """
    from orbit.utils.display import console

    if dataset_repo is None:
        console.print()
        console.print("[bold]Welcome to ORBIT![/bold]")
        console.print()
        console.print("Let's try a quick analysis on a public dataset so you can see what ORBIT does.")
        console.print()
        dataset_repo = "lerobot/pusht"
        console.print(f"[dim]Using demo dataset: {dataset_repo}[/dim]")
        console.print()

    # Run a focused analysis with sensible defaults
    from click import Context

    ctx = Context(analyze, info_name="analyze")
    ctx.invoke(
        analyze,
        dataset_repo=dataset_repo,
        episodes=50,
        no_cache=False,
        as_json=as_json,
        skip_embeddings=True,
        skip_coverage=False,
        device=None,
        max_frames=4,
        skip_signal_diagnostics=False,
        vlm=False,
        policy="auto",
        full=False,
        ai=False,
        deep=False,
        proxy=False,
        skip_proxy=True,
        source_format=None,
        action_topic=None,
        state_topic=None,
        action_key=None,
        state_key=None,
        episode_gap=None,
        fps=None,
        detail=False,
        push=False,
    )

    if not as_json:
        console.print("[bold]WHAT'S NEXT?[/bold]")
        console.print("\u2500" * 40)
        console.print(f"  orbit analyze {dataset_repo} --deep    # AI-powered deep analysis")
        console.print(f"  orbit clean {dataset_repo}             # Remove bad episodes")
        console.print(f"  orbit suggest {dataset_repo}           # Get a training command")
        console.print(f"  orbit assist                           # AI copilot for your workflow")
        console.print()
        console.print("[dim]Full docs: orbit --help | Each command: orbit <command> --help[/dim]")
        console.print()


@main.command()
@click.argument("dataset_repo")
@click.option("--output", "-o", default=None, help="Output repo/path for cleaned dataset.")
@click.option("--policy", "-p", default=None, help="Target policy (act, diffusion, bc, smolvla).")
@click.option("--dry-run", is_flag=True, help="Show what would change without executing.")
@click.option("--json", "as_json", is_flag=True, help="Output as JSON.")
def fix(dataset_repo: str, output: str | None, policy: str | None, dry_run: bool, as_json: bool):
    """Smart auto-fix: analyze, clean bad episodes, and suggest training config.

    Runs the full pipeline in one command: identifies problems, removes bad
    episodes, and generates a ready-to-run training command.

    \b
    Examples:
        orbit fix lerobot/my-dataset                    # Full auto-fix pipeline
        orbit fix lerobot/my-dataset --dry-run          # Preview without changes
        orbit fix lerobot/my-dataset -p act -o user/my-dataset-clean
    """
    from orbit.analyzer.dataset_loader import load_dataset_info, load_episode_data
    from orbit.utils.display import console

    console.print()
    console.print("[bold]ORBIT Fix Pipeline[/bold]")
    console.print("\u2500" * 40)
    console.print()

    # Step 1: Quick analysis
    console.print("[bold]Step 1/3:[/bold] Analyzing dataset...")
    try:
        info = load_dataset_info(dataset_repo)
    except (ValueError, ConnectionError) as e:
        console.print(f"\n[bold red]Error:[/bold red] {e}")
        sys.exit(1)

    console.print(f"  Found {info.num_episodes} episodes, {info.robot_type or 'unknown'} robot")

    # Step 2: Clean
    console.print()
    console.print("[bold]Step 2/3:[/bold] Cleaning bad episodes...")

    try:
        episode_data = load_episode_data(dataset_repo, max_episodes=None)
    except (ValueError, ConnectionError) as e:
        console.print(f"\n[bold red]Error:[/bold red] Could not load episodes: {e}")
        sys.exit(1)

    if episode_data is None or episode_data.empty:
        console.print("[bold red]Error:[/bold red] No episode data loaded.")
        sys.exit(1)

    from orbit.cleaner import build_cleaning_plan, execute_cleaning, print_dry_run

    plan = build_cleaning_plan(
        dataset_repo=dataset_repo,
        episode_data=episode_data,
        info=info,
        aggressive=False,
        keep_dead_joints=False,
    )

    n_remove = len(plan.episodes_to_remove)
    n_review = len(plan.episodes_to_review)

    if n_remove == 0 and n_review == 0:
        console.print("  [green]\u2713[/green] No bad episodes found! Dataset is clean.")
    else:
        console.print(f"  Found {n_remove} episodes to remove, {n_review} to review")
        if dry_run:
            print_dry_run(plan)
        else:
            result = execute_cleaning(plan, output=output)
            if result.episodes_removed > 0:
                console.print(
                    f"  [green]\u2713[/green] Removed {result.episodes_removed} episodes"
                )
                if result.output_path:
                    console.print(f"  [dim]Cleaned dataset: {result.output_path}[/dim]")

    # Step 3: Suggest training command
    console.print()
    console.print("[bold]Step 3/3:[/bold] Generating training command...")
    final_dataset = dataset_repo
    if not dry_run and n_remove > 0 and hasattr(result, "output_path") and result.output_path:
        final_dataset = result.output_path

    # Invoke suggest
    from click import Context as ClickContext

    ctx = ClickContext(suggest, info_name="suggest")
    ctx.invoke(
        suggest,
        dataset=final_dataset,
        policy=policy,
        gpu_memory=None,
        framework="lerobot",
        as_json=as_json,
    )

    if not as_json:
        console.print()
        console.print("[bold green]\u2713 Fix pipeline complete![/bold green]")
        if dry_run:
            console.print("[dim]Run without --dry-run to apply changes.[/dim]")
        console.print()


@main.command("config")
@click.option("--show", is_flag=True, help="Show current configuration.")
@click.option("--init", "do_init", is_flag=True, help="Create a default config file.")
@click.option("--set", "set_key", nargs=2, multiple=True, metavar="KEY VALUE", help="Set a config value.")
def config_cmd(show: bool, do_init: bool, set_key: tuple):
    """Manage ORBIT configuration (~/.orbit/config.yaml).

    \b
    Examples:
        orbit config --show
        orbit config --init
        orbit config --set default_policy act --set gpu_memory 24
    """
    from orbit.config import CONFIG_FILE, OrbitConfig, get_config
    from orbit.utils.display import console

    if do_init:
        if CONFIG_FILE.exists():
            console.print(f"[yellow]Config already exists:[/yellow] {CONFIG_FILE}")
        else:
            cfg = OrbitConfig()
            cfg.save()
            console.print(f"[green]\u2713[/green] Created default config at {CONFIG_FILE}")
        return

    if set_key:
        from orbit.config import validate_config_value

        cfg = get_config()
        any_errors = False
        for key, value in set_key:
            if hasattr(cfg, key):
                # Coerce types
                current = getattr(cfg, key)
                if isinstance(current, int):
                    value = int(value)
                elif isinstance(current, float):
                    value = float(value)

                # Validate before setting
                error = validate_config_value(key, value)
                if error:
                    console.print(f"  [red]\u2717[/red] {error}")
                    any_errors = True
                    continue

                setattr(cfg, key, value)
                console.print(f"  [green]\u2713[/green] {key} = {value}")
            else:
                console.print(f"  [red]\u2717[/red] Unknown key: {key}")
                any_errors = True
        if not any_errors:
            cfg.save()
            console.print(f"\n[dim]Saved to {CONFIG_FILE}[/dim]")
        else:
            console.print(f"\n[yellow]Some values were invalid. Fix errors and retry.[/yellow]")
        return

    # Default: show config
    cfg = get_config()
    console.print(f"\n[bold]ORBIT Configuration[/bold] ({CONFIG_FILE})")
    if not CONFIG_FILE.exists():
        console.print("[dim]No config file found. Run: orbit config --init[/dim]")
    else:
        for key in [
            "default_policy", "default_min_grade", "gpu_memory",
            "default_framework", "default_format", "default_device", "plugins_dir",
        ]:
            val = getattr(cfg, key, None)
            if val is not None:
                console.print(f"  {key}: [bold]{val}[/bold]")
            else:
                console.print(f"  {key}: [dim]not set[/dim]")
        if cfg.coach != {"target_episodes": 50, "poll_interval": 2.0}:
            console.print(f"  coach: [bold]{cfg.coach}[/bold]")
    console.print()


# Shell completions support
@main.command("install-completion")
@click.argument("shell", type=click.Choice(["bash", "zsh", "fish"]))
def install_completion(shell: str):
    """Install shell tab-completion for ORBIT.

    \b
    Examples:
        orbit install-completion bash
        orbit install-completion zsh
        orbit install-completion fish
    """
    import subprocess

    from orbit.utils.display import console

    env_var = "_ORBIT_COMPLETE"
    if shell == "bash":
        script = f'eval "$({env_var}=bash_source orbit)"'
        rc_file = "~/.bashrc"
    elif shell == "zsh":
        script = f'eval "$({env_var}=zsh_source orbit)"'
        rc_file = "~/.zshrc"
    elif shell == "fish":
        script = f'{env_var}=fish_source orbit | source'
        rc_file = "~/.config/fish/config.fish"

    console.print(f"\n[bold]Shell completion for {shell}[/bold]")
    console.print(f"\nAdd this to your [bold]{rc_file}[/bold]:\n")
    console.print(f"  [cyan]{script}[/cyan]")
    console.print(f"\nThen restart your shell or run: [bold]source {rc_file}[/bold]")
    console.print()


# ---------------------------------------------------------------------------
# orbit compare
# ---------------------------------------------------------------------------


@main.command()
@click.argument("dataset_a")
@click.argument("dataset_b")
@click.option(
    "--policy",
    default="auto",
    type=click.Choice(["auto", "act", "diffusion_policy", "dp3", "bc", "bc_rnn", "smolvla"]),
    help="Policy type for analysis.",
)
@click.option("--json", "as_json", is_flag=True, help="Output as JSON.")
def compare(dataset_a: str, dataset_b: str, policy: str, as_json: bool):
    """Compare two datasets side-by-side.

    Shows what improved, what degraded, and what's new between two
    versions of a dataset or two different datasets.

    \b
    Examples:
        orbit compare lerobot/pusht lerobot/xarm_lift_medium
        orbit compare ./data-v1/ ./data-v2/ --policy act
        orbit compare lerobot/my-dataset-raw lerobot/my-dataset-clean
    """
    from orbit.compare import compare_datasets, print_compare
    from orbit.utils.display import console

    if not as_json:
        console.print(f"\n[dim]Analyzing both datasets (this may take a minute)...[/dim]")

    result = compare_datasets(dataset_a, dataset_b, policy=policy)

    if as_json:
        import dataclasses
        click.echo(json.dumps(dataclasses.asdict(result), indent=2, default=_json_default))
    else:
        print_compare(result)


# ---------------------------------------------------------------------------
# orbit leaderboard
# ---------------------------------------------------------------------------


@main.command()
@click.option("--grade", default=None, type=click.Choice(["A", "B", "C", "D", "F"]), help="Filter by grade.")
@click.option("--robot", default=None, help="Filter by robot type.")
@click.option("--limit", default=20, type=int, help="Number of results to show.")
@click.option("--json", "as_json", is_flag=True, help="Output as JSON.")
def leaderboard(grade: str | None, robot: str | None, limit: int, as_json: bool):
    """Browse the quality leaderboard of scored robotics datasets.

    Shows the top-rated datasets from the community leaderboard,
    ranked by ORBIT quality score.

    \b
    Examples:
        orbit leaderboard                    # Top 20 datasets
        orbit leaderboard --grade A          # Only grade-A datasets
        orbit leaderboard --robot aloha      # Filter by robot
        orbit leaderboard --limit 50         # Show more results
    """
    from pathlib import Path

    from orbit.utils.display import console

    # Look for scores in multiple locations
    scores_path = None
    candidates = [
        Path(__file__).parent.parent.parent / "leaderboard" / "scores.json",
        Path.home() / ".cache" / "orbit" / "leaderboard" / "scores.json",
    ]
    for p in candidates:
        if p.exists():
            scores_path = p
            break

    if scores_path is None:
        console.print("[bold yellow]No leaderboard data found.[/bold yellow]")
        console.print("Run the scoring script first:")
        console.print("  python scripts/score_all_lerobot.py")
        console.print("\nOr check the online leaderboard at:")
        console.print("  https://pypi.org/project/orbit-robotics/")
        sys.exit(1)

    data = json.loads(scores_path.read_text())
    scored = [d for d in data if d.get("scored")]

    # Apply filters
    if grade:
        scored = [d for d in scored if d.get("grade") == grade]
    if robot:
        robot_lower = robot.lower()
        scored = [d for d in scored if robot_lower in (d.get("robot_type", "").lower())]

    # Sort by score
    scored.sort(key=lambda x: x.get("score", 0), reverse=True)
    scored = scored[:limit]

    if not scored:
        if as_json:
            click.echo("[]")
        else:
            console.print("[yellow]No datasets match your filters.[/yellow]")
        return

    if as_json:
        click.echo(json.dumps(scored, indent=2))
        return

    # Pretty display
    grade_emoji = {"A": "🟢", "B": "🟢", "C": "🟡", "D": "🟠", "F": "🔴"}
    total_in_db = len([d for d in data if d.get("scored")])

    console.print(f"\n[bold]ORBIT Leaderboard[/bold] — {len(scored)} of {total_in_db} datasets")
    filters = []
    if grade:
        filters.append(f"grade={grade}")
    if robot:
        filters.append(f"robot={robot}")
    if filters:
        console.print(f"[dim]Filters: {', '.join(filters)}[/dim]")
    console.print("─" * 80)

    for i, d in enumerate(scored, 1):
        g = d.get("grade", "?")
        emoji = grade_emoji.get(g, "⚪")
        repo = d.get("repo_id", "?")
        score = d.get("score", 0)
        eps = d.get("num_episodes", "?")
        pol = (d.get("policy_recommended") or "auto").upper()
        robot_name = d.get("robot_type", "—")
        if robot_name == "unknown":
            robot_name = "—"
        div = d.get("action_divergence")
        div_str = f"div={div:.2f}" if div is not None else ""

        console.print(
            f"  {i:>3}. {emoji} {g} ({score:>3}/100)  "
            f"{repo:<45} {eps:>5} eps  {pol:<8} {robot_name:<12} {div_str}"
        )

    console.print()
    console.print("[dim]Install: pip install orbit-robotics | Score your data: orbit analyze <dataset>[/dim]")
    console.print()
    console.print()