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

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

from utils import seqs, math
from utils.types import Resolution


parser = argparse.ArgumentParser()
parser.add_argument('-r', '--rot-range', nargs='+', type=int)
parser.add_argument('-t', '--trans-range', nargs='+', type=float)
parser.add_argument('--fov', type=float)
parser.add_argument('--res', type=str)
parser.add_argument('--gl', action='store_true')
parser.add_argument('-s', '--seq', type=str, required=True)
parser.add_argument('-n', '--views', type=int, required=True)
parser.add_argument('-o', '--out-desc', type=str)
parser.add_argument('--ref', type=str)
parser.add_argument('dataset', type=str)
args = parser.parse_args()


data_dir = args.dataset
os.makedirs(data_dir, exist_ok=True)
out_desc_path = os.path.join(data_dir, args.out_desc or f"{args.seq}.json")

if args.ref:
    with open(os.path.join(data_dir, args.ref), 'r') as fp:
        ref_desc = json.load(fp)
else:
    if not args.trans_range or not args.rot_range or not args.fov or not args.res:
        print('-r, -t, --fov, --res options are required if --ref is not specified')
        exit(-1)
    ref_desc = None

if args.trans_range:
    trans_range = np.array(args.trans_range * 3 if len(args.trans_range) == 1 else args.trans_range)
else:
    trans_range = np.array(ref_desc["trs_range"])
if args.rot_range:
    rot_range = np.array(args.rot_range * 2 if len(args.rot_range) == 1 else args.rot_range)
else:
    rot_range = np.array(ref_desc["rot_range"])
filter_range = np.concatenate([trans_range, rot_range])

cam_params = { "fov": args.fov } if args.fov else ref_desc["cam"]
res = Resolution.from_str(args.res or ref_desc["res"])

if args.seq == 'helix':
    centers, rots = seqs.helix(trans_range, 4, args.views)
elif args.seq == 'scan_around':
    centers, rots = seqs.scan_around(trans_range, 1, args.views)
elif args.seq == 'look_around':
    centers, rots = seqs.look_around(trans_range, args.views)

rots = np.degrees(rots)
gl = args.gl or ref_desc and ref_desc.get('gl_coord')
if gl:
    centers[:, 2] *= -1
    rots[:, 0] *= -1

dataset_desc = {
    'gl_coord': gl,
    'res': f"{res.w}x{res.h}",
    'cam': cam_params,
    "trs_range": trans_range.tolist(),
    "rot_range": rot_range.tolist(),
    'samples': [args.views],
    'centers': centers.tolist(),
    'rots': rots.tolist()
}

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