#!/usr/bin/env python

# Copyright 2025 The HuggingFace Inc. team. All rights reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

"""
Mirror a bimanual robot dataset using SLURM for distributed video processing.

This script creates a mirrored version of a dataset where:
1. Left and right arm observations/actions are swapped
2. Joint values are inverted according to a mirroring mask
3. Video frames are horizontally flipped (parallelized via SLURM)

Example usage:
```shell
# SLURM execution
python examples/port_datasets/slurm_mirror_dataset.py \
    --repo-id pepijn/openarm_bimanual \
    --output-repo-id pepijn/openarm_bimanual_mirrored \
    --logs-dir /fsx/user/logs \
    --partition hopper-cpu

# Local execution (for debugging)
python examples/port_datasets/slurm_mirror_dataset.py \
    --repo-id pepijn/openarm_bimanual \
    --output-repo-id pepijn/openarm_bimanual_mirrored \
    --slurm 0
```
"""

import argparse
import json
import logging
import subprocess
from pathlib import Path

import numpy as np
import pandas as pd
from datatrove.executor import LocalPipelineExecutor
from datatrove.executor.slurm import SlurmPipelineExecutor
from datatrove.pipeline.base import PipelineStep

logger = logging.getLogger(__name__)

OPENARM_MIRRORING_MASK = {
    "joint_1": -1,
    "joint_2": -1,
    "joint_3": -1,
    "joint_4": 1,
    "joint_5": -1,
    "joint_6": -1,
    "joint_7": -1,
    "gripper": 1,
}


def get_mirroring_mask(robot_type: str) -> dict[str, int]:
    if robot_type in ["bi_openarm_follower", "openarm_follower", "bi_openarms_follower", "openarms_follower"]:
        return OPENARM_MIRRORING_MASK
    raise ValueError(f"Unknown robot type: {robot_type}. Add a mirroring mask for this robot.")


def swap_left_right_name(name: str) -> str:
    result = name.replace("left_", "LEFT_PLACEHOLDER_")
    result = result.replace("right_", "left_")
    result = result.replace("LEFT_PLACEHOLDER_", "right_")
    return result


def mirror_feature_names(names: list[str]) -> tuple[list[str], dict[int, int]]:
    mirrored_names = [swap_left_right_name(n) for n in names]
    old_to_new_idx = {}
    for old_idx, old_name in enumerate(names):
        new_name = swap_left_right_name(old_name)
        new_idx = mirrored_names.index(new_name)
        old_to_new_idx[old_idx] = new_idx
    return mirrored_names, old_to_new_idx


def apply_mirroring_mask(value: float, feature_name: str, mirroring_mask: dict[str, int]) -> float:
    name_without_prefix = feature_name.split("_", 1)[1] if "_" in feature_name else feature_name
    joint_name = name_without_prefix.split(".")[0]
    if joint_name in mirroring_mask:
        return value * mirroring_mask[joint_name]
    return value


def mirror_array(array: np.ndarray, names: list[str], mirroring_mask: dict[str, int]) -> np.ndarray:
    mirrored_names, idx_mapping = mirror_feature_names(names)
    result = np.zeros_like(array)
    for old_idx, new_idx in idx_mapping.items():
        new_name = mirrored_names[new_idx]
        value = array[old_idx]
        mirrored_value = apply_mirroring_mask(value, new_name, mirroring_mask)
        result[new_idx] = mirrored_value
    return result


def flip_video_frames(input_path: Path, output_path: Path, fps: float, vcodec: str = "libsvtav1"):
    output_path.parent.mkdir(parents=True, exist_ok=True)
    cmd = [
        "ffmpeg", "-y", "-i", str(input_path),
        "-vf", "hflip",
        "-c:v", vcodec,
        "-g", "2",
        "-crf", "30",
        "-r", str(int(fps)),
        "-pix_fmt", "yuv420p",
        "-loglevel", "error",
    ]
    if vcodec == "libsvtav1":
        cmd.extend(["-preset", "12"])
    cmd.append(str(output_path))
    result = subprocess.run(cmd, capture_output=True, text=True)
    if result.returncode != 0:
        raise RuntimeError(f"FFmpeg failed: {result.stderr}")


