import cupy
import numpy as np
from few.waveform import GenerateEMRIWaveform, RunKerrGenericPn5Inspiral
from utility import snr, brentq_p_at_t
from pathlib import Path
from sys import stdout
import cupy as cp
from scipy.constants import year
import time
import pandas as pd

mpool = cupy.get_default_memory_pool()
mpool2 = cupy.get_default_pinned_memory_pool()

outdir = '../emri_data/schwarz_negY/{}'
Path('../emri_data/schwarz_negY/').mkdir(parents=True, exist_ok=True)

use_gpu = True
use_schwarz_separatrix = False  # set to False for negative Y, as Kerr separatrix exceeds Schwarz in this case

kerr = GenerateEMRIWaveform(
    "Pn5AAKWaveform",
    sum_kwargs=dict(pad_output=True),
    inspiral_kwargs={"max_init_len": int(1e6), "enforce_schwarz_sep": use_schwarz_separatrix},
    use_gpu=use_gpu,
    return_list=True,
)

kerr_not_list = GenerateEMRIWaveform(
    "Pn5AAKWaveform",
    sum_kwargs=dict(pad_output=True),
    inspiral_kwargs={"max_init_len": int(1e6), "enforce_schwarz_sep": use_schwarz_separatrix},
    use_gpu=use_gpu,
    return_list=False,
)

# Max waveform duration and sample spacing
T = 4  # years
Tsec = T * year
dt = 15.0

# for SNR and waveform calculations
inner_product_kwargs = dict(dt=dt, PSD="cornish_lisa_psd", use_gpu=True)
waveform_kwargs = {"T": T, "dt": dt}

# define our injection set
nM = 6  # 7
nq = 5  # 5
na = 4  # 4
ne = 4  # 4
ny = 4  # 4
nTheta = 4  # 5
nPhi = 5  # 8
nTimes = 5  # 8
nDists = 10  # dont change

Mvals = np.logspace(np.log10(8e4), np.log10(5e7), nM)
qvals = np.logspace(np.log10(5e4), np.log10(5e7), nq)
avals = np.linspace(0.5, 0.95, na)
evals = np.linspace(0.08, 0.5, ne)
Y0vals = np.linspace(-0.5, -0.99, ny)

dVals = np.linspace(0, 5, nDists + 1)[1:]
dist = dVals[1]  # our dummy distance

tplVals = np.linspace(4, 0, nTimes + 1)[:-1]  # we can fill in the last value here
thetaVals = np.linspace(0, np.pi, nTheta + 1)[:-1]  # we can fill in the last value here
phiVals = np.linspace(0, 2 * np.pi, nPhi + 1)[:-1]  # we can fill in the last value here

# initialise trajectory module
traj = RunKerrGenericPn5Inspiral(inspiral_kwargs={"max_init_len": int(1e9)}, enforce_schwarz_sep=use_schwarz_separatrix)

traj_inds = np.array([0, 1, 2, 4, 5, 11, 12, 13]).astype(np.int32)  # parameter indices needed for trajectory calc

totforms = nM * nq * na * ne * ny
num_wform_samples = int(T * year / dt)

inp_angles = []
for num1 in thetaVals:
    for num2 in thetaVals:
        for num3 in phiVals:
            inp_angles.append([num1, num2, num3])


# np.save(outdir.format('angular_param_perms.npy'),arr=np.array(inp_angles))


def wform_snr_at_t(wform, plunge_time):
    ind_to_cut = int((1 - plunge_time / 4) * num_wform_samples)
    cut_waveform = wform.copy()[ind_to_cut:]  # truncate waveforms (after a copy operation)
    snrval = snr([cp.array(cut_waveform.real), cp.array(cut_waveform.imag)],
                 **inner_product_kwargs).get()  # need to send waveform to gpu for SNR calc
    return snrval


cols = ['logM', 'logq', 'a', 'p0', 'e', 'Y0', 'thetaS', 'phiS', 'thetaK', 't', 'SNR']
try:  # pick up where we left off.
    dataframe = pd.read_csv(outdir.format('grid_dataframe.csv'))
    place = np.load(outdir.format('grid_position.npy')).item()
    wavegen_times = np.load(outdir.format('grid_waveform_times.npy')).tolist()
    process_times = np.load(outdir.format('grid_processing_times.npy')).tolist()

