from shutil import copy
import sys
import os
import json
import argparse
from typing import Mapping
from data import get_dataset_desc_path, get_data_path

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

parser = argparse.ArgumentParser()
parser.add_argument('-i', '--input', type=str, nargs='+')
parser.add_argument('output', type=str)
args = parser.parse_args()


input = [get_dataset_desc_path(path) for path in args.input]
output = get_dataset_desc_path(args.output)


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


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

n_views = 0
for i in range(len(input)):
    if i == 0:
        input_desc = dataset_desc
    else:
        with open(input[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(get_data_path(input[i], input_desc['view_file_pattern']),
                get_data_path(output, dataset_desc['view_file_pattern']),
                len(input_desc['view_centers']), n_views)
    n_views += len(input_desc['view_centers'])

dataset_desc['samples'] = [n_views]

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