def video_is_valid(path: Path) -> bool:
    """Check if a video file exists and is valid."""
    if not path.exists():
        return False
    try:
        result = subprocess.run(
            ["ffprobe", "-v", "error", "-select_streams", "v:0",
             "-show_entries", "stream=nb_frames", "-of", "csv=p=0", str(path)],
            capture_output=True, text=True, timeout=30
        )
        return result.returncode == 0 and result.stdout.strip().isdigit()
    except Exception:
        return False


class MirrorVideos(PipelineStep):
    """Pipeline step that mirrors video files for assigned episodes."""
    
    def __init__(
        self,
        repo_id: str,
        output_repo_id: str,
        root: str | None = None,
        output_root: str | None = None,
        vcodec: str = "libsvtav1",
    ):
        super().__init__()
        self.repo_id = repo_id
        self.output_repo_id = output_repo_id
        self.root = Path(root) if root else None
        self.output_root = Path(output_root) if output_root else None
        self.vcodec = vcodec

    def run(self, data=None, rank: int = 0, world_size: int = 1):
        from datasets.utils.tqdm import disable_progress_bars
        from lerobot.datasets.lerobot_dataset import LeRobotDataset
        from lerobot.utils.constants import HF_LEROBOT_HOME
        from lerobot.utils.utils import init_logging

        init_logging()
        disable_progress_bars()

        dataset = LeRobotDataset(self.repo_id, root=self.root)
        output_root = self.output_root or (HF_LEROBOT_HOME / self.output_repo_id)

        if not dataset.meta.video_keys:
            logger.info(f"Rank {rank}: No videos to process")
            return

        video_tasks = []
        for old_video_key in dataset.meta.video_keys:
            new_video_key = swap_left_right_name(old_video_key)
            for ep_idx in range(dataset.meta.total_episodes):
                try:
                    src_path = dataset.root / dataset.meta.get_video_file_path(ep_idx, old_video_key)
                    dst_relative = dataset.meta.get_video_file_path(ep_idx, old_video_key)
                    dst_relative_str = str(dst_relative).replace(old_video_key, new_video_key)
                    dst_path = output_root / dst_relative_str
                    if src_path.exists():
                        video_tasks.append((src_path, dst_path, ep_idx, old_video_key))
                except KeyError:
                    continue

        # Distribute tasks across workers
        my_tasks = [t for i, t in enumerate(video_tasks) if i % world_size == rank]
        logger.info(f"Rank {rank}/{world_size}: Processing {len(my_tasks)}/{len(video_tasks)} videos")

        for src_path, dst_path, ep_idx, video_key in my_tasks:
            if video_is_valid(dst_path):
                logger.info(f"Rank {rank}: Skipping {dst_path.name} (already done)")
                continue
            logger.info(f"Rank {rank}: Processing {src_path.name} -> {dst_path.name}")
            flip_video_frames(src_path, dst_path, dataset.meta.fps, self.vcodec)


