"""Capture first and last (success) frames from front camera for both sim tasks."""
import os
import time
import numpy as np
import cv2
from roboimi.envs.double_pos_ctrl_env import make_sim_env
from roboimi.demos.diana_policy import TestPickAndTransferPolicy
from roboimi.demos.diana_air_insert_policy import TestAirInsertPolicy
from roboimi.utils.act_ex_utils import sample_transfer_pose, sample_air_insert_socket_peg_state

OUTPUT_DIR = "sim_frames"
IMG_SIZE = (256, 256)
os.makedirs(OUTPUT_DIR, exist_ok=True)


def capture_task(task_name, policy, task_state, episode_len):
    env = make_sim_env(task_name, headless=True)
    env.reset(task_state)
    time.sleep(1)

    env._update_camera_images_sync()
    first_frame = cv2.resize(env.front, IMG_SIZE)

    last_success_frame = None
    max_reward = 0

    for step in range(episode_len):
        action = policy.predict(task_state, step)
        env.step(action)
        env._update_camera_images_sync()
        if env.rew is not None and env.rew >= max_reward:
            max_reward = env.rew
            if env.rew == env.max_reward:
                last_success_frame = cv2.resize(env.front, IMG_SIZE)

    if last_success_frame is None:
        last_success_frame = cv2.resize(env.front, IMG_SIZE)
        print(f"  [WARN] {task_name}: max_reward={max_reward}/{env.max_reward}, using last frame")
    else:
        print(f"  {task_name}: success! max_reward={max_reward}")

    return first_frame, last_success_frame


def main():
    # Transfer task (episode_len=700, policy covers full range)
    print("Running sim_transfer...")
    task_state_transfer = sample_transfer_pose()
    policy_transfer = TestPickAndTransferPolicy(inject_noise=False)
    first_t, last_t = capture_task('sim_transfer', policy_transfer, task_state_transfer, 700)
    cv2.imwrite(os.path.join(OUTPUT_DIR, "transfer_front_first.png"), first_t)
    cv2.imwrite(os.path.join(OUTPUT_DIR, "transfer_front_last.png"), last_t)

    # Air insert / peg-in-socket (policy EPISODE_END_T=600)
    print("Running sim_air_insert_socket_peg...")
    task_state_insert = sample_air_insert_socket_peg_state()
    policy_insert = TestAirInsertPolicy(inject_noise=False)
    first_i, last_i = capture_task('sim_air_insert_socket_peg', policy_insert, task_state_insert, 600)
    cv2.imwrite(os.path.join(OUTPUT_DIR, "air_insert_front_first.png"), first_i)
    cv2.imwrite(os.path.join(OUTPUT_DIR, "air_insert_front_last.png"), last_i)

    print(f"\nDone! Images saved to {OUTPUT_DIR}/")


if __name__ == '__main__':
    main()