"""Integration test: create a bad dataset, analyze, clean, analyze again.

Verifies that orbit's full pipeline — analyze → clean → analyze — actually
improves the readiness grade when bad episodes are removed.
"""

from __future__ import annotations

import json
from pathlib import Path

import numpy as np
import pandas as pd
import pytest

# ---------------------------------------------------------------------------
# Skip if lerobot not installed
# ---------------------------------------------------------------------------

lerobot_available = False
try:
    from lerobot.datasets.lerobot_dataset import LeRobotDataset
    from lerobot.datasets.dataset_tools import delete_episodes

    lerobot_available = True
except ImportError:
    pass

pytestmark = pytest.mark.skipif(
    not lerobot_available, reason="lerobot not installed"
)

# ---------------------------------------------------------------------------
# Constants
# ---------------------------------------------------------------------------

REPO_ID = "orbit_test/clean_integration"
FPS = 30
ACTION_DIM = 6
STATE_DIM = 6
N_GOOD = 7
N_BAD = 3
GOOD_EP_LEN = 80


# ---------------------------------------------------------------------------
# Helpers
# ---------------------------------------------------------------------------


def _create_test_dataset(root: Path) -> LeRobotDataset:
    """Build a 10-episode dataset: 7 good + 3 bad."""
    features = {
        "observation.state": {
            "dtype": "float32",
            "shape": (STATE_DIM,),
            "names": None,
        },
        "action": {
            "dtype": "float32",
            "shape": (ACTION_DIM,),
            "names": None,
        },
    }

    dataset = LeRobotDataset.create(
        repo_id=REPO_ID,
        fps=FPS,
        features=features,
        root=root,
    )

    rng = np.random.RandomState(42)
    task = "pick_and_place"

    for ep_idx in range(N_GOOD + N_BAD):
        dataset.create_episode_buffer()

        if ep_idx == 0:
            # Bad: all zeros (dead frames)
            for _ in range(GOOD_EP_LEN):
                dataset.add_frame({
                    "observation.state": np.zeros(STATE_DIM, dtype=np.float32),
                    "action": np.zeros(ACTION_DIM, dtype=np.float32),
                    "task": task,
                })
        elif ep_idx == 1:
            # Bad: random noise (inconsistent)
            for _ in range(GOOD_EP_LEN):
                dataset.add_frame({
                    "observation.state": rng.randn(STATE_DIM).astype(np.float32) * 5.0,
                    "action": rng.randn(ACTION_DIM).astype(np.float32) * 5.0,
                    "task": task,
                })
        elif ep_idx == 2:
            # Bad: ultra-short (5 frames)
            for _ in range(5):
                dataset.add_frame({
                    "observation.state": (rng.randn(STATE_DIM) * 0.3).astype(np.float32),
                    "action": (rng.randn(ACTION_DIM) * 0.1).astype(np.float32),
                    "task": task,
                })
        else:
            # Good: smooth sinusoidal trajectory
            phase = rng.uniform(0, 2 * np.pi)
            for t in range(GOOD_EP_LEN):
                frac = t / GOOD_EP_LEN
                base = np.sin(np.linspace(phase, phase + np.pi, ACTION_DIM))
                action = (base * 0.3 * (1 + 0.5 * frac) + rng.randn(ACTION_DIM) * 0.02).astype(np.float32)
                state = (base * 0.5 + rng.randn(STATE_DIM) * 0.02).astype(np.float32)
                dataset.add_frame({
                    "observation.state": state,
                    "action": action,
                    "task": task,
                })

        dataset.save_episode()

    dataset.finalize()
    return dataset


def _extract_episodes(dataset: LeRobotDataset):
    """Extract action/state episode arrays and a DataFrame from a LeRobotDataset."""
    # Collect all frames grouped by episode
    ep_frames: dict[int, list[tuple[np.ndarray, np.ndarray]]] = {}
    rows = []

    for i in range(dataset.num_frames):
        frame = dataset[i]
        ep_idx = frame["episode_index"].item()
        a = frame["action"].numpy()
        s = frame["observation.state"].numpy()

        if ep_idx not in ep_frames:
            ep_frames[ep_idx] = []
        ep_frames[ep_idx].append((a, s))
        rows.append({
            "episode_index": ep_idx,
            "action": a.tolist(),
            "state": s.tolist(),
        })

    action_episodes = []
    state_episodes = []
    for ep_idx in sorted(ep_frames.keys()):
        pairs = ep_frames[ep_idx]
        action_episodes.append(np.array([p[0] for p in pairs]))
        state_episodes.append(np.array([p[1] for p in pairs]))

    return action_episodes, state_episodes, pd.DataFrame(rows)


