AffordanceNet/phaseaware_data_construction/phase_annotate.py

"""Phase annotate: CLI script to add all phase/affordance columns to parquet files in a single pass.

Produces 8 columns:
  - phase.subgoal, phase.progress          (from gripper segmentation)
  - affordance.type, affordance.granularity (rule-derived from above)
  - phase.subgoal_object, phase.subgoal_target, phase.subgoal_descriptor, phase.subgoal_relation
    (structured fields resolved from parsed instructions)
"""

from __future__ import annotations

import argparse
import json
import logging
import os
import shutil
from collections import Counter, defaultdict
from multiprocessing import Pool
from pathlib import Path
from typing import Dict, List, Tuple

import numpy as np
import pandas as pd

from phaseaware_data_construction.gripper_segmenter import SegmentationResult, segment_episode
from phaseaware_data_construction.instruction_parser import Subgoal, parse_instruction
from phaseaware_data_construction.affordance_deriver import derive_affordance_columns
from phaseaware_data_construction.subgoal_columns import build_label_to_fields, resolve_frame_fields

logging.basicConfig(
    level=logging.INFO, format="%(asctime)s [%(levelname)s] %(message)s"
)
logger = logging.getLogger(__name__)

# Gripper dimension indices (from modality.json)
GRIPPER_STATE_DIM = 7   # observation.state[:, 7]
GRIPPER_ACTION_DIM = 6  # action[:, 6]


def load_task_map(src_dir: Path) -> Dict[int, str]:
    """Load task_index → instruction mapping from meta/tasks.jsonl."""
    tasks_path = src_dir / "meta" / "tasks.jsonl"
    task_map: Dict[int, str] = {}
    with open(tasks_path) as f:
        for line in f:
            d = json.loads(line.strip())
            task_map[d["task_index"]] = d["task"]
    return task_map


def parse_all_tasks(task_map: Dict[int, str]) -> Dict[int, List[Subgoal]]:
    """Pre-parse all task instructions into subgoal sequences."""
    parsed: Dict[int, List[Subgoal]] = {}
    for idx, instruction in task_map.items():
        parsed[idx] = parse_instruction(instruction)
    return parsed


def find_episode_parquets(src_dir: Path) -> List[Path]:
    """Find all episode parquet files under data/."""
    data_dir = src_dir / "data"
    parquets = sorted(data_dir.rglob("episode_*.parquet"))
    return parquets


def process_single_episode(args: Tuple) -> dict:
    """Process a single episode parquet. Designed for multiprocessing."""
    parquet_path, subgoal_map, signal_source, threshold, median_filter_size = args
    parquet_path = Path(parquet_path)

    try:
        df = pd.read_parquet(parquet_path)
        state = np.array(df["observation.state"].tolist())
        task_idx = int(df["task_index"].iloc[0])
        episode_idx = int(df["episode_index"].iloc[0])

        # Extract gripper signal based on source
        if signal_source == "action":
            action = np.array(df["action"].tolist())
            gripper_signal = action[:, GRIPPER_ACTION_DIM]
        else:
            gripper_signal = state[:, GRIPPER_STATE_DIM]

        subgoals = subgoal_map.get(task_idx)
        if subgoals is None:
            return {
                "path": str(parquet_path),
                "episode": episode_idx,
                "status": "error",
                "error": f"Unknown task_index {task_idx}",
            }

        result = segment_episode(
            gripper_signal, subgoals,
            signal_source=signal_source,
            threshold=threshold,
            median_filter_size=median_filter_size,
        )

        # Derive affordance columns (rule-based from subgoal + progress)
        aff_types, aff_grans = derive_affordance_columns(
            result.frame_subgoals, result.frame_progress
        )

        # Derive structured subgoal fields (object, target, descriptor, relation)
        label_map = build_label_to_fields(subgoals)
        sg_objects, sg_targets, sg_descriptors, sg_relations = resolve_frame_fields(
            result.frame_subgoals, label_map
        )

        return {
            "path": str(parquet_path),
            "episode": episode_idx,
            "task_index": task_idx,
            "status": "ok",
            "n_frames": len(df),
            "n_cycles_detected": result.n_cycles_detected,
            "n_cycles_expected": result.n_cycles_expected,
            "warnings": result.warnings,
            "frame_subgoals": result.frame_subgoals,
            "frame_progress": result.frame_progress,
            "aff_types": aff_types,
            "aff_grans": aff_grans,
            "sg_objects": sg_objects,
            "sg_targets": sg_targets,
            "sg_descriptors": sg_descriptors,
            "sg_relations": sg_relations,
        }
    except Exception as e:
        return {
            "path": str(parquet_path),
            "episode": -1,
            "status": "error",
            "error": str(e),
        }


