"""Training failure diagnosis.

Reads LeRobot training logs, identifies failure patterns via rules and LLM,
connects failures to data issues, and tells the user exactly what to change.
"""

from __future__ import annotations

import json
import logging
import math
import re
from pathlib import Path

logger = logging.getLogger(__name__)


# ---------------------------------------------------------------------------
# Log parsing
# ---------------------------------------------------------------------------


def _parse_lerobot_logs(training_dir: Path) -> list[dict]:
    """Parse LeRobot-style training log lines.

    Expected format (any of these):
      step:200 ... loss:0.659 grdn:6.386 lr:2.0e-05
      Step 200, loss=0.659, grad_norm=6.386
    """
    entries: list[dict] = []
    pattern = re.compile(
        r"step[:\s=]+([\d,KkMm]+)"
        r".*?loss[:\s=]+([\d.]+(?:e[+-]?\d+)?|nan|inf)"
        r"(?:.*?gr[ad]n[:\s=]+([\d.]+(?:e[+-]?\d+)?))?"
        r"(?:.*?lr[:\s=]+([\d.]+(?:e[+-]?\d+)?))?"
    )

    log_files: list[Path] = []
    for ext in ("*.log", "*.txt", "*.out"):
        log_files.extend(training_dir.rglob(ext))
    # Also check direct children with no extension (stdout captures)
    for f in training_dir.iterdir():
        if f.is_file() and f.suffix in ("", ".log", ".txt", ".out") and f not in log_files:
            log_files.append(f)

    for log_file in log_files:
        try:
            text = log_file.read_text(errors="ignore")
        except OSError:
            continue
        for line in text.splitlines():
            m = pattern.search(line)
            if not m:
                continue
            step_str = m.group(1).replace(",", "")
            step_str = re.sub(r"[Kk]$", "000", step_str)
            step_str = re.sub(r"[Mm]$", "000000", step_str)
            try:
                step = int(re.sub(r"[^\d]", "", step_str))
            except ValueError:
                continue

            loss_str = m.group(2)
            loss = float("nan") if loss_str in ("nan", "inf") else float(loss_str)

            entry: dict = {"step": step, "loss": loss, "source": str(log_file.name)}
            if m.group(3):
                try:
                    entry["grad_norm"] = float(m.group(3))
                except ValueError:
                    pass
            if m.group(4):
                try:
                    entry["lr"] = float(m.group(4))
                except ValueError:
                    pass
            entries.append(entry)

    return entries


def _parse_wandb_logs(training_dir: Path) -> list[dict]:
    """Parse wandb run history if available."""
    entries: list[dict] = []
    wandb_dir = training_dir / "wandb"
    if not wandb_dir.exists():
        # Also check for wandb inside subdirectories
        for candidate in training_dir.rglob("wandb"):
            if candidate.is_dir():
                wandb_dir = candidate
                break
        else:
            return entries

    for history_file in wandb_dir.rglob("*history*.jsonl"):
        try:
            for line in history_file.read_text(errors="ignore").splitlines():
                try:
                    row = json.loads(line)
                except json.JSONDecodeError:
                    continue
                step = row.get("_step", row.get("global_step", row.get("step")))
                loss = row.get("loss", row.get("train/loss", row.get("train_loss")))
                if step is not None and loss is not None:
                    entry: dict = {"step": int(step), "loss": float(loss), "source": "wandb"}
                    grad = row.get("grad_norm", row.get("train/grad_norm"))
                    if grad is not None:
                        entry["grad_norm"] = float(grad)
                    entries.append(entry)
        except OSError:
            continue

    return entries


def _find_config(training_dir: Path) -> dict | None:
    """Find and parse the training config."""
    candidates = [
        training_dir / "config.json",
        training_dir / "config.yaml",
        training_dir / ".hydra" / "config.yaml",
        training_dir / ".hydra" / "overrides.yaml",
    ]
    for f in training_dir.rglob("config.json"):
        if f not in candidates:
            candidates.append(f)

    for f in candidates:
        if not f.exists():
            continue
        try:
            if f.suffix == ".json":
                return json.loads(f.read_text())
            elif f.suffix in (".yaml", ".yml"):
                # Minimal YAML parsing — avoid adding PyYAML dependency
                data: dict = {}
                for line in f.read_text().splitlines():
                    line = line.strip()
                    if ":" in line and not line.startswith("#"):
                        k, _, v = line.partition(":")
                        v = v.strip()
                        if v:
                            data[k.strip()] = v
                return data
        except (OSError, json.JSONDecodeError):
            continue
    return None


