"""Task type inference engine — understands what kind of task a dataset is for.

Infers task profiles from descriptions or dataset names, enabling task-aware
coverage analysis and recommendations.
"""

from __future__ import annotations

import re
from dataclasses import dataclass, field

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


@dataclass
class TaskPhase:
    """A single phase within a task (e.g., approach, grasp, transport)."""

    name: str
    description: str
    expected_fraction: float  # fraction of episode duration
    key_action_dims: list[int] = field(default_factory=list)


@dataclass
class TaskProfile:
    """Complete profile describing a task type and its analysis requirements."""

    task_type: str  # "pick_and_place", "insertion", "pouring", "wiping", "stacking", "generic"
    phases: list[TaskPhase]
    critical_dimensions: list[str]  # which coverage dimensions matter MOST
    typical_failure_modes: list[str]  # known failure patterns
    min_recommended_episodes: int
    expected_episode_duration_range: tuple[float, float]  # seconds
    keywords: list[str]


# ---------------------------------------------------------------------------
# Task profiles
# ---------------------------------------------------------------------------

PICK_AND_PLACE = TaskProfile(
    task_type="pick_and_place",
    phases=[
        TaskPhase("approach", "Move end-effector toward the object", 0.25, [0, 1, 2]),
        TaskPhase("grasp", "Close gripper on the object", 0.15, [3, 4, 5]),
        TaskPhase("transport", "Lift and carry to target location", 0.35, [0, 1, 2]),
        TaskPhase("place", "Lower object onto target surface", 0.15, [0, 1, 2]),
        TaskPhase("release", "Open gripper and retract", 0.10, [3, 4, 5]),
    ],
    critical_dimensions=["position_diversity", "grasp_variety", "placement_variety"],
    typical_failure_modes=[
        "Fails to grasp — approach angle too narrow across demonstrations",
        "Drops during transport — no recovery demonstrations included",
        "Misses target — insufficient placement position variety",
        "Knocks object over — approach speed not varied enough",
    ],
    min_recommended_episodes=50,
    expected_episode_duration_range=(2.0, 10.0),
    keywords=[
        "pick",
        "place",
        "pick_and_place",
        "pick-and-place",
        "pickandplace",
        "grasp",
        "grab",
        "put",
        "move",
        "transfer",
        "cabinet",
        "shelf",
        "bin",
        "sort",
        "rearrange",
        "stack_and_place",
    ],
)

INSERTION = TaskProfile(
    task_type="insertion",
    phases=[
        TaskPhase("approach", "Move peg/plug toward the hole/socket", 0.25, [0, 1, 2]),
        TaskPhase("align", "Fine-tune orientation to match hole axis", 0.20, [3, 4, 5]),
        TaskPhase("insert", "Push peg into hole with controlled force", 0.35, [0, 1, 2]),
        TaskPhase("seat", "Ensure full insertion and release", 0.20, [0, 1, 2, 3, 4, 5]),
    ],
    critical_dimensions=["alignment_precision", "insertion_force_profile", "orientation_variety"],
    typical_failure_modes=[
        "Jams during insertion — demonstrations lack angular variation",
        "Misaligns — approach orientation too uniform across demos",
        "Partial insertion — episodes cut short before full seating",
        "Cross-threads — insufficient demonstrations of corrective motions",
    ],
    min_recommended_episodes=60,
    expected_episode_duration_range=(3.0, 12.0),
    keywords=[
        "insert",
        "insertion",
        "peg",
        "hole",
        "plug",
        "socket",
        "connector",
        "thread",
        "screw",
        "assemble",
        "assembly",
        "fit",
    ],
)

POURING = TaskProfile(
    task_type="pouring",
    phases=[
        TaskPhase("grasp_container", "Grasp the source container", 0.20, [3, 4, 5]),
        TaskPhase("transport", "Move container toward target", 0.20, [0, 1, 2]),
        TaskPhase("tilt_pour", "Tilt container to pour contents", 0.35, [3, 4, 5]),
        TaskPhase(
            "right_and_return",
            "Right container and return to start",
            0.25,
            [0, 1, 2, 3, 4, 5],
        ),
    ],
    critical_dimensions=[
        "wrist_rotation_variety",
        "pour_angle_coverage",
        "transport_path_diversity",
    ],
    typical_failure_modes=[
        "Spills — tilt angle too uniform across demonstrations",
        "Incomplete pour — episodes cut short before pour finishes",
        "Knocks target — transport path doesn't vary enough",
        "Over-pours — no demonstrations of partial pouring",
    ],
    min_recommended_episodes=70,
    expected_episode_duration_range=(4.0, 15.0),
    keywords=[
        "pour",
        "pouring",
        "fill",
        "liquid",
        "cup",
        "mug",
        "bottle",
        "dispense",
        "water",
        "drink",
        "serve",
    ],
)

