# Copyright 2026 Toyota Research Institute.  All rights reserved.

import os
import cv2
import h5py
import torch
import numpy as np
import imageio.v3 as iio

from glob import glob

from anydata.utils.geometry import pose_to_matrix
from anydata.utils.read import read_image, read_depth, read_json
from anydata.utils.write import write_json, write_lowdim, write_labels
from anydata.converters.utils import run, add_key_to_dict, fill_metadata, parse_dst_seq, frame_name, crawl, prepare_lowdim
from anydata.converters.utils import extract_mp4

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

def get_depth(filename):
    depth = read_depth(filename, div=1000)
    return depth
    
def h5todict(data):
    if isinstance(data, h5py._hl.dataset.Dataset):
        return data[:]
    else:
        return {key: h5todict(val) for key, val in data.items()}
    
#######################################################

# Intrinsics rescaling detection thresholds
# Based on analysis: episodes 648642-651370 have unscaled intrinsics,
# episodes 682363-685393 have properly rescaled intrinsics.
# Head camera: ppx ~648 (unscaled) vs ~326 (rescaled) for 640x480 video
# Fisheye: pu ~964 (unscaled) vs ~481 (rescaled) for 960x768 video
INTRINSICS_RESCALED_THRESHOLDS = {
    'head': {'ppx': 400},           # < 400 means rescaled
    'hand_left': {'ppx': 380},      # < 380 means rescaled
    'hand_right': {'ppx': 380},     # < 380 means rescaled
    'head_center_fisheye': {'pu': 600},  # < 600 means rescaled
    'head_left_fisheye': {'pu': 600},
    'head_right_fisheye': {'pu': 600},
    'back_left_fisheye': {'pu': 600},
    'back_right_fisheye': {'pu': 600},
}

def check_intrinsics_rescaled(intrinsics_raw, cam):
    """
    Check if intrinsics have been properly rescaled to match video resolution.

    Returns True if intrinsics are rescaled (correct), False if unscaled (need adjustment).
    """
    if cam not in INTRINSICS_RESCALED_THRESHOLDS:
        return None  # Unknown camera type

    threshold_info = INTRINSICS_RESCALED_THRESHOLDS[cam]

    # Get the principal point key based on camera type
    if 'ppx' in threshold_info:
        pp_key = 'ppx'
    elif 'pu' in threshold_info:
        pp_key = 'pu'
    else:
        return None

    if pp_key not in intrinsics_raw:
        return None

    pp_value = intrinsics_raw[pp_key]
    threshold = threshold_info[pp_key]

    return pp_value < threshold

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

