AffordanceNet/phaseaware_data_construction/gripper_segmenter.py

"""Gripper segmenter: map frames to subgoals using gripper state signal."""

from __future__ import annotations

import logging
import warnings
from dataclasses import dataclass
from typing import List, Optional, Tuple

import numpy as np

from phaseaware_data_construction.instruction_parser import Subgoal

logger = logging.getLogger(__name__)

# Subgoal types that correspond to pick-place manipulation cycles
MANIP_PICK_TYPE = "pick"
MANIP_PLACE_TYPE = "place"
# Subgoal types that are standalone interactions (no gripper cycle)
STANDALONE_TYPES = {"open", "close", "turn_on", "turn_off", "push"}


@dataclass
class SegmentationResult:
    frame_subgoals: List[str]   # per-frame subgoal label
    frame_progress: List[str]   # per-frame "early" or "late"
    n_cycles_detected: int
    n_cycles_expected: int
    warnings: List[str]


def segment_episode(
    gripper_signal: np.ndarray,
    subgoals: List[Subgoal],
    signal_source: str = "action",
    threshold: float = 0.02,
    median_filter_size: int = 5,
    min_cycle_frames: int = 10,
) -> SegmentationResult:
    """Segment an episode into subgoal phases using gripper signal.

    Args:
        gripper_signal: [T] array.
            If signal_source='action': action[:, 6], binary 1=open 0=closed.
            If signal_source='state': observation.state[:, 7], continuous ~-0.04(open) to ~-0.005(closed).
        subgoals: ordered subgoal list from instruction_parser.
        signal_source: 'action' (default, more reliable) or 'state'.
        threshold: binarization threshold (only used for 'state' source).
        median_filter_size: kernel size for median filtering.
        min_cycle_frames: minimum cycle duration in frames; shorter cycles are filtered as noise.

    Returns:
        SegmentationResult with per-frame labels.
    """
    T = len(gripper_signal)
    warn_msgs: List[str] = []

    # --- Step 1: Binarize & detect events ---
    g_bin = _binarize(gripper_signal, signal_source, threshold, median_filter_size)
    close_events, open_events = _detect_events(g_bin)
    cycles = _pair_cycles(close_events, open_events)

    # If the episode ends while gripper is still closed (no final open event),
    # the last close event is unpaired — add a virtual cycle ending at T.
    paired_closes = {c for c, _ in cycles}
    unpaired = [c for c in close_events if c not in paired_closes]
    if unpaired and g_bin[-1] == 1:
        # Only add the *last* unpaired close (the one still held at episode end)
        cycles.append((unpaired[-1], T))

    # Filter out short cycles (noise / gripper flickers)
    if min_cycle_frames > 0:
        cycles = [(c, o) for c, o in cycles if (o - c) >= min_cycle_frames]

    # --- Step 2: Classify subgoals into pre_manip, manip_pairs, post_manip ---
    pre_manip, manip_pairs, post_manip = _classify_subgoals(subgoals)
    n_expected = len(manip_pairs)
    n_detected = len(cycles)

    if n_detected != n_expected:
        warn_msgs.append(
            f"Cycle mismatch: detected {n_detected}, expected {n_expected}"
        )

    # --- Step 3: Assign frames ---
    if n_expected == 0 and len(pre_manip) == 0 and len(post_manip) == 0:
        # Edge case: only standalone subgoals with no manipulation
        # (shouldn't happen but handle gracefully)
        frame_subgoals = [subgoals[0].label()] * T if subgoals else ["unknown"] * T
    elif n_expected == 0:
        # No manipulation pairs — assign all frames to standalone subgoals
        all_standalone = pre_manip + post_manip
        frame_subgoals = _assign_frames_evenly(T, all_standalone, 0, T)
    else:
        frame_subgoals = _assign_with_cycles(
            T, pre_manip, manip_pairs, post_manip, cycles, n_expected, warn_msgs
        )

    # --- Step 4: Compute progress ---
    frame_progress = _compute_progress(frame_subgoals, T)

    return SegmentationResult(
        frame_subgoals=frame_subgoals,
        frame_progress=frame_progress,
        n_cycles_detected=n_detected,
        n_cycles_expected=n_expected,
        warnings=warn_msgs,
    )


