"""Recommendation engine for improving robot training datasets."""

from __future__ import annotations

from dataclasses import dataclass

from orbit.analyzer.coverage import FullCoverageReport
from orbit.analyzer.dataset_loader import DatasetInfo
from orbit.analyzer.quality import DEFAULT_EXPECTED_EPISODES, QualityReport

MAX_RECOMMENDATIONS = 6


@dataclass
class Recommendation:
    """A specific, actionable recommendation for improving a dataset."""

    priority: str  # "HIGH" | "MEDIUM" | "LOW"
    category: str  # e.g., "position_diversity", "episode_count"
    description: str  # human-readable, specific
    expected_impact: float  # estimated quality score improvement
    effort: str  # e.g., "collect N episodes" or "adjust setup"


def _is_replay_buffer(info: DatasetInfo) -> bool:
    """Detect if a dataset is likely RL replay buffer data (not human demos).

    Only triggers on strong indicators: 'replay' in dataset name, or both
    reward signals AND RL-like naming patterns (e.g., 'medium', 'expert',
    'random', 'mixed').  Reward signals alone are not sufficient — many
    human demo datasets include reward labels for evaluation purposes.
    """
    name_lower = info.repo_id.lower()
    if "replay" in name_lower:
        return True
    # Require BOTH reward signals AND RL dataset naming convention
    has_reward = any("reward" in k.lower() for k in info.shapes)
    rl_name_patterns = ("medium", "expert", "random", "mixed", "d4rl", "offline_rl")
    has_rl_name = any(p in name_lower for p in rl_name_patterns)
    if has_reward and has_rl_name:
        return True
    return False


def generate_recommendations(
    info: DatasetInfo,
    coverage: FullCoverageReport,
    quality: QualityReport,
) -> list[Recommendation]:
    """Generate specific, actionable recommendations sorted by expected impact.

    Args:
        info: Dataset metadata.
        coverage: Full coverage analysis report.
        quality: Quality report with component scores.

    Returns:
        List of up to 6 recommendations sorted by expected_impact descending.
    """
    recs: list[Recommendation] = []
    gap: float
    impact: float
    is_replay = _is_replay_buffer(info)

    # a) Episode count too low
    if info.num_episodes < DEFAULT_EXPECTED_EPISODES and not is_replay:
        gap = float(DEFAULT_EXPECTED_EPISODES - info.num_episodes)
        impact = 0.25 * (gap / DEFAULT_EXPECTED_EPISODES)
        recs.append(
            Recommendation(
                priority="HIGH",
                category="episode_count",
                description=(
                    f"Collect {int(gap)} more episodes to reach the minimum "
                    f"recommended count of {DEFAULT_EXPECTED_EPISODES}"
                ),
                expected_impact=round(impact, 2),
                effort=f"collect {int(gap)} episodes",
            )
        )

    # b) Position diversity low
    if coverage.position_diversity < 0.6:
        gap = 0.6 - coverage.position_diversity
        impact = 0.05 + gap * 0.25  # 0.05-0.15 range
        impact = min(0.15, impact)
        sparse_desc = ""
        if coverage.sparse_areas:
            area = coverage.sparse_areas[0]
            sparse_desc = f" with the object/robot in {area}"
        n_needed = max(3, int(gap * 30))
        recs.append(
            Recommendation(
                priority="HIGH",
                category="position_diversity",
                description=(
                    f"Collect {n_needed} episodes{sparse_desc} — "
                    f"currently only {coverage.position_diversity:.0%} of the "
                    f"position space is covered"
                ),
                expected_impact=round(impact, 2),
                effort=f"collect {n_needed} episodes",
            )
        )

    # c) Action diversity low
    if coverage.action_diversity < 0.5:
        gap = 0.5 - coverage.action_diversity
        impact = 0.05 + gap * 0.14  # 0.05-0.12 range
        impact = min(0.12, impact)
        dim_desc = ""
        if coverage.low_variance_dims:
            dim_desc = f" — {coverage.low_variance_dims[0]} actions are too uniform"
        uniform_pct = coverage.uniform_action_pct if coverage.uniform_action_pct > 0 else 50.0
        recs.append(
            Recommendation(
                priority="HIGH",
                category="action_diversity",
                description=(
                    f"Vary your approach angle/speed/grip{dim_desc}"
                    f" across {uniform_pct:.0f}% of episodes"
                ),
                expected_impact=round(impact, 2),
                effort="adjust setup",
            )
        )

    # d) No failure recovery examples (heuristic: check if consistency is
    #    very high, implying all demos are clean successes)
    if coverage.episode_consistency > 0.95 and info.num_episodes >= 10 and not is_replay:
        recs.append(
            Recommendation(
                priority="MEDIUM",
                category="failure_recovery",
                description=("Add 3-5 demonstrations where the robot nearly fails but recovers"),
                expected_impact=0.05,
                effort="collect 5 episodes",
            )
        )

    # Replay buffer specific: suggest filtering by reward
    if is_replay:
        recs.append(
            Recommendation(
                priority="MEDIUM",
                category="data_filtering",
                description=(
                    "This appears to be RL replay buffer data — consider filtering "
                    "episodes by reward to keep only successful trajectories"
                ),
                expected_impact=0.08,
                effort="filter dataset",
            )
        )

    # e) Episode consistency issues (outliers found)
    if coverage.outlier_episodes:
        episode_ids = ", ".join(str(e) for e in coverage.outlier_episodes[:5])
        impact = 0.02 + min(0.02, len(coverage.outlier_episodes) * 0.005)
        recs.append(
            Recommendation(
                priority="MEDIUM",
                category="episode_consistency",
                description=(
                    f"Review episodes {episode_ids} — they appear significantly "
                    f"different from others (possible recording errors)"
                ),
                expected_impact=round(impact, 2),
                effort="adjust setup",
            )
        )

    # f) Temporal coverage low
    if coverage.temporal_coverage < 0.5:
        gap = 0.5 - coverage.temporal_coverage
        impact = 0.04 + gap * 0.08  # 0.04-0.08 range
        impact = min(0.08, impact)
        path_info = ""
        if coverage.dominant_path_pct > 50:
            path_info = f" — {coverage.dominant_path_pct:.0f}% of episodes follow the same path"
        recs.append(
            Recommendation(
                priority="MEDIUM",
                category="temporal_coverage",
                description=(
                    f"Try different strategies for the task — most episodes "
                    f"follow the same motion pattern{path_info}"
                ),
                expected_impact=round(impact, 2),
                effort="adjust setup",
            )
        )

    # Sort by expected_impact descending
    recs.sort(key=lambda r: r.expected_impact, reverse=True)

    # Limit to top MAX_RECOMMENDATIONS
    return recs[:MAX_RECOMMENDATIONS]