class MirrorDataAndMetadata(PipelineStep):
    """Pipeline step that mirrors parquet data and metadata (runs once on rank 0)."""
    
    def __init__(
        self,
        repo_id: str,
        output_repo_id: str,
        root: str | None = None,
        output_root: str | None = None,
    ):
        super().__init__()
        self.repo_id = repo_id
        self.output_repo_id = output_repo_id
        self.root = Path(root) if root else None
        self.output_root = Path(output_root) if output_root else None

    def run(self, data=None, rank: int = 0, world_size: int = 1):
        if rank != 0:
            return

        from datasets.utils.tqdm import disable_progress_bars
        from lerobot.datasets.lerobot_dataset import LeRobotDataset, LeRobotDatasetMetadata
        from lerobot.datasets.utils import DATA_DIR, DEFAULT_DATA_PATH, write_info, write_stats, write_tasks
        from lerobot.utils.constants import HF_LEROBOT_HOME
        from lerobot.utils.utils import init_logging

        init_logging()
        disable_progress_bars()

        dataset = LeRobotDataset(self.repo_id, root=self.root)
        output_root = self.output_root or (HF_LEROBOT_HOME / self.output_repo_id)
        
        done_marker = output_root / ".data_mirrored"
        if done_marker.exists():
            logger.info("Data and metadata already mirrored, skipping")
            return

        robot_type = dataset.meta.robot_type or "bi_openarms_follower"
        mirroring_mask = get_mirroring_mask(robot_type)

        mirrored_features = {}
        for key, feat in dataset.meta.features.items():
            new_key = swap_left_right_name(key)
            new_feat = feat.copy()
            if "names" in new_feat and new_feat["names"]:
                new_feat["names"] = [swap_left_right_name(n) for n in new_feat["names"]]
            mirrored_features[new_key] = new_feat

        new_meta = LeRobotDatasetMetadata.create(
            repo_id=self.output_repo_id,
            fps=dataset.meta.fps,
            features=mirrored_features,
            robot_type=dataset.meta.robot_type,
            root=output_root,
            use_videos=len(dataset.meta.video_keys) > 0,
        )

        if dataset.meta.tasks is not None:
            write_tasks(dataset.meta.tasks, new_meta.root)

        # Mirror parquet data
        data_dir = dataset.root / DATA_DIR
        parquet_files = sorted(data_dir.glob("*/*.parquet"))
        action_names = dataset.meta.features.get("action", {}).get("names", [])
        state_names = dataset.meta.features.get("observation.state", {}).get("names", [])

        for src_path in parquet_files:
            df = pd.read_parquet(src_path).reset_index(drop=True)
            relative_path = src_path.relative_to(dataset.root)
            chunk_dir = relative_path.parts[1]
            file_name = relative_path.parts[2]
            chunk_idx = int(chunk_dir.split("-")[1])
            file_idx = int(file_name.split("-")[1].split(".")[0])

            if "action" in df.columns and action_names:
                actions = np.stack(df["action"].values)
                mirrored_actions = np.array([mirror_array(row, action_names, mirroring_mask) for row in actions])
                df["action"] = list(mirrored_actions)

            if "observation.state" in df.columns and state_names:
                states = np.stack(df["observation.state"].values)
                mirrored_states = np.array([mirror_array(row, state_names, mirroring_mask) for row in states])
                df["observation.state"] = list(mirrored_states)

            dst_path = new_meta.root / DEFAULT_DATA_PATH.format(chunk_index=chunk_idx, file_index=file_idx)
            dst_path.parent.mkdir(parents=True, exist_ok=True)
            df.to_parquet(dst_path, index=False)

        # Copy episodes metadata
        episodes_dir = dataset.root / "meta/episodes"
        dst_episodes_dir = new_meta.root / "meta/episodes"
        if episodes_dir.exists():
            dst_episodes_dir.mkdir(parents=True, exist_ok=True)
            for src_parquet in episodes_dir.glob("*/*.parquet"):
                df = pd.read_parquet(src_parquet)
                columns_to_rename = {}
                for col in df.columns:
                    if col.startswith("videos/"):
                        parts = col.split("/")
                        if len(parts) >= 2:
                            video_key = parts[1]
                            new_video_key = swap_left_right_name(video_key)
                            new_col = col.replace(f"videos/{video_key}/", f"videos/{new_video_key}/")
                            columns_to_rename[col] = new_col
                if columns_to_rename:
                    df = df.rename(columns=columns_to_rename)
                dst_parquet = dst_episodes_dir / src_parquet.relative_to(episodes_dir)
                dst_parquet.parent.mkdir(parents=True, exist_ok=True)
                df.to_parquet(dst_parquet, index=False)

        new_meta.info.update({
            "total_episodes": dataset.meta.total_episodes,
            "total_frames": dataset.meta.total_frames,
            "total_tasks": dataset.meta.total_tasks,
            "total_videos": dataset.meta.total_videos,
            "total_chunks": dataset.meta.total_chunks,
        })
        write_info(new_meta.info, new_meta.root)

        if dataset.meta.stats is not None:
            mirrored_stats = mirror_stats(dataset.meta.stats)
            write_stats(mirrored_stats, new_meta.root)

        done_marker.touch()
        logger.info(f"Data and metadata mirrored to {output_root}")


