"""Coverage analysis for LeRobot datasets."""

from __future__ import annotations

from dataclasses import dataclass, field

import numpy as np
import pandas as pd

from orbit.analyzer.dataset_loader import DatasetInfo


@dataclass
class FullCoverageReport:
    """Coverage analysis results across multiple dimensions."""

    position_diversity: float  # 0.0 to 1.0
    action_diversity: float  # 0.0 to 1.0
    episode_consistency: float  # 0.0 to 1.0
    temporal_coverage: float  # 0.0 to 1.0

    # Detail fields for recommendations
    sparse_areas: list[str] = field(default_factory=list)
    low_variance_dims: list[str] = field(default_factory=list)
    outlier_episodes: list[int] = field(default_factory=list)
    uniform_action_pct: float = 0.0
    dominant_path_pct: float = 0.0


def compute_coverage(info: DatasetInfo, episode_data: pd.DataFrame) -> FullCoverageReport:
    """Compute coverage metrics from episode data.

    Args:
        info: Dataset metadata.
        episode_data: DataFrame with episode_index, action, state columns.

    Returns:
        FullCoverageReport with all coverage metrics.
    """
    if episode_data.empty:
        return FullCoverageReport(
            position_diversity=0.0,
            action_diversity=0.0,
            episode_consistency=0.0,
            temporal_coverage=0.0,
        )

    position_div, sparse_areas = _compute_position_diversity(episode_data)
    action_div, low_var_dims, uniform_pct = _compute_action_diversity(episode_data)
    consistency, outliers = _compute_episode_consistency(episode_data)
    temporal, dominant_pct = _compute_temporal_coverage(episode_data)

    return FullCoverageReport(
        position_diversity=position_div,
        action_diversity=action_div,
        episode_consistency=consistency,
        temporal_coverage=temporal,
        sparse_areas=sparse_areas,
        low_variance_dims=low_var_dims,
        outlier_episodes=outliers,
        uniform_action_pct=uniform_pct,
        dominant_path_pct=dominant_pct,
    )


def _extract_array(series: pd.Series) -> np.ndarray | None:
    """Extract numeric array from a column that may contain lists or arrays."""
    if series.empty:
        return None
    first = series.iloc[0]
    if first is None:
        return None
    if isinstance(first, (list, np.ndarray)):
        try:
            return np.array(series.tolist(), dtype=float)
        except (ValueError, TypeError):
            return None
    return None


def _compute_position_diversity(df: pd.DataFrame) -> tuple[float, list[str]]:
    """Compute how well the dataset covers the position space."""
    state_arr = _extract_array(df["state"]) if "state" in df.columns else None
    if state_arr is None or state_arr.size == 0:
        return 0.5, []

    # Use ALL state dimensions for diversity — there is no universal state
    # ordering (e.g. Austin BUDS: dims 0-2 are joint angles, not XYZ;
    # DLR Sara Pour: dims 0-2 ARE XYZ).  Using the full state vector
    # measures state-space coverage without assuming structure.
    if state_arr.ndim == 1:
        pos = state_arr.reshape(-1, 1)
    else:
        pos = state_arr

    # Discretize into grid and measure coverage
    n_bins = 5
    sparse_areas = []
    coverages = []

    for dim_idx in range(pos.shape[1]):
        col = pos[:, dim_idx]
        col_range = col.max() - col.min()
        if col_range < 1e-8:
            coverages.append(0.0)
            continue
        bins = np.linspace(col.min(), col.max(), n_bins + 1)
        counts, _ = np.histogram(col, bins=bins)
        filled = float(np.sum(counts > 0) / n_bins)
        coverages.append(filled)

        # Identify sparse bins
        threshold = len(col) / (n_bins * 3)
        for b_idx, count in enumerate(counts):
            if count < threshold:
                labels = ["left", "center-left", "center", "center-right", "right"]
                if b_idx < len(labels):
                    sparse_areas.append(f"the {labels[b_idx]} region (dimension {dim_idx})")

    diversity = float(np.mean(coverages)) if coverages else 0.5
    return min(1.0, diversity), sparse_areas