def write_annotated_parquet(
    src_path: Path, output_path: Path, result: dict
):
    """Read src parquet, add all 8 phase/affordance columns, write to output."""
    df = pd.read_parquet(src_path)
    df["phase.subgoal"] = result["frame_subgoals"]
    df["phase.progress"] = result["frame_progress"]
    df["affordance.type"] = result["aff_types"]
    df["affordance.granularity"] = result["aff_grans"]
    df["phase.subgoal_object"] = result["sg_objects"]
    df["phase.subgoal_target"] = result["sg_targets"]
    df["phase.subgoal_descriptor"] = result["sg_descriptors"]
    df["phase.subgoal_relation"] = result["sg_relations"]
    output_path.parent.mkdir(parents=True, exist_ok=True)
    df.to_parquet(output_path, index=False)


def validate_output(output_dir: Path, subgoal_map: Dict[int, List[Subgoal]]):
    """Validate the annotated output parquets."""
    parquets = find_episode_parquets(output_dir)
    if not parquets:
        logger.error("No parquets found in output directory")
        return

    total_frames = 0
    total_episodes = 0
    suspicious = []
    order_mismatches = []
    sg_frame_counts: Counter = Counter()
    progress_counts: Counter = Counter()
    type_counts: Counter = Counter()
    gran_counts: Counter = Counter()
    sg_per_episode: defaultdict = defaultdict(list)

    for p in parquets:
        df = pd.read_parquet(p)
        total_episodes += 1
        total_frames += len(df)

        if "phase.subgoal" not in df.columns or "phase.progress" not in df.columns:
            logger.error(f"Missing phase columns in {p}")
            continue

        # Check all 8 columns exist
        all_cols = [
            "phase.subgoal", "phase.progress",
            "affordance.type", "affordance.granularity",
            "phase.subgoal_object", "phase.subgoal_target",
            "phase.subgoal_descriptor", "phase.subgoal_relation",
        ]
        missing_cols = [c for c in all_cols if c not in df.columns]
        if missing_cols:
            logger.warning(f"{p}: missing columns {missing_cols}")

        null_sg = df["phase.subgoal"].isna().sum()
        null_pr = df["phase.progress"].isna().sum()
        if null_sg > 0 or null_pr > 0:
            logger.warning(f"{p}: {null_sg} null subgoals, {null_pr} null progress")

        # Count subgoal types and check order
        episode_sgs = df["phase.subgoal"].tolist()
        episode_pr = df["phase.progress"].tolist()

        # Extract unique subgoals in order
        seen_order = []
        prev = None
        for sg in episode_sgs:
            if sg != prev:
                seen_order.append(sg)
                prev = sg

        ep_idx = int(df["episode_index"].iloc[0])
        task_idx = int(df["task_index"].iloc[0])
        sg_per_episode[task_idx].append(len(seen_order))

        # --- Semantic order check against parser expectation ---
        if task_idx in subgoal_map:
            expected_labels = [sg.label() for sg in subgoal_map[task_idx]]
            # Check 1: every expected subgoal should appear in seen_order
            missing = [lbl for lbl in expected_labels if lbl not in seen_order]
            # Check 2: relative order should be preserved (allow extra labels but
            # expected labels must appear in the same order)
            ei = 0
            out_of_order = False
            for lbl in seen_order:
                if ei < len(expected_labels) and lbl == expected_labels[ei]:
                    ei += 1
            if ei < len(expected_labels):
                out_of_order = True

            if missing or out_of_order:
                order_mismatches.append((
                    p.name, ep_idx, task_idx,
                    expected_labels, seen_order,
                    missing, out_of_order,
                ))

        for sg in episode_sgs:
            sg_frame_counts[sg] += 1
        for pr in episode_pr:
            progress_counts[pr] += 1
        if "affordance.type" in df.columns:
            for t in df["affordance.type"]:
                type_counts[t] += 1
        if "affordance.granularity" in df.columns:
            for g in df["affordance.granularity"]:
                gran_counts[g] += 1

        # Check for suspicious episodes
        sg_counts_ep = Counter(episode_sgs)
        n = len(df)
        for sg, cnt in sg_counts_ep.items():
            if cnt < 5:
                suspicious.append((p.name, ep_idx, sg, cnt, "too_short"))
            if cnt > 0.8 * n and len(sg_counts_ep) > 1:
                suspicious.append((p.name, ep_idx, sg, cnt, "too_long"))

    # Print summary
    logger.info(f"\n{'='*60}")
    logger.info(f"Validation Summary")
    logger.info(f"{'='*60}")
    logger.info(f"Total episodes: {total_episodes}")
    logger.info(f"Total frames: {total_frames}")
    logger.info(f"\nProgress distribution:")
    for pr, cnt in progress_counts.most_common():
        logger.info(f"  {pr}: {cnt} ({cnt/total_frames*100:.1f}%)")

    logger.info(f"\nSubgoal type frame counts (top 20):")
    for sg, cnt in sg_frame_counts.most_common(20):
        logger.info(f"  {sg}: {cnt}")

    if type_counts:
        logger.info(f"\naffordance.type distribution:")
        for t, cnt in type_counts.most_common():
            logger.info(f"  {t}: {cnt} ({cnt/total_frames*100:.1f}%)")

    if gran_counts:
        logger.info(f"\naffordance.granularity distribution:")
        for g, cnt in gran_counts.most_common():
            logger.info(f"  {g}: {cnt} ({cnt/total_frames*100:.1f}%)")

    # Semantic order check results
    if order_mismatches:
        logger.warning(f"\nSubgoal order mismatches ({len(order_mismatches)}):")
        for name, ep, tidx, expected, seen, missing, ooo in order_mismatches:
            logger.warning(f"  {name} ep={ep} task={tidx}:")
            logger.warning(f"    expected : {expected}")
            logger.warning(f"    seen     : {seen}")
            if missing:
                logger.warning(f"    missing  : {missing}")
            if ooo:
                logger.warning(f"    out-of-order: expected labels not in correct relative order")
    else:
        logger.info(f"\nAll episodes match expected subgoal order.")

    if suspicious:
        logger.warning(f"\nSuspicious episodes ({len(suspicious)}):")
        for name, ep, sg, cnt, reason in suspicious[:20]:
            logger.warning(f"  {name} ep={ep}: {sg} has {cnt} frames ({reason})")
    else:
        logger.info(f"\nNo suspicious episodes found.")