WIPING = TaskProfile(
    task_type="wiping",
    phases=[
        TaskPhase("approach_surface", "Move tool to contact the surface", 0.15, [0, 1, 2]),
        TaskPhase("contact", "Press tool against surface with appropriate force", 0.10, [2]),
        TaskPhase("sweep", "Move tool across surface in sweeping motions", 0.55, [0, 1]),
        TaskPhase("lift_and_check", "Lift tool and inspect surface", 0.20, [0, 1, 2]),
    ],
    critical_dimensions=["sweep_path_diversity", "contact_force_consistency", "area_coverage"],
    typical_failure_modes=[
        "Misses spots — sweep patterns don't cover full surface area",
        "Streaks — contact pressure varies too much across demonstrations",
        "Lifts tool mid-wipe — demonstrations not smooth enough",
        "Only wipes in one direction — needs multi-directional patterns",
    ],
    min_recommended_episodes=40,
    expected_episode_duration_range=(3.0, 12.0),
    keywords=[
        "wipe",
        "wiping",
        "clean",
        "scrub",
        "sweep",
        "polish",
        "table",
        "surface",
        "cloth",
        "sponge",
        "mop",
    ],
)

STACKING = TaskProfile(
    task_type="stacking",
    phases=[
        TaskPhase("approach", "Move to pick up source block/object", 0.20, [0, 1, 2]),
        TaskPhase("grasp", "Grasp the object securely", 0.10, [3, 4, 5]),
        TaskPhase("lift_and_align", "Lift object and align over target stack", 0.30, [0, 1, 2]),
        TaskPhase("stack", "Lower object carefully onto the stack", 0.25, [0, 1, 2]),
        TaskPhase("release", "Release object and retract", 0.15, [3, 4, 5]),
    ],
    critical_dimensions=["vertical_precision", "grasp_stability", "alignment_accuracy"],
    typical_failure_modes=[
        "Stack topples — placement alignment insufficient across demos",
        "Drops during lift — grasp demonstrations too uniform",
        "Misaligns stack — not enough variety in approach angles",
        "Collapses on contact — insufficient demonstrations of gentle placement",
    ],
    min_recommended_episodes=60,
    expected_episode_duration_range=(3.0, 12.0),
    keywords=[
        "stack",
        "stacking",
        "tower",
        "block",
        "cube",
        "pile",
        "build",
        "layer",
        "balance",
    ],
)

GENERIC = TaskProfile(
    task_type="generic",
    phases=[
        TaskPhase("start", "Initial phase of the task", 0.25, []),
        TaskPhase("main_action", "Primary manipulation action", 0.50, []),
        TaskPhase("finish", "Task completion and retraction", 0.25, []),
    ],
    critical_dimensions=["position_diversity", "action_diversity", "temporal_coverage"],
    typical_failure_modes=[
        "Policy fails in unseen regions of workspace",
        "Insufficient demonstration variety for robust generalization",
        "Episode durations vary too much — inconsistent task execution",
    ],
    min_recommended_episodes=50,
    expected_episode_duration_range=(2.0, 15.0),
    keywords=[],
)

# Registry of all profiles (order matters — specific before generic)
_ALL_PROFILES: list[TaskProfile] = [
    PICK_AND_PLACE,
    INSERTION,
    POURING,
    WIPING,
    STACKING,
]


# ---------------------------------------------------------------------------
# Inference
# ---------------------------------------------------------------------------


def _tokenize(text: str) -> set[str]:
    """Lowercase and split text into tokens, including underscored variants."""
    text = text.lower().strip()
    # Split on whitespace, hyphens, underscores, slashes
    tokens = set(re.split(r"[\s\-_/]+", text))
    # Also add the full string as a token (for compound keywords)
    tokens.add(text)
    return tokens


def infer_task_profile(
    task_description: str | None = None,
    dataset_name: str | None = None,
) -> TaskProfile:
    """Infer the best-matching TaskProfile from a description or dataset name.

    Parameters
    ----------
    task_description:
        Free-text task description (e.g., "pick up mugs and place on shelf").
    dataset_name:
        Dataset name or repo ID (e.g., "lerobot/aloha_sim_insertion_human").

    Returns
    -------
    TaskProfile:
        Best matching profile, or GENERIC if no match.
    """
    # Combine all available text
    parts: list[str] = []
    if task_description:
        parts.append(task_description)
    if dataset_name:
        parts.append(dataset_name)

    if not parts:
        return GENERIC

    combined = " ".join(parts)
    tokens = _tokenize(combined)

    best_profile: TaskProfile | None = None
    best_score = 0

    for profile in _ALL_PROFILES:
        score = 0
        for keyword in profile.keywords:
            kw_tokens = _tokenize(keyword)
            # Check if any keyword token appears in the input tokens
            for kw_tok in kw_tokens:
                if kw_tok in tokens:
                    score += 1
                # Also check substring matching for compound words
                elif any(kw_tok in t for t in tokens):
                    score += 0.5

        if score > best_score:
            best_score = score
            best_profile = profile

    if best_profile is not None and best_score > 0:
        return best_profile

    return GENERIC


def get_profile_by_name(task_type: str) -> TaskProfile:
    """Look up a profile by its task_type string."""
    for profile in _ALL_PROFILES:
        if profile.task_type == task_type:
            return profile
    return GENERIC
