# Copyright 2026 Toyota Research Institute.  All rights reserved.

import os
import numpy as np
import pandas as pd

from tqdm import tqdm
from glob import glob

from anydata.utils.read import read_numpy, read_yaml, read_image, read_depth
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 geometry_from_colmap, read_dense_colmap_binary
from anydata.utils.write import write_jpg_from_mp4

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

def get_sequences(args):
    # seqs = crawl(args.src, 'sparse')
    seqs = crawl(args.src, 'cameras.bin', avoid='Sparse')
    seqs = [os.path.dirname(seq) for seq in seqs]
    return seqs


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

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

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

    ### Get COLMAP data
    cameras_binary, points_binary, images_binary = read_dense_colmap_binary(seq)

    # print('starting')
    # folder_rgb_out = f'{dst_all}/rgb/video2'
    # video = '/data/cv_downloaded/EpicKitchens100/P01/videos/P01_01.MP4'
    # resolution_rgb, num_rgb_frames = write_jpg_from_mp4(folder_rgb_out, video)
    # print(resolution_rgb, num_rgb_frames)
    # import sys
    # sys.exit()

    # Get train/val annotations
    ann_train = '/data/cv_downloaded/EpicKitchens100/annotations/EPIC_100_train.csv'
    ann_train = pd.read_csv(ann_train, index_col="participant_id")
    ann_val = '/data/cv_downloaded/EpicKitchens100/annotations/EPIC_100_validation.csv'
    ann_val = pd.read_csv(ann_val, index_col="participant_id")

    # create list of start/end frames for each action
    actions = []
    seq_id = os.path.basename(seq)
    for id, st, fn, narration in zip(
            ann_train['video_id'], 
            ann_train['start_frame'], 
            ann_train['stop_frame'], 
            ann_train['narration']):
        if id == seq_id: 
            actions.append([st, fn, narration])
    for id, st, fn, narration in zip(
            ann_val['video_id'], 
            ann_val['start_frame'], 
            ann_val['stop_frame'], 
            ann_val['narration']
        ):
        if id == seq_id: 
            actions.append([st, fn, narration])
    actions = sorted(actions, key=lambda x: x[0])

    ############ LOOP OVER CAMERAS
    values = sorted(images_binary.values(), key=lambda x: x.name)

    extrinsics_breaks = []
    for i in range(len(values)-1):
        v0 = int(values[i].name.split('_')[-1].split('.')[0])
        v1 = int(values[i+1].name.split('_')[-1].split('.')[0])
        if v1 - v0 > 2:
            extrinsics_breaks.append([v0+1,v1-1])

    filtered_actions = []
    for act in actions:
        invalid = False
        st, fn = act[:2]
        for missing in extrinsics_breaks:
            invalid = st <= missing[0] <= fn or \
                      st <= missing[1] <= fn or \
                      (missing[0] <= st and missing[1] >= fn)
            break
        if not invalid:
            filtered_actions.append(act)
    actions = filtered_actions
    
    dsts, cnt = [], 0
    for action in tqdm(actions, ncols=96, leave=False):
    # for j in tqdm(range(len(extrinsics_breaks) - 1), ncols=96, leave=False):\

        print(action)

        ### Initialize lists and dicts
        cameras = ['0']
        num_frames = {cam: dict() for cam in cameras}
        resolution = {cam: dict() for cam in cameras}
        labels, lowdim = [], {}
        dense_labels = ['rgb','depth']

        cam = cameras[0]

        st, fn = action[:2]
        dst = f'{dst_all}/%03d_%03d' % (st, fn) # Sequence name
        dsts.append(dst)

        cnt = 0
        for i, v in tqdm(enumerate(range(st,fn, 2)), ncols=96, leave=False):
            dense = {key: dict() for key in dense_labels}

            frame = frame_name(i // 2)
            val = values[v]
            cnt += 1

            ### Extract COLMAP information
            intrinsics, extrinsics, depth, hw = geometry_from_colmap(
                val, cameras_binary, points_binary)

            ######## RGB
            base = os.path.basename(seq)
            first = base.split('_')[0]
            filename_rgb = seq.replace('EpicFields', 'EpicKitchens100')
            filename_rgb = '/'.join(filename_rgb.split('/')[:4])        
            filename_rgb = f'{filename_rgb}/{first}/rgb_frames/{base}/{val.name}'
            rgb = np.array(read_image(filename_rgb))
            dense['rgb'][frame] = rgb

            ######## LOWDIM RGB
            prepare_lowdim(lowdim, dst, cam, frame)

            ######## DEPTH
            dense['depth'][frame] = depth

            ######## INTRINSICS + EXTRINSICS
            filename_lowdim = add_key_to_dict(lowdim, f'{dst}/lowdim/{cam}/{frame}.npz')
            lowdim[filename_lowdim]['extrinsics'] = extrinsics
            lowdim[filename_lowdim]['intrinsics'] = intrinsics

            prompt = []
            timestamp = int(val.name.split('_')[-1].split('.')[0])
            for act in actions:
                if act[0] <= timestamp <= act[1]:
                    prompt.append(act[2])
            if len(prompt) > 0:
                filename_lowdim = add_key_to_dict(lowdim, f'{dst}/lowdim/{cam}/{frame}.npz')
                lowdim[filename_lowdim]['language'] = dict(prompt=prompt)

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

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

    ############ METADATA 
        filename = f'{dst}/metadata.json'
        seq_metadata = fill_metadata(
            args=args,
            info=dict(
                name='EpicFields',
                tags=['real','dynamic','egocentric','human'],
                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='pinhole'),
            extrinsics=dict(transform='cam2world',metric=False),
            depth=dict(extension='npz',metric=False,sparse=True),
            semantic=None,
            action=None,
            language=None,
            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)

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