def _run_analysis(action_episodes, state_episodes, episode_df, num_episodes, num_frames) -> dict:
    """Run the orbit analysis pipeline on extracted data."""
    from orbit.analyzer.coverage import compute_coverage
    from orbit.analyzer.dataset_loader import DatasetInfo
    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.success_predictor import DatasetReadinessGrader
    from orbit.analyzer.policy_fit import PolicyFitAnalyzer

    info = DatasetInfo(
        repo_id=REPO_ID,
        num_episodes=num_episodes,
        num_frames=num_frames,
        fps=FPS,
        robot_type="so100",
        camera_names=[],
        video_keys=[],
        shapes={"action": [ACTION_DIM], "observation.state": [STATE_DIM]},
    )

    coverage = compute_coverage(info, episode_df)
    signal_diag = compute_signal_diagnostics(action_episodes)
    episode_ranking = compute_episode_ranking(action_episodes, state_episodes)

    quality = compute_quality_score(
        info, coverage,
        action_episodes=action_episodes,
        state_episodes=state_episodes,
        signal_diagnostics_summary=(
            max(0.0, 1.0 - len(signal_diag.blockers) * 0.15
                - min(len(signal_diag.warnings), 6) * 0.05)
        ),
    )

    pf = PolicyFitAnalyzer()
    pf_result = pf.analyze(
        action_episodes,
        policy_type="act",
        episode_count=num_episodes,
    )

    grader = DatasetReadinessGrader()
    sig_health = {
        "dead_dimensions": signal_diag.dead_dimensions,
        "blockers": signal_diag.blockers,
        "warnings": signal_diag.warnings,
        "total_joints": len(signal_diag.joint_diagnostics),
        "clipping_joints": [
            jd.joint_index for jd in signal_diag.joint_diagnostics if jd.is_clipping
        ],
    }
    er_dict = {
        "total": len(episode_ranking.episodes),
        "remove_count": episode_ranking.remove_count,
        "review_count": episode_ranking.review_count,
    }

    readiness = grader.grade(
        quality_score=quality.overall_score * 100,
        episode_consistency=coverage.episode_consistency,
        signal_health=sig_health,
        policy_fit_score=pf_result.fit_score,
        episode_ranking=er_dict,
        action_divergence=0.0,
        episode_count=num_episodes,
        policy_type="act",
    )

    return {
        "num_episodes": num_episodes,
        "quality_score": round(quality.overall_score, 4),
        "grade": readiness.grade,
        "readiness_score": readiness.score,
        "remove_count": episode_ranking.remove_count,
        "remove_indices": episode_ranking.remove_indices,
        "dead_joints": signal_diag.dead_dimensions,
        "problems": readiness.problems,
    }


# ---------------------------------------------------------------------------
# Test
# ---------------------------------------------------------------------------


class TestCleanIntegration:
    """End-to-end: create bad dataset -> analyze -> clean -> analyze again."""

    def test_analyze_clean_analyze(self, tmp_path):
        """Full pipeline should detect bad episodes and improve the grade after cleaning."""
        root = tmp_path / "datasets" / REPO_ID

        # Step 1: Create the deliberately bad dataset
        dataset = _create_test_dataset(root)
        assert dataset.meta.total_episodes == N_GOOD + N_BAD

        # Step 2: Analyze BEFORE cleaning
        actions_b, states_b, df_b = _extract_episodes(dataset)
        num_frames_b = sum(len(a) for a in actions_b)
        before = _run_analysis(actions_b, states_b, df_b, N_GOOD + N_BAD, num_frames_b)

        print(f"\n--- BEFORE CLEANING ---")
        print(f"  Episodes:  {before['num_episodes']}")
        print(f"  Quality:   {before['quality_score']}")
        print(f"  Grade:     {before['grade']} ({before['readiness_score']}/100)")
        print(f"  To remove: {before['remove_count']} episodes {before['remove_indices']}")
        print(f"  Dead:      {before['dead_joints']}")
        print(f"  Problems:  {before['problems']}")

        assert before["remove_count"] > 0, "Should flag at least one bad episode for removal"

        # Step 3: Clean via the real lerobot API
        remove_indices = before["remove_indices"]
        clean_root = tmp_path / "datasets" / f"{REPO_ID}-clean"
        clean_repo = f"{REPO_ID}-clean"

        cleaned = delete_episodes(
            dataset,
            episode_indices=remove_indices,
            output_dir=clean_root,
            repo_id=clean_repo,
        )

        cleaned_count = cleaned.meta.total_episodes
        assert cleaned_count == (N_GOOD + N_BAD) - len(remove_indices), (
            f"Expected {(N_GOOD + N_BAD) - len(remove_indices)} episodes, got {cleaned_count}"
        )

        # Step 4: Analyze AFTER cleaning
        actions_a, states_a, df_a = _extract_episodes(cleaned)
        num_frames_a = sum(len(a) for a in actions_a)
        after = _run_analysis(actions_a, states_a, df_a, cleaned_count, num_frames_a)

        print(f"\n--- AFTER CLEANING ---")
        print(f"  Episodes:  {after['num_episodes']}")
        print(f"  Quality:   {after['quality_score']}")
        print(f"  Grade:     {after['grade']} ({after['readiness_score']}/100)")
        print(f"  To remove: {after['remove_count']} episodes")
        print(f"  Problems:  {after['problems']}")

        # Assertions
        grade_order = {"F": 0, "D": 1, "C": 2, "B": 3, "A": 4}
        assert grade_order[after["grade"]] >= grade_order[before["grade"]], (
            f"Grade should improve: {before['grade']} -> {after['grade']}"
        )

        assert after["quality_score"] >= before["quality_score"] - 0.01, (
            f"Quality should improve: {before['quality_score']} -> {after['quality_score']}"
        )

        assert after["remove_count"] <= before["remove_count"], (
            f"Fewer bad episodes after cleaning: {before['remove_count']} -> {after['remove_count']}"
        )

        print(f"\n--- RESULT ---")
        print(f"  Grade:   {before['grade']} -> {after['grade']}")
        print(f"  Quality: {before['quality_score']} -> {after['quality_score']}")
        print(f"  Removed {len(remove_indices)} episodes, {after['num_episodes']} remain")
