"""Real-time data collection coach.

Watches a dataset directory as new episodes are added, analyzes each
new episode against the existing dataset, and gives real-time feedback.
"""

from __future__ import annotations

import logging
import time
from dataclasses import dataclass, field
from pathlib import Path

import numpy as np

logger = logging.getLogger(__name__)


@dataclass
class CoachFeedback:
    """Feedback on a newly collected episode."""

    episode_index: int
    verdict: str  # "good", "review", "redo"
    consistency_score: float
    issues: list[str] = field(default_factory=list)
    suggestions: list[str] = field(default_factory=list)
    dataset_progress: str = ""


@dataclass
class CoachState:
    """Running state of the collection coach."""

    target_episodes: int
    policy: str
    current_count: int = 0
    total_good: int = 0
    total_review: int = 0
    total_redo: int = 0
    baseline_metrics: dict = field(default_factory=dict)
    recent_jerk_trend: list[float] = field(default_factory=list)


def _compute_episode_metrics(episode: np.ndarray) -> dict:
    """Compute per-episode metrics for coaching comparison."""
    metrics = {}

    if episode.ndim != 2 or len(episode) < 3:
        return metrics

    # Path length
    diffs = np.diff(episode, axis=0)
    metrics["path_length"] = float(np.sum(np.sqrt(np.sum(diffs**2, axis=1))))

    # Jerk (smoothness indicator)
    if len(episode) >= 4:
        vel = np.diff(episode, axis=0)
        acc = np.diff(vel, axis=0)
        jerk = np.diff(acc, axis=0)
        metrics["jerk"] = float(np.mean(np.sqrt(np.sum(jerk**2, axis=1))))
    else:
        metrics["jerk"] = 0.0

    # Duration (frame count as proxy)
    metrics["length"] = len(episode)

    # Action range utilization
    action_range = np.ptp(episode, axis=0)
    metrics["range_util"] = float(np.mean(action_range))

    return metrics


def _compute_baseline(action_episodes: list[np.ndarray]) -> dict:
    """Compute baseline statistics from existing episodes."""
    all_metrics: dict[str, list[float]] = {}

    for ep in action_episodes:
        m = _compute_episode_metrics(ep)
        for k, v in m.items():
            all_metrics.setdefault(k, []).append(v)

    baseline = {}
    for k, values in all_metrics.items():
        arr = np.array(values)
        median = float(np.median(arr))
        mad = float(np.median(np.abs(arr - median)))
        baseline[k] = {"median": median, "mad": max(mad, 1e-8)}

    return baseline


def _score_episode(
    episode: np.ndarray,
    baseline: dict,
) -> CoachFeedback:
    """Score a new episode against the baseline."""
    metrics = _compute_episode_metrics(episode)
    issues = []
    z_scores = {}

    for k, v in metrics.items():
        if k in baseline:
            b = baseline[k]
            z = abs(v - b["median"]) / (b["mad"] * 1.4826)  # MAD-based z-score
            z_scores[k] = z

            if k == "jerk" and z > 2.0:
                ratio = v / max(b["median"], 1e-8)
                issues.append(f"Jerk {ratio:.1f}x higher than median — jerky motion detected")
            elif k == "path_length" and z > 2.0:
                ratio = v / max(b["median"], 1e-8)
                if ratio > 1:
                    issues.append(
                        f"Path length {(ratio - 1) * 100:.0f}% longer than typical — "
                        "consider more direct approach"
                    )
                else:
                    issues.append(
                        f"Path length {(1 - ratio) * 100:.0f}% shorter than typical — "
                        "may be incomplete"
                    )
            elif k == "length" and z > 2.5:
                ratio = v / max(b["median"], 1e-8)
                if ratio > 1.5:
                    issues.append(f"Episode {ratio:.1f}x longer than median — possible hesitation")
                elif ratio < 0.5:
                    issues.append(f"Episode {ratio:.1f}x shorter than median — may be incomplete")

    # Determine verdict
    max_z = max(z_scores.values()) if z_scores else 0
    if max_z > 3.5:
        verdict = "redo"
    elif max_z > 2.0 or len(issues) >= 2:
        verdict = "review"
    else:
        verdict = "good"

    # Consistency score (inverse of mean z-score)
    mean_z = np.mean(list(z_scores.values())) if z_scores else 0
    consistency = max(0.0, min(1.0, 1.0 - mean_z / 5.0))

    suggestions = []
    if "jerk" in z_scores and z_scores["jerk"] > 2.0:
        suggestions.append("Slow down near the end of the trajectory for smoother motion")
    if "path_length" in z_scores and z_scores["path_length"] > 2.0:
        suggestions.append("Try to follow the same path as your earlier demonstrations")

    return CoachFeedback(
        episode_index=-1,  # set by caller
        verdict=verdict,
        consistency_score=round(consistency, 2),
        issues=issues,
        suggestions=suggestions,
    )