except:
    dataframe = pd.DataFrame(columns=cols)
    place = -1
    wavegen_times = []
    process_times = []

'''
We now iterate through the grid dimensions. To divide things up, we process the extrinsic parameters in a block before saving to the array.
This is a nice balance between caching data (in case of error) and not spending too long appending to the dataframe.
'''

injection_params = np.zeros(14)
for i in range(nM):
    for j in range(nq):
        for k in range(na):
            for l in range(ne):
                for m in range(ny):
                    num = i * nq * na * ne * ny + j * na * ne * ny + k * ne * ny + l * ny + m
                    if num <= place:
                        continue
                    stdout.write(
                        f'\rWaveform sets generated: {num + 1:d} out of {totforms:d}, or {100 * num / (totforms - 1):.2f}% done.')
                    stdout.flush()

                    # update injection parameters
                    injection_params[0] = Mvals[i]
                    injection_params[1] = injection_params[0] / qvals[j]
                    injection_params[2] = avals[k]
                    injection_params[4] = evals[l]
                    injection_params[5] = Y0vals[m]
                    injection_params[6] = dist

                    root = brentq_p_at_t(traj, T, traj_args=np.take(injection_params, traj_inds).tolist(),
                                         traj_kwargs={'max_init_len': int(1e9)},
                                         kerr_separatrix=not use_schwarz_separatrix,
                                         xtol=1e-3)  # find initial semi-latus rectum for T year orbit
                    injection_params[3] = root

                    out_list = []
                    wave_time = 0
                    process_time = 0

                    this_chunk = np.zeros((int(len(inp_angles) * nTimes), 11))
                    this_chunk[:, 0] = np.log10(injection_params[0])
                    this_chunk[:, 1] = np.log10(qvals[j])
                    this_chunk[:, 2:6] = injection_params[2:6]

                    for iterno, inp_set in enumerate(inp_angles):
                        stdout.write(
                            f' \rWaveform sets generated: {num + 1:d} out of {totforms:d}, or {100 * num / (totforms - 1):.2f}% done. Computing waveforms... ({iterno} of {nTheta * nTheta * nPhi})')
                        stdout.flush()
                        injection_params[7] = inp_set[0]
                        injection_params[9] = inp_set[1]
                        injection_params[8] = inp_set[2]

                        wave_start = time.perf_counter()
                        check_sig = kerr_not_list(*injection_params, **waveform_kwargs).get()
                        wave_end = time.perf_counter()
                        wave_time += wave_end - wave_start

                        process_start = time.perf_counter()
                        for time_ind, pt in enumerate(tplVals):
                            pos = int(iterno * nTimes + time_ind)
                            snr_out = wform_snr_at_t(check_sig, pt)
                            this_chunk[pos, 6:9] = injection_params[7:10]
                            this_chunk[pos, 9] = pt
                            this_chunk[pos, 10] = snr_out
                            # clear memory (or else INTERNAL_FFT_ERROR bug)
                            cp.fft.config.get_plan_cache().clear()
                            mpool.free_all_blocks()
                            mpool2.free_all_blocks()

                        process_end = time.perf_counter()
                        process_time += process_end - process_start

                    wavegen_times.append(wave_time)
                    stdout.write(f' done. Processing...')
                    stdout.flush()

                    little_dataframe = pd.DataFrame(this_chunk, columns=cols)
                    dataframe = dataframe.append(little_dataframe, ignore_index=True)
                    process_times.append(process_time)

                    # cache our progress
                    dataframe.to_csv(outdir.format('grid_dataframe.csv'), index=False)
                    np.save(outdir.format('grid_waveform_times.npy'), arr=np.array(wavegen_times))
                    np.save(outdir.format('grid_processing_times.npy'), arr=np.array(process_times))
                    np.save(outdir.format('grid_position.npy'), arr=np.array(num))
