print("--- INITIALIZING LIBRARIES ---")
import torch
from torch import nn

# GENERAL PURPOSE
import numpy as np
from datetime import datetime
import pandas as pd
import os
import tomli as tl

# FOR TRAINING
import torch.optim as optim
import utils.models as models
from utils.training import train_model, test_model, initalize_dataloaders, plot_results
print("--- LIBRARIES INITIALIZED ---")

#GET CONFIG SETTINGS
if os.getenv('ADL_CONFIG_PATH') is None:
    with open ('config.toml', 'rb') as f:
        config = tl.load(f)
else:
    with open(os.getenv('ADL_CONFIG_PATH'), 'rb') as f:
        config = tl.load(f)

cuda_device = torch.device(config['cuda']['device'])
torch.set_default_device(cuda_device)

model_name = config['model']['name']

print("--- RUNNING ---")
print("Pytorch Version: " + torch. __version__)

# data & training properties:
val_split = config['dataset']['validation_split']     # % of val and test, rest will be train
runs = config['training']['runs']
epochs = config['training']['epochs']
seeds = [np.random.randint(0, 1000) for _ in range(runs)]

#paths
mri_path = config['paths']['mri_data']
xls_path = config['paths']['xls_data']
saved_model_path = config['paths']['model_output']
plot_path = config['paths']['plot_output']
training_record_path = config['paths']['training_record_output']

#model
model_CNN = models.CNN_Net(1, 2, config['model']['droprate']).to(cuda_device)
criterion = nn.BCELoss()
optimizer = optim.Adam(model_CNN.parameters(), lr=config['training']['learning_rate'])


for seed in seeds:
    #get time stamp for model
    time_stamp = datetime.now().strftime('%Y%m%d+%H%M%S')

    #initialize dataloaders, train model, and test model
    train_loader, val_loader, test_loader, _ = initalize_dataloaders(mri_path, xls_path, val_split, seed, cuda_device=cuda_device, batch_size=config['training']['batch_size'])
    

    print("--- TRAINING MODEL ---")
    print("Seed: ", seed)
    
        
    train_results = train_model(model_CNN, seed, time_stamp, epochs, train_loader, val_loader, saved_model_path, model_name, optimizer, criterion, cuda_device=cuda_device)
    print("--- TESTING MODEL ---")
    predicted, actual, correct, incorrect, _, _ = test_model(model_CNN, test_loader, cuda_device=cuda_device)
    
    print("Accuracy: " + str(correct / (correct + incorrect)))
    
    #Plot results and confusion matrix
    plot_results(train_results["train_acc"], train_results["train_loss"], train_results["val_acc"], train_results["val_loss"], model_name, time_stamp, plot_path)
    
    #Save training results
    if not os.path.exists(training_record_path):
        os.makedirs(training_record_path)
    train_results.to_csv(training_record_path + model_name + "_t-" + time_stamp + "_s-" + str(seed) + "_e-" + str(epochs) + ".csv")
    
    
    

print("--- END ---")