import numpy as np
import json
import cv2
from sim.main import InteractiveDigitalWorld
from sim.simulator import GenieSimulator, ReplaySimulator
from sim.policy import ReplayPolicy


# NOTE: ad hoc
def normalize_frames(frames):
    new_frames = []
    for frame in frames:
        H, W = frame.shape[:2]
        if H < W:
            Hnew, Wnew = 256, int(W * 256 / H)
        else:
            Hnew, Wnew = int(H * 256 / W), 256
        frame = cv2.resize(frame, (Wnew, Hnew))
        H, W = frame.shape[:2]
        Hstart = (H - 256) // 2
        Wstart = (W - 256) // 2
        frame = frame[Hstart:Hstart+256, Wstart:Wstart+256]
        new_frames.append(frame)
    return np.stack(new_frames, axis=0)


if __name__ == '__main__':

    prompt_horizon = 11
    action_stride = 1

    dataset_dir = "data/langtable_raw_train"
    metadata = json.load(open(f"{dataset_dir}/metadata.json"))
    h, w = metadata['h'], metadata['w']
    action_dim = metadata['action_dim']
    num_images = metadata['num_images']
    actions = np.fromfile(f"{dataset_dir}/actions/actions.bin", dtype=np.float32).reshape(num_images, action_dim)
    frames = np.fromfile(f"{dataset_dir}/video.bin", dtype=np.uint8).reshape(num_images, h, w, 3)
    # frames = normalize_frames(frames)
    segment_ids = np.fromfile(f"{dataset_dir}/segment_ids.bin", dtype=np.int32)
    print(f"{actions.shape=}, {frames.shape=}, {segment_ids.shape=}")

    # get chunks' start and end
    chunks = []     # [start_index, end_index)
    start_index = 0
    end_index = 0
    while end_index < len(segment_ids):
        while end_index < len(segment_ids) and segment_ids[end_index] == segment_ids[start_index]:
            end_index += 1
        if end_index - start_index > prompt_horizon * 2:
            chunks.append((start_index, end_index))
        start_index = end_index
    print(f"there're {len(chunks)} chunks")

    for eps_idx, chunk in enumerate(chunks):
        start_idx, end_idx = chunk
        this_frames = frames[start_idx:end_idx]
        this_actions = actions[start_idx:end_idx]
        print(f"processing chunk {eps_idx} with {len(this_frames)} frames")

        replay_simulator = ReplaySimulator(frames=this_frames, prompt_horizon=prompt_horizon)
        replay_policy = ReplayPolicy(actions=this_actions, prompt_horizon=prompt_horizon, action_stride=action_stride)
        assert len(replay_policy) == len(replay_simulator)

        genie_simulator = GenieSimulator(
            image_encoder_type='temporalvae',
            image_encoder_ckpt='stabilityai/stable-video-diffusion-img2vid',
            quantize=False,
            backbone_type="stmar",
            backbone_ckpt="data/mar_ckpt/langtable",
            prompt_horizon=prompt_horizon,
            action_stride=action_stride,
            domain='language_table',
            physics_simulator=replay_simulator,
            compute_psnr=False,
            compute_delta_psnr=False,
            allow_external_prompt=True
        )

        # use whatever current state is as the initial state
        image_prompt = replay_simulator.prompt()
        action_prompt = replay_policy.prompt()
        genie_simulator.set_initial_state((image_prompt, action_prompt))
        playground = InteractiveDigitalWorld(
            simulator=genie_simulator,
            policy=replay_policy,
            offscreen=True,
            window_size=(512 * 2, 512)  # [genie image | GT image] side-by-side
        )

        for _ in range(len(replay_policy)):
            playground.step()

        save_video_path = f'data/langtable_train_videos/{eps_idx}.mp4'
        print(f"Saving video to {save_video_path}")
        playground.save_video(save_path=save_video_path, as_gif=False)
        playground.close()