def _compute_action_diversity(df: pd.DataFrame) -> tuple[float, list[str], float]:
    """Compute how diverse the actions are across episodes."""
    action_arr = _extract_array(df["action"]) if "action" in df.columns else None
    if action_arr is None or action_arr.size == 0:
        return 0.5, [], 0.0

    n_dims = action_arr.shape[1] if action_arr.ndim == 2 else 1
    low_var_dims = []
    diversity_scores = []

    for dim_idx in range(n_dims):
        col = action_arr[:, dim_idx] if action_arr.ndim == 2 else action_arr
        std = np.std(col)
        col_range = np.ptp(col)
        # Coefficient of variation relative to range
        if col_range > 1e-8:
            cv = std / col_range
            diversity_scores.append(min(1.0, cv * 4))  # Scale so 0.25 cv -> 1.0
        else:
            diversity_scores.append(0.0)
            low_var_dims.append(f"dimension {dim_idx}")

    overall = float(np.mean(diversity_scores)) if diversity_scores else 0.5
    # Compute uniform action percentage
    uniform_count = sum(1 for s in diversity_scores if s < 0.2)
    uniform_pct = uniform_count / max(1, len(diversity_scores)) * 100

    return min(1.0, overall), low_var_dims, uniform_pct


def _compute_episode_consistency(df: pd.DataFrame) -> tuple[float, list[int]]:
    """Compute consistency across episodes and find outliers."""
    if "episode_index" not in df.columns:
        return 0.5, []

    episodes = df.groupby("episode_index")
    lengths = episodes.size()

    if len(lengths) < 2:
        return 1.0, []

    mean_len = lengths.mean()
    std_len = lengths.std()

    # Outliers are episodes whose length deviates > 2 std from the mean
    outliers = []
    if std_len > 0:
        z_scores = np.abs((lengths - mean_len) / std_len)
        outliers = [int(idx) for idx, z in z_scores.items() if z > 2.0]

    # Consistency = 1 - normalized coefficient of variation
    cv = std_len / mean_len if mean_len > 0 else 0
    consistency = max(0.0, min(1.0, 1.0 - cv))

    return consistency, outliers


def _compute_temporal_coverage(df: pd.DataFrame) -> tuple[float, float]:
    """Compute temporal/trajectory diversity using numpy (no scipy dependency)."""
    action_arr = _extract_array(df["action"]) if "action" in df.columns else None
    if action_arr is None or "episode_index" not in df.columns:
        return 0.5, 50.0

    episodes = df.groupby("episode_index")
    if len(episodes) < 2:
        return 0.5, 100.0

    # Compute per-episode mean action as a trajectory fingerprint
    fingerprints = []
    for _, ep_df in episodes:
        ep_actions = _extract_array(ep_df["action"])
        if ep_actions is not None and ep_actions.size > 0:
            fingerprints.append(np.mean(ep_actions, axis=0))

    if len(fingerprints) < 2:
        return 0.5, 100.0

    fingerprints_arr = np.array(fingerprints)
    n = len(fingerprints_arr)

    # Subsample for speed: cap at 100 fingerprints for pairwise computation
    if n > 100:
        rng = np.random.RandomState(42)
        idx = rng.choice(n, 100, replace=False)
        fingerprints_arr = fingerprints_arr[idx]
        n = 100

    # Pairwise distance using numpy broadcasting (avoids scipy import)
    # diff[i,j] = fingerprints[i] - fingerprints[j]
    diff = fingerprints_arr[:, np.newaxis, :] - fingerprints_arr[np.newaxis, :, :]
    dist_matrix = np.sqrt(np.sum(diff ** 2, axis=2))

    # Extract upper triangle (pairwise distances without duplicates/diagonal)
    triu_idx = np.triu_indices(n, k=1)
    distances = dist_matrix[triu_idx]

    if len(distances) == 0:
        return 0.5, 100.0

    # Diversity = normalized spread of distances
    mean_dist = np.mean(distances)
    max_possible = np.max(distances) if np.max(distances) > 0 else 1.0
    temporal_score = min(1.0, mean_dist / max_possible) if max_possible > 0 else 0.0

    # Dominant path: count how many episodes are within small distance of mode
    median_dist = np.median(distances)
    close_threshold = median_dist * 0.5
    close_counts = np.sum(dist_matrix < close_threshold, axis=1)
    dominant_pct = float(np.max(close_counts) / n * 100)

    return temporal_score, dominant_pct
