"""Dataset cleaning module — removes bad episodes and flags dead joints.

Runs signal diagnostics and episode ranking on every episode, then produces
a cleaned dataset with bad episodes removed and dead joints documented in
a sidecar report.
"""

from __future__ import annotations

import json
import logging
import subprocess
import sys
from dataclasses import dataclass, field
from datetime import datetime, timezone
from pathlib import Path

import numpy as np

logger = logging.getLogger(__name__)


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


@dataclass
class CleaningPlan:
    """What cleaning would do (for --dry-run and actual execution)."""

    original_dataset: str
    original_episodes: int
    episodes_to_remove: list[dict]  # [{"index": 3, "reason": "..."}, ...]
    episodes_to_keep: int
    dead_joints: list[dict]  # [{"index": 3, "name": "joint_3", "alive_ratio": 0.02}, ...]
    total_joints: int  # total action dimensions
    quality_before: float
    quality_after: float  # estimated
    cleaning_mode: str  # "standard" or "aggressive"


@dataclass
class CleaningResult:
    """Result of a cleaning operation."""

    plan: CleaningPlan
    output_path: str | None  # where the cleaned dataset ended up
    report_path: str | None  # path to the JSON sidecar report
    method: str  # "lerobot_api", "lerobot_cli", "report_only"


# ---------------------------------------------------------------------------
# Thresholds
# ---------------------------------------------------------------------------

DEAD_JOINT_ALIVE_RATIO = 0.05
DEAD_FRAME_THRESHOLD = 0.50  # >50% dead frames = remove
TEMPORAL_LAG_THRESHOLD = 3  # frames
AGGRESSIVE_CLIPPING_RATE = 0.15
AGGRESSIVE_MIN_LENGTH_RATIO = 0.20
AGGRESSIVE_MAX_LENGTH_RATIO = 3.00


# ---------------------------------------------------------------------------
# Main cleaning logic
# ---------------------------------------------------------------------------


