import numpy as np


ROT_LIGHT_NUM_LIGHT = 240
ROT_LIGHT_NUM_LOOP = 8


def gen_rot_light__light_pos(num_light, num_loop):
    theta = np.linspace(0, np.pi, num_light)
    phi = np.linspace(0, 2 * np.pi * num_loop, num_light)
    x = np.sin(theta) * np.cos(phi)
    z = np.sin(theta) * np.sin(phi)
    y = np.cos(theta)
    return np.stack([x, y, z], axis=-1)


def load_cam_traj(cam_traj_f, gt_intr):
    if cam_traj_f.suffix == '.npz':
        info = np.load(cam_traj_f)
        return info['traj'], info['intr']
    else:
        assert cam_traj_f.suffix == '.npy', f'Unsupported {cam_traj_f=}'
        return np.load(cam_traj_f), gt_intr


def get_contour_line_files(summaries, key, num_high_res_to_exc=0):
    def iterate_summary_json(summary, ret, keys=[]):
        if isinstance(summary, dict):
            for k, v in summary.items():
                iterate_summary_json(v, ret, keys + [k])
        else:
            ret.append((keys, '/'.join(summary.split('/')[-2:])))

    files1, files2 = [], []
    iterate_summary_json(summaries[0][key], files1)
    iterate_summary_json(summaries[1][key], files2)
    files1 = sorted(files1, key=lambda x: [float(y) for y in x[0]])
    files2 = sorted(files2, key=lambda x: [float(y) for y in x[0]])
    files1 = files1[:len(files1) - num_high_res_to_exc]
    files2 = files2[:len(files2) - num_high_res_to_exc]
    assert len(files1) == len(files2)
    for i in range(len(files1)):
        assert files1[i][1] == files2[i][1]
    files1 = [x[1] for x in files1]
    files2 = [x[1] for x in files2]
    return files1, files2
