import json
import os
import h5py
import numpy as np
from PIL import Image
from io import BytesIO
import argparse
from glob import glob
import random
from tqdm import tqdm

def decode_and_resize_images(image_bytes_array, size=256):
    resized = []
    for img_bytes in image_bytes_array:
        img_pil = Image.open(BytesIO(img_bytes.tobytes())).convert("RGB")
        img_resized = img_pil.resize((size, size), resample=Image.BICUBIC)
        resized.append(np.array(img_resized))
    return np.stack(resized)

def load_instruction(instruction_dir, episode_idx):
    json_path = os.path.join(instruction_dir, f"episode_{episode_idx}.json")
    if not os.path.exists(json_path):
        raise FileNotFoundError(f"Instruction file not found: {json_path}")
    with open(json_path, "r") as f:
        data = json.load(f)
    candidates = data.get("seen", []) + data.get("unseen", [])
    if not candidates:
        raise ValueError(f"No instructions found in {json_path}")
    return random.choice(candidates)
def process_one_episode(input_path, output_path, episode_idx, instruction_dir, resize_size=256):
    with h5py.File(input_path, "r") as f:
        action = f["joint_action/vector"][()]
        rel_action = np.zeros_like(action)
        rel_action[:-1] = action[1:] - action[:-1]
        rel_action[-1] = rel_action[-2]

        head = decode_and_resize_images(f["observation/head_camera/rgb"][()], size=resize_size)
        left = decode_and_resize_images(f["observation/left_camera/rgb"][()], size=resize_size)
        right = decode_and_resize_images(f["observation/right_camera/rgb"][()], size=resize_size)
        front = decode_and_resize_images(f["observation/front_camera/rgb"][()], size=resize_size)

    # 读取 instruction JSON
    json_path = os.path.join(instruction_dir, f"episode{episode_idx}.json")
    with open(json_path, "r") as f:
        inst_data = json.load(f)
    seen_list = inst_data.get("seen", [])
    unseen_list = inst_data.get("unseen", [])

    os.makedirs(os.path.dirname(output_path), exist_ok=True)
    with h5py.File(output_path, "w") as f:
        f.create_dataset("head_camera_image", data=head, dtype="uint8", chunks=(1, resize_size, resize_size, 3))
        f.create_dataset("left_wrist_image", data=left, dtype="uint8", chunks=(1, resize_size, resize_size, 3))
        f.create_dataset("right_wrist_image", data=right, dtype="uint8", chunks=(1, resize_size, resize_size, 3))
        f.create_dataset("low_cam_image", data=front, dtype="uint8", chunks=(1, resize_size, resize_size, 3))
        f.create_dataset("action", data=action)
        f.create_dataset("relative_action", data=rel_action)
        f.create_dataset("seen", data=np.array(seen_list, dtype=h5py.string_dtype(encoding="utf-8")))
        f.create_dataset("unseen", data=np.array(unseen_list, dtype=h5py.string_dtype(encoding="utf-8")))


def main(args):
    input_dir = args.dataset_path
    output_base = args.out_base_dir
    resize_size = args.img_resize_size
    instruction_dir = args.instruction_dir

    all_eps = sorted(glob(os.path.join(input_dir, "*.hdf5")))[:100]
    random.seed(42)
    random.shuffle(all_eps)

    n_val = int(len(all_eps) * args.percent_val)
    train_eps = all_eps[:-n_val]
    val_eps = all_eps[-n_val:]

    print(f"Total episodes: {len(all_eps)}")
    print(f"Train: {len(train_eps)}, Val: {len(val_eps)}")

    for split_name, split_eps in [("train", train_eps), ("val", val_eps)]:
        out_dir = os.path.join(output_base, split_name)
        os.makedirs(out_dir, exist_ok=True)
        for i, ep in enumerate(tqdm(split_eps, desc=f"Processing {split_name}")):
            ep_name = f"episode_{i}.hdf5"
            out_path = os.path.join(out_dir, ep_name)
            try:
                process_one_episode(ep, out_path, i, instruction_dir, resize_size=resize_size)
            except Exception as e:
                print(f"[ERROR] Failed to process {ep}: {e}")

if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    parser.add_argument("--dataset_path", type=str, required=True,
                        help="Path to RoboTwin hdf5 files")
    parser.add_argument("--out_base_dir", type=str, required=True,
                        help="Output dir for processed OpenVLA-compatible dataset")
    parser.add_argument("--instruction_dir", type=str, required=True,
                        help="Directory containing episode_*.json instruction files")
    parser.add_argument("--percent_val", type=float, default=0.05,
                        help="Fraction of data to use as validation")
    parser.add_argument("--img_resize_size", type=int, default=256,
                        help="Final size for RGB images")
    args = parser.parse_args()
    main(args)


"""
python preprocess_aloha.py   --dataset_path /mnt/data/VLA_flowmatching/RoboTwin/data/place_object_scale/demo_randomized/data   --out_base_dir /mnt/data/VLA_flowmatching/RoboTwin/data/place_object_scale/processed_openvla/   --percent_val 0.05 --instruction_dir /mnt/data/VLA_flowmatching/RoboTwin/data/place_object_scale/demo_randomized/instructions
"""