def _detect_fatigue(state: CoachState) -> str | None:
    """Detect operator fatigue from jerk trends."""
    recent = state.recent_jerk_trend
    if len(recent) < 10:
        return None

    last_5 = recent[-5:]
    first_5 = recent[:5]
    mean_last = np.mean(last_5)
    mean_first = np.mean(first_5)

    if mean_last > mean_first * 1.5:
        return (
            "Your last 5 episodes have higher jerk than your first 5. "
            "Take a short break — fatigue affects demonstration quality."
        )
    return None


def watch_and_coach(
    dataset_path: str,
    target_episodes: int = 50,
    policy: str = "act",
    poll_interval: float = 2.0,
) -> None:
    """Watch a dataset directory and coach the operator in real-time.

    Args:
        dataset_path: Path to dataset directory.
        target_episodes: Target number of episodes to collect.
        policy: Target policy for quality context.
        poll_interval: Seconds between directory polls.
    """
    from orbit.utils.display import console

    ds_path = Path(dataset_path)
    if not ds_path.exists():
        console.print(f"[red]Dataset path not found: {ds_path}[/red]")
        return

    console.print(f"\n[bold]Collection Coach[/bold] — Target: {target_episodes} episodes for {policy.upper()}")
    console.print(f"[dim]Watching: {ds_path}[/dim]\n")

    # Load existing episodes
    action_episodes = _load_existing_episodes(ds_path)
    state = CoachState(
        target_episodes=target_episodes,
        policy=policy,
        current_count=len(action_episodes),
    )

    if action_episodes:
        state.baseline_metrics = _compute_baseline(action_episodes)
        console.print(f"[dim]Loaded {len(action_episodes)} existing episodes as baseline[/dim]")
    else:
        console.print("[dim]No existing episodes — will build baseline as you collect[/dim]")

    seen_files: set[str] = set()
    for f in ds_path.rglob("*"):
        if f.suffix in (".parquet", ".npz", ".npy"):
            seen_files.add(str(f))

    try:
        while state.current_count < target_episodes:
            new_files = []
            for f in ds_path.rglob("*"):
                fstr = str(f)
                if fstr not in seen_files and f.suffix in (".parquet", ".npz", ".npy"):
                    new_files.append(f)
                    seen_files.add(fstr)

            for new_file in sorted(new_files):
                ep = _load_single_episode(new_file)
                if ep is None:
                    continue

                state.current_count += 1
                action_episodes.append(ep)

                # Update baseline periodically
                if state.current_count % 10 == 0 and len(action_episodes) >= 5:
                    state.baseline_metrics = _compute_baseline(action_episodes)

                if state.baseline_metrics:
                    feedback = _score_episode(ep, state.baseline_metrics)
                    feedback.episode_index = state.current_count
                    feedback.dataset_progress = (
                        f"{state.current_count}/{target_episodes} episodes "
                        f"({state.current_count * 100 // target_episodes}%)"
                    )

                    _print_feedback(console, feedback, state)

                    # Track jerk for fatigue detection
                    metrics = _compute_episode_metrics(ep)
                    if "jerk" in metrics:
                        state.recent_jerk_trend.append(metrics["jerk"])

                    # Fatigue check
                    fatigue_msg = _detect_fatigue(state)
                    if fatigue_msg:
                        console.print(f"\n  [yellow]Tip: {fatigue_msg}[/yellow]\n")
                else:
                    console.print(
                        f"  Episode {state.current_count} recorded "
                        f"({len(ep)} frames) — building baseline..."
                    )
                    if len(action_episodes) >= 5:
                        state.baseline_metrics = _compute_baseline(action_episodes)
                        console.print("[dim]  Baseline computed from first 5 episodes[/dim]")

            time.sleep(poll_interval)

    except KeyboardInterrupt:
        console.print("\n[dim]Coach stopped.[/dim]")

    # Final summary
    console.print(f"\n[bold]Collection Complete[/bold]")
    console.print(f"  Episodes: {state.current_count}/{target_episodes}")
    console.print(f"  Good: {state.total_good}  Review: {state.total_review}  Redo: {state.total_redo}")
    if state.current_count >= target_episodes:
        console.print(f"\n  Next: orbit analyze {dataset_path} --policy {policy}")
    console.print()


