import torch
import numpy as np
import matplotlib.pyplot as plt
import time
from EMRI_DET.utilities import norm, norm_inputs, unnorm
from EMRI_DET.nn.model_creation import load_mlp

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

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

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)
normed_input = normed_input.float().to(device)
ynorm = norm(ydata, ymeanstd[0], ymeanstd[1])

st = time.perf_counter()
nb = 50

with torch.no_grad():
    out = []
    for i in range(nb):
        output = mlp(normed_input[i * ydata.size // nb : (i+1)*ydata.size // nb])
        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]))

threshold = np.median(np.exp(ydata))

out_classified = np.zeros(shape=ydata.size)
out_classified[out_unnorm.flatten() >= threshold] = 1

truth_classified = np.zeros(shape=ydata.size)
truth_classified[np.exp(ydata) >= threshold] = 1

plt.hist(out_classified,bins=2,histtype='step')
plt.hist(truth_classified, bins=2,histtype='step')
plt.show()

print('Root mean square error (unnorm): ',np.sqrt(np.mean((out_unnorm.flatten() - np.exp(ydata))**2)))
print('Root mean square error (norm): ',np.sqrt(np.mean((output.flatten() - ynorm)**2)))
print('Accuracy: ', 1 - np.mean(np.abs(out_classified - truth_classified)))


# out_plot = out_unnorm.detach().numpy().flatten()
# diffs = abs(np.log10(abs(out_plot/ydata)))
#
# plt.hist(diffs, bins='auto')
# plt.xlabel('|log10(ratio)| between prediction and truth')
# plt.show()
#
# diffs2 = np.log10(abs(out_plot-ydata))
#
# plt.hist(diffs2, bins='auto')
# plt.xlabel('log10(|difference|) between prediction and truth')
# plt.show()

# plt.hist(np.exp(out_plot), bins='auto', label='NN', alpha=0.5, density=True)
# plt.hist(np.exp(ydata), bins='auto', label='Truth', alpha=0.5, density=True)
# plt.legend()
# plt.xlabel('Mock SNR')
# plt.show()