def mirror_stats(stats: dict) -> dict:
    mirrored = {}
    for key, value in stats.items():
        new_key = swap_left_right_name(key)
        if isinstance(value, dict):
            mirrored[new_key] = mirror_stats(value)
        else:
            mirrored[new_key] = value
    return mirrored


def get_num_video_tasks(repo_id: str, root: str | None = None) -> int:
    """Count total video files to process."""
    from lerobot.datasets.lerobot_dataset import LeRobotDataset
    dataset = LeRobotDataset(repo_id, root=root)
    count = 0
    for video_key in dataset.meta.video_keys:
        for ep_idx in range(dataset.meta.total_episodes):
            try:
                src_path = dataset.root / dataset.meta.get_video_file_path(ep_idx, video_key)
                if src_path.exists():
                    count += 1
            except KeyError:
                continue
    return count


def make_mirror_executor(
    repo_id: str,
    output_repo_id: str,
    root: str | None,
    output_root: str | None,
    vcodec: str,
    job_name: str,
    logs_dir: Path,
    workers: int,
    partition: str,
    cpus_per_task: int,
    mem_per_cpu: str,
    time_limit: str,
    slurm: bool = True,
):
    num_tasks = get_num_video_tasks(repo_id, root) if slurm else 1
    num_tasks = max(1, num_tasks)

    kwargs = {
        "pipeline": [
            MirrorDataAndMetadata(repo_id, output_repo_id, root, output_root),
            MirrorVideos(repo_id, output_repo_id, root, output_root, vcodec),
        ],
        "logging_dir": str(logs_dir / job_name),
    }

    if slurm:
        kwargs.update({
            "job_name": job_name,
            "tasks": num_tasks,
            "workers": min(workers, num_tasks),
            "time": time_limit,
            "partition": partition,
            "cpus_per_task": cpus_per_task,
            "sbatch_args": {
                "mem-per-cpu": mem_per_cpu,
                "requeue": True,
                "signal": "USR1@30",
            },
        })
        return SlurmPipelineExecutor(**kwargs)
    else:
        kwargs.update({"tasks": 1, "workers": 1})
        return LocalPipelineExecutor(**kwargs)


def main():
    logging.basicConfig(level=logging.INFO, format="%(asctime)s - %(name)s - %(levelname)s - %(message)s")
    parser = argparse.ArgumentParser(description="Mirror a bimanual robot dataset using SLURM")
    parser.add_argument("--repo-id", type=str, required=True, help="Source dataset repo_id")
    parser.add_argument("--output-repo-id", type=str, required=True, help="Output dataset repo_id")
    parser.add_argument("--root", type=str, default=None, help="Source dataset root directory")
    parser.add_argument("--output-root", type=str, default=None, help="Output dataset root directory")
    parser.add_argument("--vcodec", type=str, default="libsvtav1", help="Video codec")
    parser.add_argument("--logs-dir", type=Path, default=Path("logs"), help="Directory for datatrove logs")
    parser.add_argument("--job-name", type=str, default="mirror_dataset", help="SLURM job name")
    parser.add_argument("--slurm", type=int, default=1, help="Use SLURM (1) or local (0)")
    parser.add_argument("--workers", type=int, default=64, help="Number of SLURM workers")
    parser.add_argument("--partition", type=str, default="hopper-cpu", help="SLURM partition")
    parser.add_argument("--cpus-per-task", type=int, default=4, help="CPUs per task")
    parser.add_argument("--mem-per-cpu", type=str, default="2G", help="Memory per CPU")
    parser.add_argument("--time-limit", type=str, default="04:00:00", help="SLURM time limit")

    args = parser.parse_args()

    executor = make_mirror_executor(
        repo_id=args.repo_id,
        output_repo_id=args.output_repo_id,
        root=args.root,
        output_root=args.output_root,
        vcodec=args.vcodec,
        job_name=args.job_name,
        logs_dir=args.logs_dir,
        workers=args.workers,
        partition=args.partition,
        cpus_per_task=args.cpus_per_task,
        mem_per_cpu=args.mem_per_cpu,
        time_limit=args.time_limit,
        slurm=args.slurm == 1,
    )
    executor.run()


if __name__ == "__main__":
    main()