# Copyright 2026 Toyota Research Institute.  All rights reserved.

import os
import cv2
import numpy as np

import pickle
from glob import glob
from datetime import datetime

from anydata.utils.read import read_numpy, read_yaml, read_image, read_depth
from anydata.utils.write import write_image, write_npz, write_json
from anydata.converters.utils import run, add_key_to_dict, fill_metadata, parse_dst_seq, frame_name, crawl

#######################################################
    
def process_state_and_time(path):
    fp = os.path.join(path, "obs_dict.pkl")
    with open(fp, "rb") as f:
        x = pickle.load(f)
    return x["full_state"], x["time_stamp"]

def process_actions(path):  # gets actions
    fp = os.path.join(path, "policy_out.pkl")
    with open(fp, "rb") as f:
        act_list = pickle.load(f)
    if isinstance(act_list[0], dict):
        act_list = [x["actions"] for x in act_list]
    return act_list

#######################################################

def get_sequences(args):
    seqs = crawl(args.src, 'obs_dict.pkl')
    seqs = [os.path.dirname(seq) for seq in seqs]
    return seqs


def parse_sequence(seq, args):
    return parse_dst_seq(seq, args, remove=[6])

#######################################################

def process_sequence(i, seq, dst, args):

    ### Initialize lists and dicts
    cameras = sorted(glob(f'{seq}/images*'))
    cameras = [os.path.basename(c) for c in cameras]
    num_frames = {cam: dict() for cam in cameras}
    resolution = {cam: dict() for cam in cameras}
    labels, lowdim = [], {}

    # Data collected prior to 7-23 has a delay of 1, otherwise a delay of 0
    date_time = datetime.strptime(seq.split("/")[-4], "%Y-%m-%d_%H-%M-%S")
    latency_shift = date_time < datetime(2021, 7, 23)

    action = process_actions(seq)
    state, timestamp = process_state_and_time(seq)

    # # shift the actions according to camera latency (NOT USED CURRENTLY)
    # if latency_shift:
    #     out["observations"] = out["observations"][1:]
    #     out["next_observations"] = out["next_observations"][1:]
    #     out["actions"] = out["actions"][:-1]
    #     out["terminals"] = term[:-1]

    language = f'{seq}/lang.txt'
    if os.path.exists(language):
        with open(language) as file:
            prompt = [l.strip() for l in list(file) if "confidence" not in l]
            language = dict(prompt=prompt)
            if 'language' not in labels: labels.append('language')
    else:
        language = None

    ############ LOOP OVER CAMERAS
    for cam in cameras:

        ### Get filenames
        filename_rgbs = sorted(glob(f'{seq}/{cam}/*.jpg'))

        # Check if everything has the correct length        
        assert len(filename_rgbs) == len(state) and \
               len(filename_rgbs) == len(timestamp), 'Wrong length'

        ######## RGB FILENAMES
        for i, filename_rgb in enumerate(filename_rgbs):
            frame = frame_name(i)

            ######## RGB
            if 'rgb' not in labels: labels.append('rgb')
            rgb = np.array(read_image(filename_rgb).resize((640, 480)))
            filename_rgb_out = f'{dst}/rgb/{cam}/{frame}.jpg'
            write_image(filename_rgb_out, rgb)

            ######## LOWDIM RGB             
            filename_lowdim = add_key_to_dict(lowdim, f'{dst}/lowdim/{cam}/{frame}.npz')
            lowdim[filename_lowdim]['camera'] = cam
            lowdim[filename_lowdim]['timestep'] = timestamp[i]

        if len(filename_rgbs) > 0:
            resolution[cam]['rgb'] = rgb.shape[:2]
            num_frames[cam]['rgb'] = len(filename_rgbs)

        ######## ACTION
        for i in range(len(state)):             
            frame = frame_name(i)

            if 'action' not in labels: labels.append('action')
            filename_lowdim = add_key_to_dict(lowdim, f'{dst}/lowdim/{cam}/{frame}.npz')
            lowdim[filename_lowdim]['action'] = {'state': state[i]}
            if i < len(action): 
                lowdim[filename_lowdim]['action']['action'] = action[i]
        num_frames[cam]['action'] = len(state)


    ######## WRITE LOWDIM
    for key, val in lowdim.items():
        write_npz(key, val)

############ METADATA 
    filename = f'{dst}/metadata.json'
    seq_metadata = fill_metadata(
        args=args,
        info=dict(
            name='BridgeV2',
            tags=['real','robotics','tabletop'],
            raw_id=seq.replace(f'{args.src}/', ''),
        ),
        labels=labels,
        cameras=cameras,
        resolution=resolution,
        num_frames=num_frames,
        framerate=10,
        rgb=dict(extension='jpg'),
        intrinsics=None,
        extrinsics=None,
        depth=None,
        semantic=None,
        action=None,
        language=language,
        specific=None,
    )
    write_json(filename, seq_metadata)

    return dst

#######################################################

if __name__ == '__main__':
    converter = os.path.basename(__file__)
    run(converter, get_sequences, parse_sequence, process_sequence)

#######################################################