def build_cleaning_plan(
    dataset_repo: str,
    episode_data,
    info,
    aggressive: bool = False,
    keep_dead_joints: bool = False,
    policy: str | None = None,
) -> CleaningPlan:
    """Analyze a dataset and build a plan for what to clean.

    Args:
        dataset_repo: HuggingFace repo ID or local path.
        episode_data: DataFrame with episode_index, action, state columns.
        info: DatasetInfo from dataset_loader.
        aggressive: If True, also remove 'review' episodes and apply stricter thresholds.
        keep_dead_joints: If True, don't flag dead joints.
        policy: Optional target policy name for policy-aware cleaning.

    Returns:
        CleaningPlan with all decisions made.
    """
    from orbit.analyzer.coverage import _extract_array, compute_coverage
    from orbit.analyzer.episode_ranker import compute_episode_ranking
    from orbit.analyzer.quality import compute_quality_score
    from orbit.analyzer.signal_diagnostics import compute_signal_diagnostics
    from orbit.analyzer.temporal_alignment import compute_temporal_alignment

    # Build per-episode action/state arrays
    action_episodes = []
    state_episodes = []

    if episode_data is None or episode_data.empty or "episode_index" not in episode_data.columns:
        return CleaningPlan(
            original_dataset=dataset_repo,
            original_episodes=info.num_episodes,
            episodes_to_remove=[],
            episodes_to_keep=info.num_episodes,
            dead_joints=[],
            total_joints=0,
            quality_before=0.0,
            quality_after=0.0,
            cleaning_mode="aggressive" if aggressive else "standard",
        )

    ep_groups = episode_data.groupby("episode_index")
    episode_indices = sorted(episode_data["episode_index"].unique())

    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)
            else:
                action_episodes.append(np.zeros((len(ep_df), 1)))
        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)

    if not action_episodes:
        return CleaningPlan(
            original_dataset=dataset_repo,
            original_episodes=info.num_episodes,
            episodes_to_remove=[],
            episodes_to_keep=info.num_episodes,
            dead_joints=[],
            total_joints=0,
            quality_before=0.0,
            quality_after=0.0,
            cleaning_mode="aggressive" if aggressive else "standard",
        )

    # Run diagnostics
    signal_diag = compute_signal_diagnostics(action_episodes)
    episode_ranking = compute_episode_ranking(
        action_episodes,
        state_episodes=state_episodes if state_episodes else None,
    )
    temporal_result = compute_temporal_alignment(
        action_episodes,
        state_episodes=state_episodes if state_episodes else None,
    )

    # Coverage and quality for before/after estimate
    coverage = compute_coverage(info, episode_data)
    quality_before = compute_quality_score(
        info, coverage, action_episodes=action_episodes, state_episodes=state_episodes or None,
    )

    # --- Decide what to remove ---
    remove_set: dict[int, str] = {}  # episode_index -> reason

    # Always: episodes classified as "remove"
    for ep_rank in episode_ranking.episodes:
        if ep_rank.verdict == "remove":
            remove_set[ep_rank.episode_index] = (
                f"Low quality (score: {ep_rank.composite_score:.2f})"
            )

    # Always: episodes with >50% dead frames
    for i, ep_actions in enumerate(action_episodes):
        if i < len(episode_indices):
            ep_idx = episode_indices[i]
        else:
            ep_idx = i
        if ep_actions.ndim == 1:
            ep_actions = ep_actions.reshape(-1, 1)
        dead_frame_ratio = float(np.mean(np.all(np.abs(ep_actions) < 0.01, axis=1)))
        if dead_frame_ratio > DEAD_FRAME_THRESHOLD and ep_idx not in remove_set:
            remove_set[ep_idx] = (
                f"Dead frames: {dead_frame_ratio:.0%} of frames have no motion"
            )

    # Always: episodes with temporal misalignment > 3 frames
    if abs(temporal_result.estimated_lag) > TEMPORAL_LAG_THRESHOLD:
        # This is a global issue, not per-episode — flag it as a warning but
        # don't remove episodes (the lag affects all episodes equally)
        pass

    # Aggressive mode
    if aggressive:
        # Also remove "review" episodes
        for ep_rank in episode_ranking.episodes:
            if ep_rank.verdict == "review" and ep_rank.episode_index not in remove_set:
                remove_set[ep_rank.episode_index] = (
                    f"Review quality (score: {ep_rank.composite_score:.2f}) — aggressive mode"
                )

        # Episodes with clipping rate > 15% on any joint
        for ed in signal_diag.episode_diagnostics:
            if ed.episode_index not in remove_set:
                # Check if this episode has clipping issues
                for issue in ed.issues:
                    if "clipping" in issue.lower():
                        remove_set[ed.episode_index] = f"Clipping detected — {issue}"
                        break

        # Episodes shorter than 20% of median or longer than 300% of median
        lengths = [len(ep) for ep in action_episodes]
        if lengths:
            median_len = float(np.median(lengths))
            for i, length in enumerate(lengths):
                ep_idx = episode_indices[i] if i < len(episode_indices) else i
                if ep_idx not in remove_set:
                    if length < AGGRESSIVE_MIN_LENGTH_RATIO * median_len:
                        remove_set[ep_idx] = (
                            f"Too short: {length} frames (median: {median_len:.0f}), likely aborted"
                        )
                    elif length > AGGRESSIVE_MAX_LENGTH_RATIO * median_len:
                        remove_set[ep_idx] = (
                            f"Too long: {length} frames (median: {median_len:.0f}), likely stuck/frozen"
                        )

    # --- Dead joints ---
    dead_joints = []
    if not keep_dead_joints:
        for jd in signal_diag.joint_diagnostics:
            alive_ratio = 1.0 - jd.zero_velocity_ratio
            if alive_ratio < DEAD_JOINT_ALIVE_RATIO:
                dead_joints.append({
                    "index": jd.joint_index,
                    "name": jd.joint_name or f"joint_{jd.joint_index}",
                    "alive_ratio": round(alive_ratio, 4),
                })

    # --- Quality estimate after cleaning ---
    kept_indices = set(range(len(action_episodes))) - {
        i for i, idx in enumerate(episode_indices) if idx in remove_set
    }
    if kept_indices:
        kept_actions = [action_episodes[i] for i in sorted(kept_indices)]
        kept_states = (
            [state_episodes[i] for i in sorted(kept_indices) if i < len(state_episodes)]
            if state_episodes else None
        )
        # Use a copy of info with the correct episode count so that
        # compute_quality_score's episode_count_score reflects kept episodes.
        import copy
        import dataclasses
        try:
            info_after = dataclasses.replace(
                info,
                num_episodes=info.num_episodes - len(remove_set),
            )
        except TypeError:
            # info may not be a dataclass (e.g. in tests with mocks)
            info_after = copy.copy(info)
            info_after.num_episodes = info.num_episodes - len(remove_set)
        quality_after = compute_quality_score(
            info_after, coverage,
            action_episodes=kept_actions,
            state_episodes=kept_states,
        )
        quality_after_score = quality_after.overall_score
    else:
        quality_after_score = quality_before.overall_score

    episodes_to_remove = [
        {"index": idx, "reason": reason}
        for idx, reason in sorted(remove_set.items())
    ]

    # Determine total action dimensions for display
    total_joints = action_episodes[0].shape[1] if action_episodes and action_episodes[0].ndim == 2 else 0

    return CleaningPlan(
        original_dataset=dataset_repo,
        original_episodes=info.num_episodes,
        episodes_to_remove=episodes_to_remove,
        episodes_to_keep=info.num_episodes - len(episodes_to_remove),
        dead_joints=dead_joints,
        total_joints=total_joints,
        quality_before=round(quality_before.overall_score * 100, 1),
        quality_after=round(quality_after_score * 100, 1),
        cleaning_mode="aggressive" if aggressive else "standard",
    )


