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
import pandas as pd

if __name__ == '__main__':
    device = "cuda:0"
    fp = '../schwarz_negY/{}'
    #['logM', 'logq', 'a', 'p0', 'e', 'Y0', 'thetaS', 'phiS', 'thetaK', 't', 'SNR']
    train_inds = [0,1,2,4,5,6,7,8,9]
    test_inds = [10]

    traindata = pd.read_csv(fp.format('samp_dataframe.csv'))
    xtrain = traindata.iloc[:,train_inds].to_numpy()
    ytrain = traindata.iloc[:,test_inds].to_numpy().flatten()

    testdata = pd.read_csv(fp.format('grid_dataframe.csv'))
    xtest = testdata.iloc[:, train_inds].to_numpy(  )
    ytest = testdata.iloc[:, test_inds].to_numpy().flatten()

    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')

    data = [xtrain, ytrain, xtest, ytest]

    loss_function = nn.MSELoss()
    LR = 2e-3

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