def main():
    parser = argparse.ArgumentParser(
        description="Add all phase/affordance columns (8 total) to LeRobot parquet files."
    )
    parser.add_argument("--src_dir", type=str, required=True, help="Source dataset directory")
    parser.add_argument("--output_dir", type=str, required=True, help="Output directory")
    parser.add_argument("--num_workers", type=int, default=1, help="Number of parallel workers")
    parser.add_argument("--signal_source", type=str, default="action",
                        choices=["action", "state"],
                        help="Gripper signal source: 'action' (binary, more reliable) or 'state' (continuous)")
    parser.add_argument("--gripper_threshold", type=float, default=0.02,
                        help="Gripper binarization threshold (state dim 7: > -threshold → closed)")
    parser.add_argument("--median_filter_size", type=int, default=5,
                        help="Median filter kernel size")
    parser.add_argument("--dry_run", action="store_true",
                        help="Don't write files, only print stats")
    parser.add_argument("--validate", action="store_true",
                        help="Validate existing output instead of annotating")
    args = parser.parse_args()

    src_dir = Path(args.src_dir)
    output_dir = Path(args.output_dir)

    # --- Validate mode ---
    if args.validate:
        task_map = load_task_map(src_dir)
        subgoal_map = parse_all_tasks(task_map)
        validate_output(output_dir, subgoal_map)
        return

    # --- Annotate mode ---
    logger.info(f"Source: {src_dir}")
    logger.info(f"Output: {output_dir}")

    # Load and parse tasks
    task_map = load_task_map(src_dir)
    subgoal_map = parse_all_tasks(task_map)
    logger.info(f"Parsed {len(task_map)} tasks:")
    for idx, sgs in subgoal_map.items():
        logger.info(f"  task {idx}: {task_map[idx]} → {[sg.label() for sg in sgs]}")

    # Find episodes
    parquets = find_episode_parquets(src_dir)
    logger.info(f"Found {len(parquets)} episode parquets")

    if not parquets:
        logger.error("No episode parquets found!")
        return

    # Prepare arguments for workers
    work_args = [
        (str(p), subgoal_map, args.signal_source, args.gripper_threshold, args.median_filter_size)
        for p in parquets
    ]

    # Process
    if args.num_workers > 1:
        with Pool(args.num_workers) as pool:
            results = pool.map(process_single_episode, work_args)
    else:
        results = [process_single_episode(a) for a in work_args]

    # Summarize and write
    ok_count = 0
    err_count = 0
    warn_count = 0
    cycle_match = 0
    cycle_mismatch = 0

    for r in results:
        if r["status"] == "error":
            err_count += 1
            logger.error(f"Episode {r['episode']}: {r.get('error', 'unknown error')}")
            continue
        ok_count += 1
        if r["warnings"]:
            warn_count += 1
            for w in r["warnings"]:
                logger.warning(f"Episode {r['episode']}: {w}")
        if r["n_cycles_detected"] == r["n_cycles_expected"]:
            cycle_match += 1
        else:
            cycle_mismatch += 1

    logger.info(f"\n{'='*60}")
    logger.info(f"Processing Summary")
    logger.info(f"{'='*60}")
    logger.info(f"OK: {ok_count}, Errors: {err_count}, With warnings: {warn_count}")
    logger.info(f"Cycle match: {cycle_match}, Cycle mismatch: {cycle_mismatch}")
    logger.info(f"Match rate: {cycle_match/(ok_count or 1)*100:.1f}%")

    if args.dry_run:
        logger.info("Dry run — no files written.")
        return

    # --- Write output ---
    logger.info(f"\nWriting annotated parquets to {output_dir}")

    # Copy meta directory
    src_meta = src_dir / "meta"
    dst_meta = output_dir / "meta"
    if src_meta.exists():
        if dst_meta.exists():
            shutil.rmtree(dst_meta)
        shutil.copytree(src_meta, dst_meta)
        logger.info(f"Copied meta/ directory")

    # Symlink video directory
    src_video = src_dir / "video"
    dst_video = output_dir / "video"
    if src_video.exists() and not dst_video.exists():
        dst_video.parent.mkdir(parents=True, exist_ok=True)
        os.symlink(src_video.resolve(), dst_video)
        logger.info(f"Symlinked video/ directory")

    # Write annotated parquets
    written = 0
    for r in results:
        if r["status"] != "ok":
            continue
        src_path = Path(r["path"])
        rel_path = src_path.relative_to(src_dir)
        out_path = output_dir / rel_path
        write_annotated_parquet(src_path, out_path, r)
        written += 1

    logger.info(f"Written {written} annotated parquets")


if __name__ == "__main__":
    main()