"""Task context analyzer — estimates task difficulty from action data and metadata.

Provides heuristic signals for task complexity without requiring a VLM:
action dimensionality, episode length patterns, coordination complexity,
gripper transitions, bimanual detection, and trajectory multimodality.
"""

from __future__ import annotations

import logging
from dataclasses import dataclass, field

import numpy as np

logger = logging.getLogger(__name__)


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


@dataclass
class TaskContext:
    """Estimated task difficulty from action data analysis."""

    task_complexity_score: float  # 0-1 overall difficulty
    dimension_complexity: float  # Higher action dims = harder
    temporal_complexity: float  # Longer/variable horizons = harder
    coordination_complexity: float  # Multi-joint coordination = harder
    object_interaction_complexity: float  # Gripper transitions, contact phases
    trajectory_multimodality: float  # Multiple strategies = harder to learn
    is_bimanual: bool  # Detected bimanual manipulation
    action_dim: int  # Detected action dimensionality
    mean_episode_length: float  # Average episode length in steps
    estimated_task_type: str  # "simple", "moderate", "complex", "very_complex"
    complexity_breakdown: dict[str, float] = field(default_factory=dict)


# ---------------------------------------------------------------------------
# Complexity scoring functions
# ---------------------------------------------------------------------------

# Action dim thresholds
_DIM_SCORES = [
    (2, 0.10),  # 2D (e.g., pusht) — very simple
    (3, 0.20),  # 3D position only
    (6, 0.35),  # 6D single arm (pos + rot) or 7D with gripper
    (7, 0.40),  # 7D single arm + gripper
    (12, 0.60),  # Bimanual without grippers
    (14, 0.70),  # Bimanual with grippers
    (20, 0.85),  # High-DOF (humanoid hands)
    (30, 0.95),  # Very high DOF
]


def _dimension_complexity(action_dim: int) -> float:
    """Map action dimensionality to a 0-1 complexity score."""
    if action_dim <= 0:
        return 0.5

    prev_dim, prev_score = 1, 0.05
    for dim_thresh, score in _DIM_SCORES:
        if action_dim <= dim_thresh:
            # Interpolate
            frac = (action_dim - prev_dim) / max(1, dim_thresh - prev_dim)
            return prev_score + frac * (score - prev_score)
        prev_dim, prev_score = dim_thresh, score

    return min(1.0, 0.95 + (action_dim - 30) * 0.005)


def _temporal_complexity(episode_lengths: list[int]) -> float:
    """Score temporal complexity from episode length distribution.

    Short & consistent = simple task. Long & variable = complex.
    """
    if not episode_lengths:
        return 0.5

    lengths = np.array(episode_lengths, dtype=float)
    mean_len = float(np.mean(lengths))
    std_len = float(np.std(lengths))
    cv = std_len / max(1.0, mean_len)  # coefficient of variation

    # Length component: longer episodes = harder
    # 50 steps ~ simple, 200 ~ moderate, 500+ ~ complex
    length_score = float(np.clip(mean_len / 500.0, 0.0, 1.0))

    # Variability component: high CV = harder (more variable task execution)
    variability_score = float(np.clip(cv / 0.5, 0.0, 1.0))

    return float(np.clip(length_score * 0.6 + variability_score * 0.4, 0.0, 1.0))


def _coordination_complexity(action_episodes: list[np.ndarray]) -> float:
    """Score coordination complexity from simultaneous multi-joint movement.

    High variance in multiple joints simultaneously = coordinated manipulation.
    """
    if not action_episodes:
        return 0.5

    coordination_scores = []
    for ep in action_episodes:
        if ep.ndim == 1:
            ep = ep.reshape(-1, 1)
        if len(ep) < 3:
            continue

        # Compute per-step velocity
        velocity = np.diff(ep, axis=0)

        # For each timestep, count how many joints are active
        active_threshold = np.std(velocity) * 0.1 + 1e-8
        active_joints_per_step = np.sum(np.abs(velocity) > active_threshold, axis=1)
        action_dim = ep.shape[1]

        # Fraction of joints active simultaneously, averaged over time
        if action_dim > 0:
            mean_active_frac = float(np.mean(active_joints_per_step / action_dim))
            coordination_scores.append(mean_active_frac)

    if not coordination_scores:
        return 0.5

    mean_coord = float(np.mean(coordination_scores))
    # Map: 0.3 (few joints at once) -> low, 0.7+ (most joints) -> high
    return float(np.clip((mean_coord - 0.1) / 0.6, 0.0, 1.0))


