"""Create train splits for OOD object position experiments on the 2view dataset.

Symlinks demos from /data/libero/ood_objpos_task0_2view into a left-half train split.
For test positions, we reuse the existing /data/libero/ood_objpos_splits/exp3_right_test.npy
(physical dx/dy positions are dataset-independent), so 1view vs 2view evals are fair.

Usage:
    python create_splits_2view.py \\
        --data_root /data/libero/ood_objpos_task0_2view \\
        --splits_root /data/libero/ood_objpos_splits_2view
"""
import argparse
import os
import shutil
from pathlib import Path
import numpy as np


def create_split(splits_root, split_name, data_root, demo_indices):
    split_dir = splits_root / f"{split_name}_train" / "libero_spatial" / "task_0"
    if split_dir.exists():
        shutil.rmtree(split_dir)
    split_dir.mkdir(parents=True)

    n_made = 0
    for new_idx, real_idx in enumerate(sorted(demo_indices)):
        src = data_root / "libero_spatial" / "task_0" / f"demo_{real_idx}"
        dst = split_dir / f"demo_{new_idx}"
        if src.exists():
            os.symlink(str(src.resolve()), str(dst))
            n_made += 1
    print(f"  {split_name}: requested={len(demo_indices)} created={n_made} -> {split_dir}")


def main():
    parser = argparse.ArgumentParser()
    parser.add_argument("--data_root", type=str, required=True)
    parser.add_argument("--splits_root", type=str, required=True)
    args = parser.parse_args()

    data_root = Path(args.data_root)
    splits_root = Path(args.splits_root)
    splits_root.mkdir(parents=True, exist_ok=True)

    meta = np.load(data_root / "libero_spatial" / "task_0" / "grid_meta.npz")
    dx_vals = meta["dx_vals"]
    dy_vals = meta["dy_vals"]
    Nx = len(dx_vals)
    Ny = len(dy_vals)
    print(f"Grid: {Nx}x{Ny}, dx=[{dx_vals[0]:.3f}, {dx_vals[-1]:.3f}], dy=[{dy_vals[0]:.3f}, {dy_vals[-1]:.3f}]")

    # Left-half training split (same indexing as create_splits.py)
    exp3_left = [i * Ny + j for i in range(Nx) for j in range(Ny // 2)]
    create_split(splits_root, "exp3_left", data_root, exp3_left)
    print("\nDone. Use existing /data/libero/ood_objpos_splits/exp3_right_test.npy for eval.")


if __name__ == "__main__":
    main()