def execute_cleaning(
    plan: CleaningPlan,
    output: str | None = None,
) -> CleaningResult:
    """Execute a cleaning plan, producing a cleaned dataset.

    Tries three methods in order:
    1. lerobot Python API (programmatic)
    2. lerobot-edit-dataset CLI (subprocess)
    3. Report-only fallback (JSON report + printed command)

    Args:
        plan: The cleaning plan to execute.
        output: Output repo ID or local path. Defaults to <dataset>-clean.

    Returns:
        CleaningResult with output paths and method used.
    """
    if not plan.episodes_to_remove:
        report_path = _write_report(plan, output)
        return CleaningResult(
            plan=plan,
            output_path=None,
            report_path=report_path,
            method="no_changes_needed",
        )

    remove_indices = [ep["index"] for ep in plan.episodes_to_remove]
    output = output or _default_output_name(plan.original_dataset)

    # Try Option A: lerobot Python API
    try:
        return _clean_via_api(plan, output, remove_indices)
    except Exception as e:
        logger.debug("lerobot API cleaning failed: %s", e)

    # Try Option B: lerobot-edit-dataset CLI
    try:
        return _clean_via_cli(plan, output, remove_indices)
    except Exception as e:
        logger.debug("lerobot CLI cleaning failed: %s", e)

    # Option C: Report only
    return _clean_report_only(plan, output, remove_indices)


def _default_output_name(dataset: str) -> str:
    if "/" in dataset:
        return dataset + "-clean"
    return str(Path(dataset).with_name(Path(dataset).name + "-clean"))


def _clean_via_api(plan, output, remove_indices) -> CleaningResult:
    """Use lerobot's dataset_tools.delete_episodes to remove episodes."""
    from lerobot.datasets.dataset_tools import delete_episodes
    from lerobot.datasets.lerobot_dataset import LeRobotDataset

    # Determine if dataset is local (path exists on disk) or remote (HF Hub)
    local_path = Path(plan.original_dataset)
    if local_path.exists() and local_path.is_dir():
        # Check for episodes metadata — required by delete_episodes.
        # Datasets created programmatically need finalize() called first.
        episodes_dir = local_path / "meta" / "episodes"
        if not episodes_dir.exists():
            raise RuntimeError(
                f"Missing {episodes_dir} — dataset may not be finalized. "
                "If you created this dataset programmatically, call "
                "dataset.finalize() before cleaning."
            )
        dataset = LeRobotDataset(
            repo_id=local_path.name,
            root=local_path,
        )
    else:
        dataset = LeRobotDataset(plan.original_dataset)

    delete_episodes(dataset, episode_indices=remove_indices, repo_id=output)

    report_path = _write_report(plan, output)
    return CleaningResult(
        plan=plan, output_path=output, report_path=report_path, method="lerobot_api",
    )


