import torch
import numpy as np
import matplotlib.pyplot as plt
from EMRI_DET.utilities import norm, norm_inputs, unnorm_inputs, unnorm, get_script_path
from sys import stdout


def model_train_test(data, model, device, n_epochs, n_batches, loss_function, learning_rate, verbose=False, return_losses=False):
    xtrain, ytrain, xtest, ytest = data

    name = model.name
    path = get_script_path()

    np.save(path+'/../models/'+name+'/xdata_mean_std.npy',np.array([xtrain.mean(axis=0), xtrain.std(axis=0)]))
    np.save(path+'/../models/'+name+'/ydata_mean_std.npy',np.array([ytrain.mean(), ytrain.std()]))

    xtest = torch.from_numpy(norm_inputs(xtest, ref_dataframe=xtrain)).to(device).float()
    ytest = torch.from_numpy(norm(ytest, ref_dataframe=ytrain)).to(device).float()
    xtrain = torch.from_numpy(norm_inputs(xtrain, ref_dataframe=xtrain)).to(device).float()
    ytrain = torch.from_numpy(norm(ytrain, ref_dataframe=ytrain)).to(device).float()

    ytrainsize = len(ytrain)
    ytestsize = len(ytest)

    optimizer = torch.optim.Adam(model.parameters(), lr=learning_rate)
    scheduler = torch.optim.lr_scheduler.ExponentialLR(optimizer, gamma=0.97)  # to settle into the loss 'valley'
    train_losses = []
    test_losses = []
    rate = []
    # Run the training loop

    datasets = {"train": [xtrain, ytrain], "test": [xtest, ytest]}

    cutoff_LR = n_epochs - 50
    for epoch in range(n_epochs):  # 5 epochs at maximum
        # Print epoch
        for phase in ['train','test']:
            if phase == 'train':
                model.train(True)
                shuffled_inds = torch.randperm(ytrainsize)

                # Set current loss value
                current_loss = 0.0

                # Iterate over the DataLoader for training data
                # Get and prepare inputs
                inputs, targets = datasets[phase]
                inputs = inputs[shuffled_inds]
                targets = targets[shuffled_inds]

                targets = targets.reshape((targets.shape[0], 1))

                for i in range(n_batches):
                    for param in model.parameters():
                        param.grad = None
                    outputs = model(inputs[i * ytrainsize // n_batches:(i+1)*ytrainsize // n_batches])
                    loss = torch.sqrt(loss_function(outputs, targets[i * ytrainsize // n_batches: (i+1)*ytrainsize // n_batches]))
                    loss.backward()
                    optimizer.step()
                    current_loss += loss.item()

                train_losses.append(current_loss / n_batches)

            else:
                with torch.no_grad():
                    model.train(False)
                    shuffled_inds = torch.randperm(ytestsize)
                    current_loss = 0.0
                    inputs, targets = datasets[phase]
                    inputs = inputs[shuffled_inds]
                    targets = targets[shuffled_inds]

                    targets = targets.reshape((targets.shape[0], 1))

                    for i in range(n_batches):
                        outputs = model(inputs[i * ytestsize // n_batches: (i+1)*ytestsize // n_batches])
                        loss = torch.sqrt(loss_function(outputs, targets[i * ytestsize // n_batches: (i+1)*ytestsize // n_batches]))
                        current_loss += loss.item()

                    test_losses.append(current_loss / n_batches)

        if epoch >= cutoff_LR:
            scheduler.step()
            rate.append(scheduler.get_last_lr()[0])
        else:
            rate.append(learning_rate)
        if verbose:
            stdout.write(f'\rEpoch: {epoch} | Train loss: {train_losses[-1]:.3e} | Test loss: {test_losses[-1]:.3e} ')

    if verbose:
        print('\nTraining complete - saving.')
    torch.save(model.state_dict(),path+'/'+name+'/model.pth')

    epochs = np.arange(n_epochs)
    plt.semilogy(epochs, train_losses, label='Train')
    plt.semilogy(epochs, test_losses, label='Test')
    plt.legend()
    plt.xlabel('Epochs')
    plt.ylabel('Loss')
    plt.title('Train and Test Loss Across Train Epochs')
    plt.savefig(path+'/../models/'+name+'/losses.png')
    #plt.show()
    plt.close()

    out = (model,)
    if return_losses:
        out += (train_losses, test_losses,)
    return out
