from torch import nn
import numpy as np
from EMRI_DET.nn.model_train_test import model_train_test
from EMRI_DET.nn.model_creation import create_mlp

if __name__ == '__main__':
    device = "cuda:0"

    in_features = 9
    out_features = 1

    layers = 4
    neurons = [256,128,64,32]
    activation = nn.Tanh
    model = create_mlp(input_features=in_features,output_features=out_features,neurons=neurons,layers=layers,activation=activation,device=device, model_name='m1')


    fp = '../_data/{}'

    xtrain = np.load(fp.format('xsamp.npy'))
    ytrain = np.log(np.load(fp.format('ysamp.npy')))
    xtest = np.load(fp.format('xgrid.npy'))
    ytest = np.log(np.load(fp.format('ygrid.npy')))

    data = [xtrain, ytrain, xtest, ytest]

    loss_function = nn.MSELoss()

    LR = 1e-3

    model_train_test(data, model, device,n_epochs=500,n_batches=50,loss_function=loss_function,learning_rate=LR,verbose=True)