def _clean_via_cli(plan, output, remove_indices) -> CleaningResult:
    """Use lerobot-edit-dataset CLI to delete episodes."""
    import shutil

    cli_path = shutil.which("lerobot-edit-dataset")
    if cli_path is None:
        raise RuntimeError("lerobot-edit-dataset CLI not found on PATH")

    indices_str = json.dumps(remove_indices)
    cmd = [
        cli_path,
        "--repo_id", plan.original_dataset,
        "--new_repo_id", output,
        "--operation.type", "delete_episodes",
        "--operation.episode_indices", indices_str,
    ]
    result = subprocess.run(cmd, capture_output=True, text=True, timeout=300)
    if result.returncode != 0:
        raise RuntimeError(f"lerobot-edit-dataset failed: {result.stderr}")

    report_path = _write_report(plan, output)
    return CleaningResult(
        plan=plan, output_path=output, report_path=report_path, method="lerobot_cli",
    )


def _clean_report_only(plan, output, remove_indices) -> CleaningResult:
    """Fallback: write a report and print the manual command."""
    report_path = _write_report(plan, output)
    return CleaningResult(
        plan=plan, output_path=None, report_path=report_path, method="report_only",
    )


def _write_report(plan: CleaningPlan, output: str | None) -> str:
    """Write the cleaning report JSON sidecar."""
    report = {
        "original_dataset": plan.original_dataset,
        "original_episodes": plan.original_episodes,
        "removed_episodes": plan.episodes_to_remove,
        "kept_episodes": plan.episodes_to_keep,
        "dead_joints": plan.dead_joints,
        "quality_before": plan.quality_before,
        "quality_after": plan.quality_after,
        "cleaning_mode": plan.cleaning_mode,
        "timestamp": datetime.now(timezone.utc).isoformat(),
    }

    # Determine report path
    if output:
        safe_name = output.replace("/", "_")
    else:
        safe_name = plan.original_dataset.replace("/", "_")
    report_path = f"{safe_name}_cleaning_report.json"

    with open(report_path, "w") as f:
        json.dump(report, f, indent=2)

    return report_path


# ---------------------------------------------------------------------------
# Rich output helpers
# ---------------------------------------------------------------------------


def print_dry_run(plan: CleaningPlan) -> None:
    """Print a rich dry-run summary of what cleaning would do."""
    from rich.panel import Panel
    from rich.table import Table

    from orbit.utils.display import console

    console.print()

    # Summary table
    table = Table(title="Cleaning Summary (Dry Run)", show_header=False, border_style="cyan")
    table.add_column("Metric", style="bold")
    table.add_column("Value")
    table.add_row("Original episodes", str(plan.original_episodes))
    table.add_row(
        "Episodes to remove",
        f"{len(plan.episodes_to_remove)} ({len(plan.episodes_to_remove) * 100 // max(1, plan.original_episodes)}%)",
    )
    table.add_row("Episodes to keep", str(plan.episodes_to_keep))
    table.add_row("Dead joints", f"{len(plan.dead_joints)} of {plan.total_joints}" if plan.total_joints else f"{len(plan.dead_joints)}")
    table.add_row("Quality", f"{plan.quality_before:.0f} \u2192 {plan.quality_after:.0f}")
    table.add_row("Mode", plan.cleaning_mode)
    console.print(table)

    # Episode removal details
    if plan.episodes_to_remove:
        console.print("\n[bold]Episodes to remove:[/bold]")
        for ep in plan.episodes_to_remove:
            console.print(f"  #{ep['index']}  - {ep['reason']}")

    # Dead joints
    if plan.dead_joints:
        joint_indices = [str(j["index"]) for j in plan.dead_joints]
        console.print(
            f"\n[yellow]\u26a0 Dead joints: [{', '.join(joint_indices)}] \u2014 "
            "these joints never move.[/yellow]"
        )
        console.print(
            "    Your policy will waste capacity learning zero outputs."
        )
        console.print("    Consider remapping your action space.")

    # Print the manual command
    if plan.episodes_to_remove:
        indices = [ep["index"] for ep in plan.episodes_to_remove]
        output_name = _default_output_name(plan.original_dataset)
        console.print(f"\n[dim]To clean manually:[/dim]")
        console.print(
            f"  lerobot-edit-dataset --repo_id {plan.original_dataset} "
            f"--new_repo_id {output_name} "
            f"--operation.type delete_episodes "
            f"--operation.episode_indices '{json.dumps(indices)}'"
        )

    console.print()


