import numpy as np
from pathlib import Path

data_directory = '../schwarz_data/{}'

snr_grid = np.load(data_directory.format('grid_snrs.npy'))
intrinsics = np.load(data_directory.format('grid_intrinsics.npy'))
angulars = np.load(data_directory.format('grid_extrinsics.npy'))
plunges = np.load(data_directory.format('grid_plunges.npy'))

intrinsics[:,:2] = np.log10(intrinsics[:,:2])

grid_x_out = np.zeros(shape=(snr_grid.size,9))

num = 0
for intrinsic_set in intrinsics:
    for angular_set in angulars:
        for time in plunges:
            out = np.zeros(9)
            out[:5] = intrinsic_set
            out[5:8] = angular_set
            out[8] = time

            grid_x_out[num,:] = out

            num += 1

save_dir = '../_data/{}'
Path('../_data/').mkdir(parents=True, exist_ok=True)

np.save(save_dir.format('xgrid.npy'),arr=grid_x_out)
np.save(save_dir.format('ygrid.npy'),arr=snr_grid.flatten() + 1e-5)  # add small offset for log of zeros

snr_list = np.load(data_directory.format('samp_snrs.npy'))
inds_to_keep = ~np.isnan(snr_list)
np.save(save_dir.format('ysamp.npy'),arr=snr_list[inds_to_keep] + 1e-5)

intrinsics = np.load(data_directory.format('samp_intrinsics.npy'))
angulars = np.load(data_directory.format('samp_extrinsics.npy'))
intrinsics[:,:2] = np.log10(intrinsics[:,:2])

samp_x_out = np.zeros(shape=(snr_list.size,9))

for i in range(snr_list.size):
    here = np.zeros(9)
    here[:5] = intrinsics[i,:5]
    here[5:8] = angulars[i,:]
    here[8] = intrinsics[i,5]

    samp_x_out[i,:] = here

np.save(save_dir.format('xsamp.npy'),arr=samp_x_out[inds_to_keep,:])
