import torch
import torch.nn as nn
from tqdm import tqdm
import random
import numpy as np
from omegaconf import OmegaConf
from AutoEncoder_models import MLPAE
from create_data_loader import anomaly_data_loader_celeba
from train_test import train_and_eval_per_epoch_anomalies
from utils import l2_norm_model_weights
import pandas as pd
import os

config = OmegaConf.load('config.yaml')
seed = config.seed
batch_size = config.hyperparameters.batch_size
lr = config.hyperparameters.lr
opt = config.hyperparameters.opt
epochs = config.hyperparameters.epochs
n_samples = config.n_samples
noise_per = config.noise.noise_per

exp_name = 'celeba_anomalies_with_anomaly_per'
for latent_dim in [25]:#range(5, 40, 5):
    for hidden_dim in range(4, 404, 4):
        torch.manual_seed(seed)
        random.seed(seed)
        np.random.seed(seed)
        train_loader, test_regular_loader, test_anomalies_loader = anomaly_data_loader_celeba(n_samples, noise_per,
                                                                                              batch_size)

        print(f'num train={len(train_loader.dataset)}')
        print(f'hidden_dim={hidden_dim}')
        print(f'latent_dim={latent_dim}')
        print(f'anomaly_per={noise_per}')
        print(f'seed={seed}')

        model = MLPAE(input_dim=train_loader.dataset[0][0].shape[0], latent_dim=latent_dim, hidden_dim=hidden_dim,
                       n_hidden_layers=0, final_activation=nn.Identity())

        total_params = sum(param.numel() for param in model.parameters())
        print(f'total_params={total_params}')
        run_results = []
        min_train_loss = np.inf
        min_test_regular_loss = np.inf
        min_test_anomaly_loss = np.inf
        min_roc_auc = np.inf
        for epoch in tqdm(range(epochs)):
            train_loss, test_regular_loss, test_anomaly_loss, roc_auc = train_and_eval_per_epoch_anomalies(train_loader,
                                                                                                 test_regular_loader,
                                                                                                 test_anomalies_loader,
                                                                                                 model, opt, lr)
            min_train_loss = min(min_train_loss, train_loss)
            min_test_regular_loss = min(min_test_regular_loss, test_regular_loss)
            min_test_anomaly_loss = min(min_test_anomaly_loss, test_anomaly_loss)
            min_roc_auc = min(min_roc_auc, roc_auc)
            l2_norm = l2_norm_model_weights(model)
            if (epoch + 1) % 10 == 0:
                print(f'Epoch: {epoch + 1}, Train loss: {round(train_loss, 3)}, clean test loss: '
                      f'{round(test_regular_loss, 3)}, anomaly test loss: {round(test_anomaly_loss, 3)}, roc auc: '
                      f'{round(roc_auc, 3)}')

            results_per_epoch = {
                'train_loss': train_loss,
                'clean_test_loss': test_regular_loss,
                'anomaly_test_loss': test_anomaly_loss,
                'roc_auc': roc_auc,
                'epoch': epoch,
                'n_model_params': total_params,
                'hidden_dim': hidden_dim,
                'latent_dim': latent_dim,
                'weights_l2_norm': l2_norm,
                'min_train_loss': min_train_loss,
                'min_test_regular_loss': min_test_regular_loss,
                'min_test_anomaly_loss': min_test_anomaly_loss,
                'min_roc_auc': min_roc_auc,
                'seed': seed,
                'anomaly_per': noise_per,
                'n_samples': n_samples,
                'batch_size': batch_size
            }
            run_results.append(results_per_epoch)

        run_results = pd.DataFrame(run_results)
        file_name = f'seed_{seed}_anomaly_per_{noise_per}.csv'
        if not os.path.exists(f'results/{exp_name}'):  # create the folder with the experiment name if it doesn't exist
            os.makedirs(f'results/{exp_name}')
        for root, dirs, files in os.walk(os.getcwd() + f'/results/{exp_name}'):
            if file_name not in files:  # if the directory is empty create the csv file
                run_results.to_csv(f'results/{exp_name}/{file_name}', index=False)
            else:  # if the csv file exists, read it, add to it new information and save it again
                results = pd.read_csv(f'results/{exp_name}/{file_name}')
                results = pd.concat([results, run_results])
                results.to_csv(f'results/{exp_name}/{file_name}', index=False)