def print_cleaning_result(result: CleaningResult) -> None:
    """Print the result of a cleaning operation."""
    from rich.table import Table

    from orbit.utils.display import console

    plan = result.plan
    console.print()

    if result.method == "no_changes_needed":
        console.print("[green]\u2713 No episodes need removal — dataset is clean![/green]")
        return

    if result.output_path:
        console.print(f"[green]\u2713 Cleaned dataset saved to {result.output_path}[/green]")
    elif result.method == "report_only":
        console.print(
            "[yellow]\u26a0 Could not modify dataset — "
            "neither the lerobot Python API nor lerobot-edit-dataset CLI were available.[/yellow]\n"
            "[yellow]  A cleaning report was generated. See the manual command below to clean the dataset yourself.[/yellow]"
        )
    else:
        console.print("[yellow]\u26a0 Could not modify dataset directly.[/yellow]")

    # Summary table
    table = Table(title="Cleaning Summary", show_header=False, border_style="cyan")
    table.add_column("Metric", style="bold")
    table.add_column("Value")
    table.add_row("Original episodes", str(plan.original_episodes))
    table.add_row(
        "Removed",
        f"{len(plan.episodes_to_remove)} ({len(plan.episodes_to_remove) * 100 // max(1, plan.original_episodes)}%)",
    )
    table.add_row("Kept", str(plan.episodes_to_keep))
    table.add_row("Dead joints", f"{len(plan.dead_joints)}")
    table.add_row("Quality", f"{plan.quality_before:.0f} \u2192 {plan.quality_after:.0f}")
    console.print(table)

    # Episode removal details
    if plan.episodes_to_remove:
        console.print("\n[bold]Removed episodes:[/bold]")
        for ep in plan.episodes_to_remove:
            console.print(f"  #{ep['index']}  - {ep['reason']}")

    # Dead joints warning
    if plan.dead_joints:
        joint_indices = [str(j["index"]) for j in plan.dead_joints]
        console.print(
            f"\n[yellow]\u26a0 Dead joints: [{', '.join(joint_indices)}] \u2014 "
            "these joints never move.[/yellow]"
        )
        console.print(
            "    Your policy will waste capacity learning zero outputs."
        )
        console.print("    Consider remapping your action space.")

    # Report and fallback command
    if result.report_path:
        console.print(f"\n[dim]Report saved to: {result.report_path}[/dim]")

    if result.method == "report_only" and plan.episodes_to_remove:
        indices = [ep["index"] for ep in plan.episodes_to_remove]
        output_name = result.plan.original_dataset.replace("/", "_") + "-clean"
        console.print(
            f"\n[bold]Run manually:[/bold]\n"
            f"  lerobot-edit-dataset --repo_id {plan.original_dataset} "
            f"--new_repo_id {output_name} "
            f"--operation.type delete_episodes "
            f"--operation.episode_indices '{json.dumps(indices)}'"
        )

    # Next step
    output_name = result.output_path or plan.original_dataset
    console.print(
        f"\n[bold]Next step:[/bold] orbit suggest {output_name}"
    )
    console.print()
