import os
import argparse

ap = argparse.ArgumentParser()
ap.add_argument("task_name")
ap.add_argument("task_config")
ap.add_argument("index", type=int, help="Index of the episode to delete")
args = ap.parse_args()

target_path = os.path.join("data", args.task_name, args.task_config)
txt_path = os.path.join(target_path, "seed.txt")
dir_path = os.path.join(target_path, "data")
replace_id = args.index

with open(txt_path, "r") as f:
    nums = list(map(int, f.read().split()))

final_id = len(nums) - 1

last_seed = nums[-1]

nums[replace_id] = last_seed
nums.pop()

with open(txt_path, "w") as f:
    f.write(" ".join(map(str, nums)))

target_file = os.path.join(dir_path, f"episode{replace_id}.pkl")
if os.path.exists(target_file):
    os.remove(target_file)

last_file = os.path.join(dir_path, f"episode{final_id}.pkl")
if os.path.exists(last_file):
    os.rename(last_file, target_file)
