from shutil import copy
import sys
import os
import json
import argparse
from typing import Mapping

sys.path.append(os.path.abspath(sys.path[0] + '/../'))

from utils import misc


def copy_images(src_path, dst_path, n, offset=0):
    misc.create_dir(os.path.dirname(dst_path))
    for i in range(n):
        copy(src_path % i, dst_path % (i + offset))


input_data_desc_paths = [
    '/home/dengnc/dvs/data/__new/barbershop_all/nerf_cvt.json',
    '/home/dengnc/dvs/data/__new/__demo/fvvdp/0816_3_barbershop_cvt.json',
    '/home/dengnc/dvs/data/__new/__demo/fvvdp/0816_1_barbershop_cvt.json',
]
output_data_desc_path = '/home/dengnc/dvs/data/__new/__demo/fvvdp/barbershop_3&1_nerf.json'
output_data_name = os.path.splitext(os.path.basename(output_data_desc_path))[0]
output_dir = os.path.dirname(output_data_desc_path)

with open(input_data_desc_paths[0], 'r') as fp:
    dataset_desc: Mapping = json.load(fp)

n_views = 0
# Copy images of the first dataset

for i in range(len(input_data_desc_paths)):
    if i == 0:
        input_desc = dataset_desc
    else:
        with open(input_data_desc_paths[i], 'r') as fp:
            input_desc: Mapping = json.load(fp)
        dataset_desc['view_centers'] += input_desc['view_centers']
        dataset_desc['view_rots'] += input_desc['view_rots']
    copy_images(
        os.path.join(os.path.dirname(input_data_desc_paths[i]), input_desc['view_file_pattern']),
        os.path.join(output_dir, output_data_name, 'view_%04d.png'),
        len(input_desc['view_centers']), n_views
    )
    n_views += len(input_desc['view_centers'])

dataset_desc['samples'] = [n_views]
dataset_desc['view_file_pattern'] = os.path.join(output_data_name, 'view_%04d.png')

with open(output_data_desc_path, 'w') as fp:
    json.dump(dataset_desc, fp, indent=4)