def _print_feedback(console, feedback: CoachFeedback, state: CoachState) -> None:
    """Print episode feedback to console."""
    verdict_colors = {"good": "green", "review": "yellow", "redo": "red"}
    verdict_icons = {"good": "\u2713", "review": "\u26a0", "redo": "\u274c"}
    color = verdict_colors.get(feedback.verdict, "white")
    icon = verdict_icons.get(feedback.verdict, "?")

    console.print(
        f"\n  Episode {feedback.episode_index} recorded"
    )
    console.print(
        f"    [{color}]{icon} Consistency: {feedback.consistency_score:.2f}[/{color}] "
        f"— [{color}]{feedback.verdict.upper()}[/{color}]"
    )

    for issue in feedback.issues:
        console.print(f"    [yellow]\u26a0[/yellow] {issue}")

    for suggestion in feedback.suggestions:
        console.print(f"    [dim]\u2192 {suggestion}[/dim]")

    # Progress bar
    pct = state.current_count * 100 // state.target_episodes
    filled = pct // 5
    bar = "\u2588" * filled + "\u2591" * (20 - filled)
    console.print(f"\n  Progress: {bar} {feedback.dataset_progress}")

    if feedback.verdict == "good":
        state.total_good += 1
    elif feedback.verdict == "review":
        state.total_review += 1
    else:
        state.total_redo += 1


def _load_existing_episodes(ds_path: Path) -> list[np.ndarray]:
    """Load existing episodes from a dataset directory."""
    episodes = []

    # Try LeRobot format first
    try:
        from orbit.analyzer.dataset_loader import load_episode_data
        from orbit.analyzer.coverage import _extract_array

        ep_data = load_episode_data(str(ds_path), max_episodes=200)
        if ep_data is not None and not ep_data.empty:
            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:
                        episodes.append(arr)
            if episodes:
                return episodes
    except Exception:
        pass

    # Fallback: scan for numpy/parquet files
    for f in sorted(ds_path.rglob("*.npz")):
        ep = _load_single_episode(f)
        if ep is not None:
            episodes.append(ep)

    for f in sorted(ds_path.rglob("*.npy")):
        ep = _load_single_episode(f)
        if ep is not None:
            episodes.append(ep)

    return episodes


def _load_single_episode(path: Path) -> np.ndarray | None:
    """Load a single episode from a file."""
    try:
        if path.suffix == ".npz":
            data = np.load(path)
            for key in ("actions", "action", "qpos", "joint_positions"):
                if key in data:
                    arr = data[key]
                    if arr.ndim == 2 and arr.shape[0] >= 3:
                        return arr
        elif path.suffix == ".npy":
            arr = np.load(path)
            if arr.ndim == 2 and arr.shape[0] >= 3:
                return arr
        elif path.suffix == ".parquet":
            import pandas as pd

            df = pd.read_parquet(path)
            if "action" in df.columns:
                from orbit.analyzer.coverage import _extract_array

                arr = _extract_array(df["action"])
                if arr is not None and arr.ndim == 2 and arr.shape[0] >= 3:
                    return arr
    except Exception:
        pass
    return None