def _binarize(
    gripper_signal: np.ndarray, signal_source: str, threshold: float, median_size: int
) -> np.ndarray:
    """Binarize gripper signal: 1=closed, 0=open.

    For 'action': action[:, 6] where 1=open, 0=closed → invert.
    For 'state': observation.state[:, 7] where ~-0.04=open, ~-0.005=closed → threshold.
    """
    if signal_source == "action":
        # action: 1=open, 0=closed → invert to get 1=closed, 0=open
        g_bin = (gripper_signal < 0.5).astype(np.int32)
    else:
        # state: more negative = more open
        g_bin = (gripper_signal > -threshold).astype(np.int32)
    if median_size > 1:
        g_bin = _median_filter_1d(g_bin, median_size)
    return g_bin


def _median_filter_1d(arr: np.ndarray, kernel: int) -> np.ndarray:
    """Simple 1D median filter (no scipy dependency)."""
    out = arr.copy()
    half = kernel // 2
    for i in range(half, len(arr) - half):
        out[i] = int(np.median(arr[i - half : i + half + 1]))
    return out


def _detect_events(g_bin: np.ndarray) -> Tuple[List[int], List[int]]:
    """Detect close (0→1) and open (1→0) transition indices."""
    close_events = []
    open_events = []
    for t in range(1, len(g_bin)):
        if g_bin[t - 1] == 0 and g_bin[t] == 1:
            close_events.append(t)
        elif g_bin[t - 1] == 1 and g_bin[t] == 0:
            open_events.append(t)
    return close_events, open_events


def _pair_cycles(
    close_events: List[int], open_events: List[int]
) -> List[Tuple[int, int]]:
    """Pair each close event with its nearest subsequent open event."""
    cycles = []
    open_idx = 0
    for ci in close_events:
        # Find next open event after this close
        while open_idx < len(open_events) and open_events[open_idx] <= ci:
            open_idx += 1
        if open_idx < len(open_events):
            cycles.append((ci, open_events[open_idx]))
            open_idx += 1
    return cycles


def _classify_subgoals(
    subgoals: List[Subgoal],
) -> Tuple[List[Subgoal], List[Tuple[Subgoal, Subgoal]], List[Subgoal]]:
    """Split subgoals into pre_manip, manip_pairs (pick+place), post_manip."""
    pre_manip: List[Subgoal] = []
    manip_pairs: List[Tuple[Subgoal, Subgoal]] = []
    post_manip: List[Subgoal] = []

    i = 0
    # Collect leading standalone subgoals
    while i < len(subgoals) and subgoals[i].type in STANDALONE_TYPES:
        pre_manip.append(subgoals[i])
        i += 1

    # Collect pick-place pairs and trailing standalone
    while i < len(subgoals):
        if subgoals[i].type == MANIP_PICK_TYPE:
            pick_sg = subgoals[i]
            if i + 1 < len(subgoals) and subgoals[i + 1].type == MANIP_PLACE_TYPE:
                place_sg = subgoals[i + 1]
                manip_pairs.append((pick_sg, place_sg))
                i += 2
            else:
                # Orphan pick — treat as standalone
                manip_pairs.append((pick_sg, pick_sg))
                i += 1
        elif subgoals[i].type in STANDALONE_TYPES:
            post_manip.append(subgoals[i])
            i += 1
        else:
            # Unexpected type order — skip
            post_manip.append(subgoals[i])
            i += 1

    return pre_manip, manip_pairs, post_manip


