import numpy as np
from scipy.interpolate import RegularGridInterpolator
import cupy as cp
from cupyx.scipy.ndimage import map_coordinates
import time
import matplotlib.pyplot as plt
import torch
from EMRI_DET.nn.model_creation import load_mlp
from EMRI_DET.utilities import norm, unnorm_inputs, norm_inputs, unnorm


xgrid = np.load('../data/xgrid.npy')
ygrid = np.load('../data/ygrid.npy')

xdata = np.load('../data/xtest.npy')
ydata = np.load('../data/ytest.npy')

nM = 8
nq = 5
na = 4
ne = 4
nY = 4
nTh = 5
nPh = 8
nt = 5

xdata_in = ()
for i in range(xgrid.shape[1]):
    xdata_in += (np.unique(xgrid[:,i]),)

y_in = ygrid.reshape((nM,nq,na,ne,nY,nTh,nPh,nTh,nt))

xdata_in2 = xgrid.copy()

print('Beginning Linear interpolator init')
st = time.perf_counter()

interp = RegularGridInterpolator(xdata_in, y_in, method='linear',bounds_error=False)
nn_interp = RegularGridInterpolator(xdata_in, y_in, method='nearest',bounds_error=False)

dat1 = interp(xdata_in2[:1000])
print('Consistency check:', np.sum(dat1 - ygrid.flatten()[:1000]))

mt = time.perf_counter()
print('Starting interpolation')

lin_values = interp(xdata)
nearest_values = nn_interp(xdata)
et = time.perf_counter()
RMSE = np.sqrt(np.nanmean((lin_values- ydata)**2))

plt.hist(np.log10(abs(lin_values-ydata)), bins='auto', alpha=0.5, label='Linear Interp. - 1e6 points',histtype='step',density=True)
plt.hist(np.log10(abs(nearest_values-ydata)), bins='auto', alpha=0.5, label='Nearest Neighbour - 1e6 points',histtype='step',density=True)

print('Linear grid interpolator RMSE: ', RMSE)
print(f'Time for [interpolator init, interpolation] | [{mt-st:.3e}, {et-mt:.3e}]')

Mvals = np.linspace(np.log10(8e4), np.log10(5e7), nM)
qvals = np.linspace(np.log10(5e4), np.log10(5e7), nq)
avals = np.linspace(0.01, 0.99, na)
evals = np.linspace(0.01, 0.5, ne)
Yvals = np.linspace(0.5, 0.99, nY)
qSvals = np.linspace(0, np.pi, nTh)
qKvals = np.linspace(0, np.pi, nTh)
phiSvals = np.linspace(0, 2 * np.pi, nPh)
tVals = np.linspace(0, 4, nt + 1)[1:]

ydata_reshaped = ygrid.reshape((nM,nq,na,ne,nY,nTh,nPh,nTh,nt))

xtest_in = []

lists = [Mvals, qvals, avals, evals, Yvals, qSvals, phiSvals, qKvals, tVals]
norms = [listy[-1] - listy[0] for listy in lists]
nums = [nM, nq, na, ne, nY, nTh, nPh, nTh, nt]

for i in range(ydata.size):
    sub = []
    for j in range(9):
        sub.append((nums[j]-1)*(xdata[i, j] - lists[j][0])/norms[j])
    xtest_in.append(sub)

xtest_in = np.asarray(xtest_in).T
mt = time.perf_counter()
print('Starting interpolation')

map_values = map_coordinates(cp.array(ydata_reshaped), cp.array(xtest_in), prefilter=True).get()
et = time.perf_counter()

RMSE = np.sqrt(np.nanmean((map_values - ydata)**2))

print('map_coordinates grid interpolator RMSE: ', RMSE)
print(f'Time for interpolation | {et-mt:.3e}')

plt.hist(np.log10(abs(map_values-ydata)), bins='auto', alpha=0.5, label='n-Cubic Spine interp. - 1e6 points',histtype='step',density=True)

ydata = np.log(ydata)  # network is trained on the log of the SNR

device = 'cuda:0'
model_name = 'm1'
mlp = load_mlp(model_name, get_state_dict=True).to(device)
mlp.eval()

xmeanstd = np.load(f'../models/{model_name}_xdata_mean_std.npy')
ymeanstd = np.load(f'../models/{model_name}_ydata_mean_std.npy')

test_input = torch.Tensor(xdata)
normed_input = norm_inputs(test_input, xmeanstd[0], xmeanstd[1]).float().to(device)
st = time.perf_counter()
with torch.no_grad():
    out = []
    for i in range(20):
        output = mlp(normed_input[i * ydata.size // 20 : (i+1)*ydata.size // 20])
        out.append(output.detach().cpu().numpy())

et = time.perf_counter()
print(f'Evaluation time: {et-st:.3e}s for {ydata.size:.1e} values, {(et-st)/ydata.size:.3e}s per data point.')
output = np.concatenate(out)

out_unnorm = np.exp(unnorm(output, ymeanstd[0], ymeanstd[1]))

network_out = out_unnorm.flatten()

plt.hist(np.log10(abs(network_out-ydata)), bins='auto', alpha=0.5,label='NN - 1e6 points', histtype='step', density=True)

RMSE = np.sqrt(np.nanmean((network_out - ydata)**2))

print('Neural network interpolator RMSE: ', RMSE)

plt.legend()
plt.xlabel('log10(|diff|) between truth and interpolated')
plt.title('Neural Network (NN) performance vs. selected interpolation methods')
plt.show()