# ---------------------------------------------------------------------------
# Training dynamics analysis
# ---------------------------------------------------------------------------


def _analyze_dynamics(entries: list[dict]) -> dict:
    """Analyze training dynamics from parsed log entries."""
    if not entries:
        return {
            "status": "no_logs",
            "message": "No training log entries found.",
            "total_steps": 0,
            "loss_curve_summary": {},
            "findings_data": {},
        }

    entries = sorted(entries, key=lambda e: e["step"])
    losses = [e["loss"] for e in entries]
    steps = [e["step"] for e in entries]
    total_steps = max(steps) if steps else 0

    # Check for NaN/inf
    nan_count = sum(1 for l in losses if math.isnan(l) or math.isinf(l))
    if nan_count > 0:
        nan_step = next(e["step"] for e in entries if math.isnan(e["loss"]) or math.isinf(e["loss"]))
        return {
            "status": "nan_loss",
            "message": f"Loss exploded to NaN/inf at step {nan_step} ({nan_count} total NaN entries).",
            "total_steps": total_steps,
            "nan_step": nan_step,
            "loss_curve_summary": _loss_summary(entries),
            "findings_data": {"nan_step": nan_step, "nan_count": nan_count},
        }

    valid_losses = [l for l in losses if not math.isnan(l) and not math.isinf(l)]
    if not valid_losses:
        return {
            "status": "no_logs",
            "message": "No valid loss values found.",
            "total_steps": total_steps,
            "loss_curve_summary": {},
            "findings_data": {},
        }

    start_loss = valid_losses[0]
    end_loss = valid_losses[-1]
    min_loss = min(valid_losses)

    # Gradient norms
    grad_norms = [e["grad_norm"] for e in entries if "grad_norm" in e]
    max_grad = max(grad_norms) if grad_norms else 0
    mean_grad = sum(grad_norms) / len(grad_norms) if grad_norms else 0

    # Loss explosion: loss increases by >5x
    if end_loss > start_loss * 5 and end_loss > 1.0:
        return {
            "status": "loss_explosion",
            "message": f"Loss exploded from {start_loss:.4f} to {end_loss:.4f}.",
            "total_steps": total_steps,
            "loss_curve_summary": _loss_summary(entries),
            "findings_data": {"start_loss": start_loss, "end_loss": end_loss},
        }

    # Gradient instability: very high gradient norms
    if max_grad > 100:
        return {
            "status": "gradient_instability",
            "message": f"Gradient norms reached {max_grad:.1f} (mean: {mean_grad:.1f}). Training unstable.",
            "total_steps": total_steps,
            "loss_curve_summary": _loss_summary(entries),
            "findings_data": {"max_grad": max_grad, "mean_grad": mean_grad},
        }

    # Plateau: loss hasn't dropped significantly
    ratio = end_loss / max(start_loss, 1e-8)
    if ratio > 0.7 and total_steps > 10_000:
        return {
            "status": "plateau",
            "message": (
                f"Loss barely changed: {start_loss:.4f} → {end_loss:.4f} "
                f"after {total_steps} steps. Likely a data problem."
            ),
            "total_steps": total_steps,
            "loss_curve_summary": _loss_summary(entries),
            "findings_data": {"start_loss": start_loss, "end_loss": end_loss, "ratio": ratio},
        }

    # Converged but high
    if end_loss > 0.15:
        return {
            "status": "converged_high",
            "message": (
                f"Training converged but loss is high ({end_loss:.4f}). "
                f"Typical healthy final loss is 0.02-0.08."
            ),
            "total_steps": total_steps,
            "loss_curve_summary": _loss_summary(entries),
            "findings_data": {"start_loss": start_loss, "end_loss": end_loss},
        }

    # Partial convergence: still dropping
    n = len(valid_losses)
    last_quarter = valid_losses[3 * n // 4:]
    if last_quarter:
        last_q_mean = sum(last_quarter) / len(last_quarter)
        if last_q_mean > min_loss * 1.5 and end_loss > 0.05:
            return {
                "status": "partial_convergence",
                "message": f"Loss is still dropping ({end_loss:.4f}). Consider training longer.",
                "total_steps": total_steps,
                "loss_curve_summary": _loss_summary(entries),
                "findings_data": {"start_loss": start_loss, "end_loss": end_loss, "min_loss": min_loss},
            }

    # Good convergence
    return {
        "status": "converged_well",
        "message": f"Training converged well: {start_loss:.4f} → {end_loss:.4f}.",
        "total_steps": total_steps,
        "loss_curve_summary": _loss_summary(entries),
        "findings_data": {"start_loss": start_loss, "end_loss": end_loss, "min_loss": min_loss},
    }


def _loss_summary(entries: list[dict]) -> dict:
    """Summarize loss curve at quartile checkpoints."""
    valid = [e for e in entries if not math.isnan(e["loss"]) and not math.isinf(e["loss"])]
    if not valid:
        return {}
    losses = [e["loss"] for e in valid]
    n = len(losses)
    return {
        "start": losses[0],
        "q25": losses[n // 4] if n > 4 else losses[0],
        "q50": losses[n // 2] if n > 2 else losses[0],
        "q75": losses[3 * n // 4] if n > 4 else losses[-1],
        "end": losses[-1],
    }


# ---------------------------------------------------------------------------
# Rule-based diagnosis
# ---------------------------------------------------------------------------


def _rule_based_diagnosis(dynamics: dict, config: dict | None, dataset_id: str | None) -> list[dict]:
    """Generate rule-based findings from training dynamics."""
    findings: list[dict] = []
    status = dynamics["status"]
    data = dynamics.get("findings_data", {})

    if status == "nan_loss":
        findings.append({
            "severity": "critical",
            "finding": f"Loss exploded to NaN at step {data.get('nan_step', '?')}",
            "cause": "Usually learning rate too high, or dead joints causing gradient explosion in action space.",
            "fix": "Reduce learning rate to 5e-5. Check for dead joints with `orbit analyze`.",
        })

    if status == "loss_explosion":
        findings.append({
            "severity": "critical",
            "finding": f"Loss diverged: {data.get('start_loss', '?'):.4f} → {data.get('end_loss', '?'):.4f}",
            "cause": "Learning rate too high, or severe data quality issues (conflicting demonstrations).",
            "fix": "Halve the learning rate. Run `orbit analyze` to check action divergence.",
        })

    if status == "gradient_instability":
        findings.append({
            "severity": "high",
            "finding": f"Gradient norms spiking to {data.get('max_grad', '?'):.1f}",
            "cause": "Action space has very different scales across dimensions, or dead joints.",
            "fix": "Enable gradient clipping (max_grad_norm=1.0). Check for dead joints with `orbit analyze`.",
        })

    if status == "plateau":
        findings.append({
            "severity": "high",
            "finding": f"Loss plateau at {data.get('end_loss', '?'):.4f} after {dynamics['total_steps']} steps",
            "cause": "Almost always a DATA problem: inconsistent demonstrations, multimodal behavior with a unimodal policy, or too few episodes.",
            "fix": "Run `orbit analyze --deep` on your dataset. If consistency < 0.65, collect cleaner data. If divergence > 0.12, switch to Diffusion Policy.",
        })

    if status == "converged_high":
        findings.append({
            "severity": "high",
            "finding": f"Converged but loss is high ({data.get('end_loss', '?'):.4f})",
            "cause": "Data quality issues preventing full convergence. Common causes: inconsistent demos, too few episodes, or multimodal demonstrations with ACT.",
            "fix": "Run `orbit clean` to remove bad episodes, then retrain. If using ACT with divergence > 0.10, switch to Diffusion Policy.",
        })

    # Config-based checks
    if config:
        batch_size = config.get("batch_size", config.get("training", {}).get("batch_size"))
        if batch_size is not None:
            try:
                bs = int(batch_size)
                if bs == 1:
                    findings.append({
                        "severity": "critical",
                        "finding": "batch_size=1 detected",
                        "cause": "batch_size=1 causes silent training failure — gradients oscillate, loss looks normal but robot won't move.",
                        "fix": "Set batch_size >= 4 (8 preferred). Use `orbit suggest` to get optimal batch size for your GPU.",
                    })
                elif bs < 4:
                    findings.append({
                        "severity": "high",
                        "finding": f"batch_size={bs} is very low",
                        "cause": "Small batch sizes cause unstable gradients and worse convergence.",
                        "fix": f"Increase to at least 8 if GPU memory allows. Run `orbit suggest --gpu-memory <GB>`.",
                    })
            except (ValueError, TypeError):
                pass

    if status == "converged_well":
        findings.append({
            "severity": "info",
            "finding": "Training converged normally",
            "cause": "If the robot still doesn't work at eval, the issue is likely deployment, not training.",
            "fix": "Check action normalization stats match between training and eval. Verify camera preprocessing is identical.",
        })

    return findings


# ---------------------------------------------------------------------------
# LLM diagnosis
# ---------------------------------------------------------------------------


def _llm_diagnosis(
    dynamics: dict,
    config: dict | None,
    findings: list[dict],
    dataset_id: str | None,
    api_key: str,
    analysis_results: dict | None = None,
) -> str | None:
    """Get LLM-powered diagnosis of training failure.

    When analysis_results is provided (from orbit analyze --json), the LLM
    can trace training failures back to specific data problems.
    """
    from orbit.analyzer.deep_analysis import TRAINING_EXPERTISE
    from orbit.llm_utils import call_gemini

    findings_text = "\n".join(
        f"- [{f['severity'].upper()}] {f['finding']}: {f['cause']}"
        for f in findings
    )
    lcs = dynamics.get("loss_curve_summary", {})
    loss_text = ""
    if lcs.get("start") is not None:
        loss_text = (
            f"Loss curve: start={lcs['start']:.4f}, "
            f"25%={lcs.get('q25', '?')}, 50%={lcs.get('q50', '?')}, "
            f"75%={lcs.get('q75', '?')}, end={lcs['end']:.4f}"
        )

    config_text = json.dumps(config, indent=2, default=str)[:2000] if config else "No config found"

    # Build dataset analysis context if available
    dataset_context = ""
    if analysis_results:
        parts = []
        if "episode_count" in analysis_results:
            parts.append(f"- {analysis_results['episode_count']} episodes")
        if "consistency" in analysis_results:
            c = analysis_results["consistency"]
            if isinstance(c, dict):
                parts.append(f"- Consistency: {c.get('overall', 'unknown')}")
                for k, v in c.items():
                    if k != "overall" and isinstance(v, (int, float)):
                        parts.append(f"  - {k}: {v:.3f}")
            else:
                parts.append(f"- Consistency: {c}")
        if "signal_health" in analysis_results:
            sh = analysis_results["signal_health"]
            if isinstance(sh, dict):
                dead = sh.get("dead_dimensions", [])
                if dead:
                    parts.append(f"- Dead joints: {dead}")
                clip = sh.get("clipping_joints", [])
                if clip:
                    parts.append(f"- Clipping joints: {clip}")
        if "action_divergence" in analysis_results:
            parts.append(f"- Action divergence: {analysis_results['action_divergence']}")
        if "outlier_episodes" in analysis_results:
            parts.append(f"- Outlier episodes: {analysis_results['outlier_episodes']}")
        if "readiness" in analysis_results:
            r = analysis_results["readiness"]
            if isinstance(r, dict):
                parts.append(f"- Readiness grade: {r.get('grade', '?')} (score: {r.get('score', '?')})")
        if parts:
            dataset_context = "\n\nDATASET ANALYSIS:\n" + "\n".join(parts)

    prompt = f"""{TRAINING_EXPERTISE}

---

A user's training run has finished. Here is the diagnosis:

Status: {dynamics['status']}
Message: {dynamics['message']}
Total steps: {dynamics['total_steps']}
Dataset: {dataset_id or 'unknown'}
{loss_text}

Config:
{config_text}

Rule-based findings:
{findings_text}
{dataset_context}

---

Give a CONCISE diagnosis (under 200 words). Structure:

**Root cause**: One sentence. What went wrong and why.

**The fix**: 2-3 bullet points. Specific actions with numbers/commands.
Include `orbit` commands where relevant.
{"Reference specific dataset metrics above when diagnosing data issues." if dataset_context else ""}

**If the fix doesn't work**: One sentence fallback.

Be direct. No hedging."""

    return call_gemini(prompt, api_key, max_tokens=2048)


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


def debug_training(
    training_dir: str,
    dataset_id: str | None = None,
    use_ai: bool = False,
    analysis_results: dict | None = None,
) -> None:
    """Diagnose a training run from its output directory.

    Args:
        training_dir: Path to training output directory.
        dataset_id: Optional dataset repo ID (auto-detected from config).
        use_ai: If True, always attempt AI diagnosis even without dataset context.
        analysis_results: Optional dict from ``orbit analyze --json`` to give
            the LLM data-aware reasoning for tracing failures to data issues.
    """
    from rich.markdown import Markdown
    from rich.table import Table

    from orbit.llm_utils import load_api_key
    from orbit.utils.display import console

    training_path = Path(training_dir)

    if not training_path.exists():
        console.print(f"[red]Training directory not found: {training_path}[/red]")
        return

    console.print(f"[dim]Scanning {training_path}...[/dim]")

    # Find config
    config = _find_config(training_path)

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

    if not all_entries:
        console.print("[red]No training logs found.[/red]")
        console.print("Expected: .log/.txt files with loss entries, or wandb/ directory")
        console.print(f"Searched: {training_path}")
        return

    console.print(
        f"[dim]Found {len(all_entries)} log entries "
        f"({len(log_entries)} from logs, {len(wandb_entries)} from wandb)[/dim]"
    )

    # Auto-detect dataset from config
    if dataset_id is None and config:
        dataset_id = (
            config.get("dataset", {}).get("repo_id")
            if isinstance(config.get("dataset"), dict)
            else config.get("dataset.repo_id")
        )
    if dataset_id is None:
        for f in training_path.rglob("*.json"):
            try:
                d = json.loads(f.read_text())
                if isinstance(d.get("dataset"), dict) and "repo_id" in d["dataset"]:
                    dataset_id = d["dataset"]["repo_id"]
                    break
            except (OSError, json.JSONDecodeError):
                continue

    # Analyze dynamics
    dynamics = _analyze_dynamics(all_entries)

    # Header
    policy_type = "unknown"
    if config:
        if isinstance(config.get("policy"), dict):
            policy_type = config["policy"].get("type", "unknown")
        elif "policy.type" in config:
            policy_type = config["policy.type"]

    console.print()
    console.print(f"\u256d{'\u2500' * 70}\u256e")
    console.print(
        f"\u2502  Training Diagnosis: {str(training_path)[:48]:<48}  \u2502"
    )
    console.print(
        f"\u2502  Dataset: {str(dataset_id or 'unknown')[:23]:<23}  "
        f"Policy: {policy_type:<10}  Steps: {str(dynamics['total_steps']):<8}  \u2502"
    )
    console.print(f"\u2570{'\u2500' * 70}\u256f")
    console.print()

    # Training outcome
    status_colors = {
        "nan_loss": "red",
        "loss_explosion": "red",
        "gradient_instability": "red",
        "converged_high": "yellow",
        "plateau": "yellow",
        "partial_convergence": "yellow",
        "converged_well": "green",
        "no_logs": "red",
    }
    color = status_colors.get(dynamics["status"], "white")
    console.print(
        f"[bold]Training Outcome:[/bold] [{color}]{dynamics['status'].upper().replace('_', ' ')}[/{color}]"
    )
    console.print(f"  {dynamics['message']}")
    console.print()

    # Loss curve summary
    lcs = dynamics.get("loss_curve_summary", {})
    if lcs.get("start") is not None:
        console.print("[bold]Loss Curve:[/bold]")
        console.print(
            f"  Start: {lcs['start']:.4f}  \u2192  "
            f"25%: {lcs.get('q25', 0):.4f}  \u2192  "
            f"50%: {lcs.get('q50', 0):.4f}  \u2192  "
            f"75%: {lcs.get('q75', 0):.4f}  \u2192  "
            f"End: {lcs['end']:.4f}"
        )
        console.print()

    # Rule-based findings
    findings = _rule_based_diagnosis(dynamics, config, dataset_id)
    if findings:
        console.print("[bold]Findings:[/bold]")
        for f in findings:
            sev = f["severity"]
            icon = "\u274c" if sev in ("critical", "high") else "\u26a0" if sev != "info" else "\u2713"
            color = "red" if sev == "critical" else "yellow" if sev == "high" else "green" if sev == "info" else "white"
            console.print(f"  {icon} [{color}][{sev.upper()}] {f['finding']}[/{color}]")
            console.print(f"     Cause: {f['cause']}")
            console.print(f"     Fix: {f['fix']}")
            console.print()

    # Auto-analyze dataset when --ai flag is used with a dataset
    if use_ai and dataset_id and analysis_results is None:
        try:
            from orbit.analyzer.dataset_loader import load_dataset_info, load_episode_data

            console.print("[dim]Loading dataset for AI-powered analysis...[/dim]")
            ds_info = load_dataset_info(dataset_id)
            ep_data = load_episode_data(dataset_id, max_episodes=100)
            # Build a lightweight analysis summary
            analysis_results = {
                "episode_count": ds_info.num_episodes,
                "fps": ds_info.fps,
                "robot_type": ds_info.robot_type,
            }
            # Add consistency if computable
            if ep_data is not None and not ep_data.empty:
                from orbit.analyzer.coverage import _extract_array

                action_episodes = []
                for _, ep_df in ep_data.groupby("episode_index"):
                    if "action" in ep_df.columns:
                        arr = _extract_array(ep_df["action"])
                        if arr is not None and arr.size > 0:
                            action_episodes.append(arr)
                if len(action_episodes) >= 2:
                    try:
                        from orbit.analyzer.consistency_metrics import compute_consistency_metrics

                        cr = compute_consistency_metrics(action_episodes)
                        analysis_results["consistency"] = {
                            "overall": round(cr.overall_consistency, 3),
                            "path_length_cv": round(cr.path_length_cv, 3),
                            "cartesian_jerk_cv": round(cr.cartesian_jerk_cv, 3),
                        }
                    except Exception:
                        pass
                    try:
                        from orbit.analyzer.signal_diagnostics import compute_signal_diagnostics

                        sig = compute_signal_diagnostics(action_episodes)
                        analysis_results["signal_health"] = {
                            "dead_dimensions": sig.dead_dimensions,
                            "clipping_joints": getattr(sig, "clipping_joints", []),
                        }
                    except Exception:
                        pass
                    try:
                        from orbit.analyzer.action_divergence import compute_action_divergence

                        div_result = compute_action_divergence(action_episodes)
                        analysis_results["action_divergence"] = round(
                            div_result.overall_divergence, 3
                        )
                    except Exception:
                        pass
                    try:
                        from orbit.analyzer.outlier_detection import detect_outliers

                        outlier_result = detect_outliers(action_episodes)
                        analysis_results["outlier_episodes"] = outlier_result.outlier_indices
                    except Exception:
                        pass
        except Exception as e:
            console.print(f"[dim]Could not analyze dataset: {e}[/dim]")

    # LLM diagnosis
    api_key = load_api_key()
    if api_key and (use_ai or analysis_results is not None):
        with console.status("[dim]Running AI diagnosis...[/dim]"):
            llm_result = _llm_diagnosis(
                dynamics, config, findings, dataset_id, api_key,
                analysis_results=analysis_results,
            )
        if llm_result:
            console.print("[bold]AI DIAGNOSIS[/bold]")
            console.print("\u2500" * 40)
            console.print(Markdown(llm_result))
    elif api_key:
        with console.status("[dim]Running AI diagnosis...[/dim]"):
            llm_result = _llm_diagnosis(dynamics, config, findings, dataset_id, api_key)
        if llm_result:
            console.print("[bold]AI DIAGNOSIS[/bold]")
            console.print("\u2500" * 40)
            console.print(Markdown(llm_result))
    else:
        if use_ai:
            console.print("[yellow]Set GOOGLE_API_KEY for AI-powered diagnosis[/yellow]")
        else:
            console.print("[dim]Add --ai for AI-powered diagnosis (requires GOOGLE_API_KEY)[/dim]")

    # Next steps
    console.print()
    console.print("[bold]Next steps:[/bold]")
    if dataset_id:
        console.print(f"  1. orbit analyze {dataset_id} --deep    # Check data quality")
        console.print(f"  2. orbit clean {dataset_id}             # Fix data issues")
        console.print(f"  3. orbit suggest {dataset_id}           # Get corrected training command")
    if findings:
        console.print(f"  4. Fix the issues above and re-train")
    console.print()
