import sys
import os
import argparse
import json
import numpy as np

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

from utils import img

parser = argparse.ArgumentParser()
parser.add_argument('dataset', type=str)
args = parser.parse_args()

data_desc_path = args.dataset
data_desc_name = os.path.splitext(os.path.basename(data_desc_path))[0]
data_dir = os.path.dirname(data_desc_path) + '/'

with open(data_desc_path, 'r') as fp:
    dataset_desc = json.load(fp)

centers = np.array(dataset_desc['view_centers'])
t_max = np.max(centers, axis=0)
t_min = np.min(centers, axis=0)
dataset_desc['range'] = {
    'min': [t_min[0], t_min[1], t_min[2], 0, 0],
    'max': [t_max[0], t_max[1], t_max[2], 0, 0]
}
dataset_desc['gl_coord'] = True

for subdir in ['images', 'images_4', 'images_8']:
    dir = os.path.join(data_dir, subdir)
    files = os.listdir(dir)
    files.sort()
    image = img.load(os.path.join(dir, files[0]))
    res = image.shape[-2:]
    for i, file in enumerate(files):
        src = os.path.join(dir, file)
        tgt = os.path.join(dir, "view_%04d.jpg" % i)
        print(f"Rename {src} to {tgt}")
        os.rename(src, tgt)
    out_desc = dataset_desc.copy()
    out_desc['color_file'] = f"{subdir}/view_%04d.jpg"
    k = res[0] / dataset_desc['view_res']['y']
    out_desc['view_res'] = {
        'x': res[1],
        'y': res[0]
    }
    out_desc['cam_params'] = {
        'fx': dataset_desc['cam_params']['fx'] * k,
        'fy': dataset_desc['cam_params']['fy'] * k,
        'cx': res[1] / 2,
        'cy': res[0] / 2,
    }
    
    with open(os.path.join(data_dir, f"{subdir}.json"), 'w') as fp:
        json.dump(out_desc, fp, indent=4)