def get_sequences(args):
    seqs = crawl(f'{args.src}/observations', 'videos')
    seqs = [os.path.dirname(seq) for seq in seqs]
    return seqs


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

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

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

    ### Initialize lists and dicts
    videos = glob(f'{seq}/videos/*.mp4')
    cameras = [os.path.splitext(vid.split('/')[-1])[0].replace('_color', '') for vid in videos]
    cameras = [c for c in cameras if 'back' not in c]
    # cameras = ['head']

    ### Get task information
    task, episode = videos[0].split('/')[-4:-2]
    all_task_params = '/'.join(videos[0].split('/')[:-4]).replace('observations','task_info') + f'/task_{task}.json'
    all_task_params = read_json(all_task_params) # [0]
    task_params = [val for val in all_task_params if val['episode_id'] == int(episode)]
    assert len(task_params) == 1, f'Error processing {seq}'
    task_params = task_params[0]

    task = task_params['task_name']
    scene = task_params['init_scene_text']

    videos = {cam: extract_mp4(f'{seq}/videos/{cam}_color.mp4') for cam in cameras}

    dsts = []
    for action_config in task_params['label_info']['action_config']:

        start_frame = action_config['start_frame']
        end_frame = action_config['end_frame']

        dst = f'{dst_all}/%d_%d' % (start_frame, end_frame)
        dsts.append(dst)

        prompt = action_config.get('action_text', None)
        skill = action_config.get('skill', None)

        num_frames = {cam: dict() for cam in cameras}
        resolution = {cam: dict() for cam in cameras}
        labels, lowdim = [], {}
        dense_labels = ['rgb','depth']

        intrinsics_model = {}     # Track camera model
        intrinsics_rescaled = {}  # Track rescaling status per camera

        ############ LOOP OVER CAMERAS
        for cam in cameras:
            dense = {label: dict() for label in dense_labels}

            ### Get filenames
            vid = f'{seq}/videos/{cam}_color.mp4'

            filename_params = '/'.join(vid.split('/')[:-2]).replace('observations','parameters')
            filename_extrinsics = f'{filename_params}/parameters/camera/{cam}_extrinsic_params_aligned.json'
            filename_intrinsics = f'{filename_params}/parameters/camera/{cam}_intrinsic_params.json'

            # ### Get intrinsics
            intrinsics = read_json(filename_intrinsics)
            if 'intrinsic' in intrinsics:
                intrinsics = intrinsics['intrinsic']
            intrinsics_model[cam] = intrinsics.pop('distortion_model')

            ### Check if intrinsics are properly rescaled to match video resolution
            intrinsics_rescaled[cam] = check_intrinsics_rescaled(intrinsics, cam)

            x_intr_ratio, y_intr_ratio = 1.0, 1.0
            if intrinsics_model[cam] == 'plumb bob':
                if not intrinsics_rescaled[cam]:
                    if cam.startswith('head'):
                        x_intr_ratio, y_intr_ratio = 2.0, 1.5
                    elif cam.startswith('hand'):
                        x_intr_ratio, y_intr_ratio = 1.325, 1.0
                intrinsics_model[cam] = 'plumb_bob'
                fx, fy, cx, cy = intrinsics['fx'], intrinsics['fy'], intrinsics['ppx'], intrinsics['ppy']       
                k1, k2, p1, p2, k3 = intrinsics['k1'], intrinsics['k2'], intrinsics['p1'], intrinsics['p2'], intrinsics['k3']     
                intrinsics = np.array([fx / x_intr_ratio, fy / y_intr_ratio, 
                                    cx / x_intr_ratio, cy / y_intr_ratio, 
                                    k1, k2, p1, p2, k3], dtype=np.float32)  
            elif intrinsics_model[cam] == 'equidistant':
                if not intrinsics_rescaled[cam]:
                    x_intr_ratio, y_intr_ratio = 2.0, 2.0
                intrinsics_model[cam] = 'fisheye'
                fx, fy, cx, cy = intrinsics['fu'], intrinsics['fv'], intrinsics['pu'], intrinsics['pv']
                k1, k2, k3, k4 = intrinsics['k1'], intrinsics['k2'], intrinsics['k3'], intrinsics['k4']     
                intrinsics = np.array([fx / x_intr_ratio, fy / y_intr_ratio, 
                                    cx / x_intr_ratio, cy / y_intr_ratio, 
                                    k1, k2, k3, k4], dtype=np.float32)  
            else:
                raise ValueError(f'INVALID INTRINSICS MODEL {intrinsics_model[cam]}')

            extrinsics_all = read_json(filename_extrinsics)

            ### Get action information
            filename_proprio = '/'.join(vid.split('/')[:-2]).replace('observations','proprio_stats')
            with h5py.File(f'{filename_proprio}/proprio_stats.h5', "r") as f:
                action, state = f['action'], f['state']
                action, state = h5todict(action), h5todict(action)

            task = task_params['task_name']
            scene = task_params['init_scene_text']
            skill = action_config.get('skill', None)
            prompt = action_config.get('action_text', None)

            start_frame = action_config['start_frame']
            end_frame = action_config['end_frame']
            video = videos[cam] # extract_mp4(vid, start_frame, end_frame)

            for i, idx in enumerate(range(start_frame, end_frame)):
                frame = frame_name(i)

                ######## RGB FILENAMES
                dense['rgb'][frame] = video[idx]

                ######## LOWDIM RGB    
                filename_lowdim = add_key_to_dict(lowdim, f'{dst}/lowdim/{cam}/{frame}.npz')
                prepare_lowdim(lowdim, dst, cam, frame)

                ######## EXTRINSICS
                if extrinsics_all is not None:
                    extrinsics_i = extrinsics_all[idx]['extrinsic']
                    rot = extrinsics_i['rotation_matrix']
                    tvec = extrinsics_i['translation_vector']
                    extrinsics = np.vstack([
                        np.hstack((rot, np.expand_dims(tvec, axis=1))),
                        np.array([0.0, 0.0, 0.0, 1.0])
                    ]).astype(np.float32)
                    lowdim[filename_lowdim]['extrinsics'] = extrinsics

                ######## INTRINSICS
                lowdim[filename_lowdim]['intrinsics'] = intrinsics

                ######## ACTION
                lowdim[filename_lowdim]['action'] = dict(
                    action={key1: {key2: val2[idx] for key2, val2 in val1.items() if len(val2.shape) > 1} 
                            for key1, val1 in action.items()},
                    state={key1: {key2: val2[idx] for key2, val2 in val1.items() if len(val2.shape) > 1} 
                            for key1, val1 in action.items()},
                )

                ######## DEPTH
                if cam == 'head':
                    filename_depth = os.path.dirname(vid).replace('/videos','/depth') + '/head_depth_%06d.png' % idx 
                    depth = read_depth(filename_depth, div=1000)
                    dense['depth'][frame] = depth

            ######## WRITE LABELS
            write_labels(dst, cam, args.storage, dense, labels, resolution, num_frames)

        ######## WRITE LOWDIM
        write_lowdim(args, dst, labels, num_frames, lowdim)

        ############ METADATA
        # Determine overall intrinsics rescaling status for tags
        # True if ALL cameras have rescaled intrinsics, False if ANY camera has unscaled
        all_rescaled = all(v is True for v in intrinsics_rescaled.values() if v is not None)
        any_unscaled = any(v is False for v in intrinsics_rescaled.values())
        intrinsics_status = 'rescaled' if all_rescaled and not any_unscaled else 'unscaled'

        filename = f'{dst}/metadata.json'
        seq_metadata = fill_metadata(
            args=args,
            info=dict(
                name='AgiBotWorld',
                tags=['real','dynamic','robotics','egocentric','wrist'],
                raw_id=seq.replace(f'{args.src}/', ''),
            ),
            labels=labels,
            cameras=cameras,
            resolution=resolution,
            num_frames=num_frames,
            framerate=10,
            rgb=dict(extension='jpg'),
            intrinsics=dict(model=intrinsics_model, rescaled=intrinsics_rescaled, status=intrinsics_status),
            extrinsics=dict(transform='cam2world',metric=True),
            depth=dict(extension='npz',metric=True,sparse=True),
            semantic=None,
            action=dict(format='AgiBotWorld'),
            language=dict(task=task,scene=scene,skill=skill,prompt=prompt),
            specific=None,
        )
        write_json(filename, seq_metadata)

    return dsts

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

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

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