# Copyright 2026 Toyota Research Institute.  All rights reserved.

import os
import json
import numpy as np

from glob import glob

from anydata.utils.read import read_numpy, read_yaml, 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.utils.geometry import invert_extrinsics
from anydata.utils.colmap import qvec2rotmat

import re

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

def get_rgb(filename):
    return np.array(read_image(filename))


def get_depth(filename):
    with open(filename, "rb") as file:

        header = file.readline().rstrip()
        if header.decode("ascii") == "PF":
            color = True
        elif header.decode("ascii") == "Pf":
            color = False
        else:
            raise Exception("Not a PFM file: " + filename)

        dim_match = re.match(r"^(\d+)\s(\d+)\s$", file.readline().decode("ascii"))
        if dim_match:
            width, height = list(map(int, dim_match.groups()))
        else:
            raise Exception("Malformed PFM header.")

        scale = float(file.readline().decode("ascii").rstrip())
        if scale < 0:
            # little-endian
            endian = "<"
            scale = -scale
        else:
            # big-endian
            endian = ">"

        data = np.fromfile(file, endian + "f")
        shape = (height, width, 3) if color else (height, width)

        data = np.reshape(data, shape)
        data = np.flipud(data)

        return data * scale


def get_intrinsics(filename):

    with open(filename, 'r') as f:
        data = f.readlines()

    line1 = data[7].replace('\n', '').split(' ')[:3]
    line2 = data[8].replace('\n', '').split(' ')[:3]
    line3 = data[9].replace('\n', '').split(' ')[:3]

    line1 = [float(l) for l in line1]
    line2 = [float(l) for l in line2]
    line3 = [float(l) for l in line3]

    return np.array([line1, line2, line3])

def get_extrinsics(filename):

    with open(filename, 'r') as f:
        data = f.readlines()

    line1 = data[1].replace('\n', '').split(' ')[:4]
    line2 = data[2].replace('\n', '').split(' ')[:4]
    line3 = data[3].replace('\n', '').split(' ')[:4]
    line4 = data[4].replace('\n', '').split(' ')[:4]

    line1 = [float(l) for l in line1]
    line2 = [float(l) for l in line2]
    line3 = [float(l) for l in line3]
    line4 = [float(l) for l in line4]

    extrinsics = np.array([line1, line2, line3, line4])
    extrinsics = invert_extrinsics(extrinsics)

    return extrinsics

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

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


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

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

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

    ### 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']

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

        ### Get filenames
        filename_rgbs = sorted(glob(f'{seq}/blended_images/*.jpg'))
        filename_masks = [f for f in filename_rgbs if 'masked' in f]
        filename_rgbs = [f for f in filename_rgbs if 'masked' not in f]

        filename_depths = sorted(glob(f'{seq}/rendered_depth_maps/*.pfm'))
        filename_cameras = sorted(glob(f'{seq}/cams/*_cam.txt'))

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

            ######## RGB
            rgb = get_rgb(filename_rgb)
            dense['rgb'][frame] = rgb

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

        ######## CAMERA FILENAMES
        for i, filename_camera in enumerate(filename_cameras):
            frame = frame_name(i)

            filename_lowdim = add_key_to_dict(lowdim, f'{dst}/lowdim/{cam}/{frame}.npz')
            lowdim[filename_lowdim]['intrinsics'] = get_intrinsics(filename_camera)
            lowdim[filename_lowdim]['extrinsics'] = get_extrinsics(filename_camera)

        ######## DEPTH FILENAMES
        for i, filename_depth in enumerate(filename_depths):
            frame = frame_name(i)

            ### DEPTH
            depth = get_depth(filename_depth)
            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 
    filename = f'{dst}/metadata.json'
    seq_metadata = fill_metadata(
        args=args,
        info=dict(
            name='BlendedMVG',
            tags=['real','static'],
            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 dst

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

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

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