def _object_interaction_complexity(
    action_episodes: list[np.ndarray],
    gripper_dim: int | None = None,
) -> float:
    """Estimate object interaction complexity from gripper transitions.

    More open/close cycles per episode = more object interactions = harder.
    """
    if not action_episodes:
        return 0.5

    action_dim = action_episodes[0].shape[1] if action_episodes[0].ndim == 2 else 1

    # Auto-detect gripper dim if not provided
    if gripper_dim is None:
        # Gripper is typically last dim, with bimodal distribution
        if action_dim >= 2:
            gripper_dim = action_dim - 1
        else:
            return 0.3  # Can't detect gripper in 1D

    transitions_per_ep = []
    for ep in action_episodes:
        if ep.ndim == 1:
            continue
        if gripper_dim >= ep.shape[1]:
            continue

        gripper = ep[:, gripper_dim]
        if len(gripper) < 2:
            continue

        # Detect transitions: sign changes or large jumps in gripper value
        gripper_range = float(np.max(gripper) - np.min(gripper))
        if gripper_range < 1e-6:
            transitions_per_ep.append(0)
            continue

        # Normalize gripper to [0, 1]
        gripper_norm = (gripper - np.min(gripper)) / gripper_range
        # Threshold at 0.5 to binarize
        binary = (gripper_norm > 0.5).astype(int)
        transitions = int(np.sum(np.abs(np.diff(binary))))
        transitions_per_ep.append(transitions)

    if not transitions_per_ep:
        return 0.3

    mean_transitions = float(np.mean(transitions_per_ep))
    # 0 transitions = no grasping, 1-2 = simple pick-place, 3+ = complex manipulation
    return float(np.clip(mean_transitions / 6.0, 0.0, 1.0))