def _assign_with_cycles(
    T: int,
    pre_manip: List[Subgoal],
    manip_pairs: List[Tuple[Subgoal, Subgoal]],
    post_manip: List[Subgoal],
    cycles: List[Tuple[int, int]],
    n_expected: int,
    warn_msgs: List[str],
) -> List[str]:
    """Assign frames to subgoals using detected gripper cycles."""
    frame_subgoals = [""] * T

    # When more cycles detected than expected, use the LAST n_expected cycles.
    # Rationale: extra early cycles likely belong to pre_manip actions (e.g. turn_on_stove
    # triggers a gripper cycle before the actual pick-place).
    if len(cycles) > n_expected and n_expected > 0:
        selected_cycles = cycles[-n_expected:]
        discarded_cycles = cycles[:-n_expected]
    else:
        selected_cycles = cycles[:min(len(cycles), n_expected)]
        discarded_cycles = []
    n_usable = len(selected_cycles)

    # Determine boundaries
    if n_usable > 0:
        first_close = selected_cycles[0][0]
    else:
        # No usable cycles — divide evenly among all subgoals
        all_sgs = (
            pre_manip
            + [sg for pair in manip_pairs for sg in pair]
            + post_manip
        )
        return _assign_frames_evenly(T, all_sgs, 0, T)

    # --- pre_manip: [0, pre_end) ---
    # If there are discarded leading cycles (from pre_manip actions like turn_on),
    # use the open event of the last discarded cycle as pre_end.
    # This ensures the pick phase gets the approach frames between pre_manip and grasp.
    if discarded_cycles and pre_manip:
        pre_end = discarded_cycles[-1][1]  # open event of last discarded cycle
    elif pre_manip:
        # Share [0, first_close) between pre_manip subgoals and the first pick's
        # approach phase. Give pre_manip a proportional share so pick isn't starved.
        n_pre = len(pre_manip)
        pre_end = first_close * n_pre // (n_pre + 1)
    else:
        pre_end = 0

    if pre_manip and pre_end > 0:
        pre_labels = _assign_frames_evenly(T, pre_manip, 0, pre_end)
        for t in range(0, pre_end):
            frame_subgoals[t] = pre_labels[t]
    manip_start = pre_end

    # --- manipulation pairs: use cycles ---
    prev_boundary = manip_start
    for pair_idx in range(n_expected):
        pick_sg, place_sg = manip_pairs[pair_idx]
        if pair_idx < n_usable:
            close_t, open_t = selected_cycles[pair_idx]
            # Pick: from prev_boundary to close_t
            for t in range(prev_boundary, close_t):
                frame_subgoals[t] = pick_sg.label()
            # Place: from close_t to open_t
            for t in range(close_t, open_t):
                frame_subgoals[t] = place_sg.label()
            prev_boundary = open_t
        else:
            # No cycle for this pair — will be handled in post section
            break

    # --- post_manip / remaining: [last_open, T) ---
    post_start = prev_boundary

    # Collect remaining subgoals (unmatched manip pairs + post_manip)
    remaining_sgs: List[Subgoal] = []
    for pair_idx in range(n_usable, n_expected):
        pick_sg, place_sg = manip_pairs[pair_idx]
        remaining_sgs.extend([pick_sg, place_sg])
    remaining_sgs.extend(post_manip)

    if remaining_sgs and post_start < T:
        post_labels = _assign_frames_evenly(T, remaining_sgs, post_start, T)
        for t in range(post_start, T):
            frame_subgoals[t] = post_labels[t]
    elif post_start < T:
        # No remaining subgoals — extend last assigned label
        last_label = frame_subgoals[post_start - 1] if post_start > 0 else "unknown"
        for t in range(post_start, T):
            frame_subgoals[t] = last_label

    # Fill any gaps
    for t in range(T):
        if not frame_subgoals[t]:
            frame_subgoals[t] = "unknown"

    return frame_subgoals


def _assign_frames_evenly(
    T: int, subgoals: List[Subgoal], start: int, end: int
) -> List[str]:
    """Evenly distribute frame range [start, end) among subgoals."""
    labels = [""] * T
    n = len(subgoals)
    if n == 0 or start >= end:
        return labels
    span = end - start
    per_sg = max(1, span // n)
    for i, sg in enumerate(subgoals):
        sg_start = start + i * per_sg
        sg_end = start + (i + 1) * per_sg if i < n - 1 else end
        for t in range(sg_start, min(sg_end, end)):
            labels[t] = sg.label()
    return labels


def _compute_progress(frame_subgoals: List[str], T: int) -> List[str]:
    """Assign 'early'/'late' progress to each frame within its subgoal span."""
    progress = ["early"] * T
    if T == 0:
        return progress

    # Find contiguous subgoal spans
    spans: List[Tuple[int, int, str]] = []  # (start, end, label)
    current_label = frame_subgoals[0]
    span_start = 0
    for t in range(1, T):
        if frame_subgoals[t] != current_label:
            spans.append((span_start, t, current_label))
            current_label = frame_subgoals[t]
            span_start = t
    spans.append((span_start, T, current_label))

    for s, e, _ in spans:
        mid = (s + e) // 2
        for t in range(s, mid):
            progress[t] = "early"
        for t in range(mid, e):
            progress[t] = "late"

    return progress