def _trajectory_multimodality(action_episodes: list[np.ndarray]) -> float:
    """Estimate whether demonstrations show multiple distinct strategies.

    Multiple strategies for the same task = harder to learn with BC.
    Uses pairwise trajectory distance to detect clusters.
    """
    if len(action_episodes) < 3:
        return 0.3

    # Normalize episodes to same length for comparison
    target_len = 50
    normalized = []
    for ep in action_episodes:
        if ep.ndim == 1:
            ep = ep.reshape(-1, 1)
        if len(ep) < 2:
            continue
        # Resample to target_len
        indices = np.linspace(0, len(ep) - 1, target_len).astype(int)
        normalized.append(ep[indices].flatten())

    if len(normalized) < 3:
        return 0.3

    normalized = np.array(normalized)

    # Compute pairwise L2 distances
    # Normalize by trajectory scale
    scale = np.std(normalized) + 1e-8
    normalized = normalized / scale

    n = len(normalized)
    max_pairs = min(n * (n - 1) // 2, 500)  # Cap for performance

    distances = []
    rng = np.random.RandomState(42)
    if n * (n - 1) // 2 <= max_pairs:
        for i in range(n):
            for j in range(i + 1, n):
                distances.append(float(np.linalg.norm(normalized[i] - normalized[j])))
    else:
        for _ in range(max_pairs):
            i, j = rng.choice(n, 2, replace=False)
            distances.append(float(np.linalg.norm(normalized[i] - normalized[j])))

    distances = np.array(distances)

    # High coefficient of variation in pairwise distances suggests multimodality
    mean_dist = float(np.mean(distances))
    std_dist = float(np.std(distances))
    cv = std_dist / max(1e-8, mean_dist)

    # CV > 0.5 suggests multiple clusters of trajectories
    return float(np.clip(cv / 0.8, 0.0, 1.0))


def _detect_bimanual(action_dim: int) -> bool:
    """Heuristic: bimanual if action dim >= 12 (2x 6D arms or 2x 7D with grippers)."""
    return action_dim >= 12


def _estimate_task_type(complexity: float) -> str:
    """Map complexity score to a human-readable task type."""
    if complexity < 0.25:
        return "simple"
    elif complexity < 0.50:
        return "moderate"
    elif complexity < 0.75:
        return "complex"
    else:
        return "very_complex"


# ---------------------------------------------------------------------------
# Main analyzer
# ---------------------------------------------------------------------------

# Weights for combining complexity components
_COMPLEXITY_WEIGHTS = {
    "dimension": 0.20,
    "temporal": 0.20,
    "coordination": 0.20,
    "object_interaction": 0.20,
    "multimodality": 0.20,
}


class TaskContextAnalyzer:
    """Estimates task difficulty from action data and metadata."""

    def __init__(self, weights: dict[str, float] | None = None):
        self.weights = weights or dict(_COMPLEXITY_WEIGHTS)

    def analyze(
        self,
        action_episodes: list[np.ndarray],
        metadata: dict | None = None,
        gripper_dim: int | None = None,
    ) -> TaskContext:
        """Analyze task context from action episodes.

        Args:
            action_episodes: List of per-episode action arrays (T, action_dim).
            metadata: Optional dict with keys like 'robot_type', 'task_description'.
            gripper_dim: Optional gripper dimension index.

        Returns:
            TaskContext with complexity scores and breakdown.
        """
        metadata = metadata or {}

        if not action_episodes:
            return TaskContext(
                task_complexity_score=0.5,
                dimension_complexity=0.5,
                temporal_complexity=0.5,
                coordination_complexity=0.5,
                object_interaction_complexity=0.5,
                trajectory_multimodality=0.5,
                is_bimanual=False,
                action_dim=0,
                mean_episode_length=0.0,
                estimated_task_type="moderate",
                complexity_breakdown={},
            )

        # Determine action dim
        first = action_episodes[0]
        action_dim = first.shape[1] if first.ndim == 2 else 1

        # Episode lengths
        episode_lengths = [len(ep) for ep in action_episodes]
        mean_ep_len = float(np.mean(episode_lengths))

        # Compute each complexity component
        dim_score = _dimension_complexity(action_dim)
        temp_score = _temporal_complexity(episode_lengths)
        coord_score = _coordination_complexity(action_episodes)
        obj_score = _object_interaction_complexity(action_episodes, gripper_dim)
        multi_score = _trajectory_multimodality(action_episodes)

        # Bimanual detection
        is_bimanual = _detect_bimanual(action_dim)
        if is_bimanual:
            # Bimanual tasks get a complexity boost
            coord_score = min(1.0, coord_score + 0.15)

        # Override from metadata if available
        if metadata.get("is_bimanual") is not None:
            is_bimanual = bool(metadata["is_bimanual"])

        # Weighted combination
        breakdown = {
            "dimension": dim_score,
            "temporal": temp_score,
            "coordination": coord_score,
            "object_interaction": obj_score,
            "multimodality": multi_score,
        }

        overall = sum(self.weights.get(k, 0.2) * v for k, v in breakdown.items())
        overall = float(np.clip(overall, 0.0, 1.0))

        return TaskContext(
            task_complexity_score=overall,
            dimension_complexity=dim_score,
            temporal_complexity=temp_score,
            coordination_complexity=coord_score,
            object_interaction_complexity=obj_score,
            trajectory_multimodality=multi_score,
            is_bimanual=is_bimanual,
            action_dim=action_dim,
            mean_episode_length=mean_ep_len,
            estimated_task_type=_estimate_task_type(overall),
            complexity_breakdown=breakdown,
        )
