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 gaussian_data_loader
from train_test import train_and_eval_per_epoch
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
higher_dim = config.higher_dim
noise_per = config.noise.noise_per
snr_db = config.noise.snr_db
scenario = config.scenario
shift = config.hyperparameters.shift

exp_name = 'sample_wise_gaussian_feature_noise'#'sample_wise_gaussian_sample_noise'
for n_samples in [100, 200, 500, 1000, 1500, 2000, 2500, 3000, 3500, 4000, 4500, 5000, 5500, 6000, 7000, 8000, 9000,
                  10000, 15000, 20000, 25000, 30000, 35000, 40000, 45000, 50000]:
    for latent_dim in [35]:
        for hidden_dim in [150]:#range(4, 504, 4):
            torch.manual_seed(seed)
            random.seed(seed)
            np.random.seed(seed)
            train_loader, test_loader = gaussian_data_loader(n_samples, higher_dim, 20, batch_size, scenario, noise_per,
                                                             snr_db, shift)

            print(f'hidden_dim={hidden_dim} \n latent_dim={latent_dim} \n n_samples={len(train_loader.dataset)} '
                  f'\n noise_per={noise_per} \n snr_db={snr_db} \n seed={seed}')
            if scenario == 'domain_shift':
                print(f'shift={shift}')

            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_loss = np.inf
            for epoch in tqdm(range(epochs)):
                train_loss, test_loss, _, _ = train_and_eval_per_epoch(train_loader, test_loader, model, opt, lr)
                min_train_loss = min(min_train_loss, train_loss)
                min_test_loss = min(min_test_loss, test_loss)
                l2_norm = l2_norm_model_weights(model)
                if (epoch + 1) % 10 == 0:
                    print(f'Epoch: {epoch + 1}, Train loss: {train_loss}, test loss: {test_loss}')

                results_per_epoch = {
                    'train_loss': train_loss,
                    'test_loss': test_loss,
                    'epoch': epoch,
                    'n_model_params': total_params,
                    'hidden_dim': hidden_dim,
                    'latent_dim': latent_dim,
                    'n_samples': len(train_loader.dataset),
                    'weights_l2_norm': l2_norm,
                    'min_train_loss': min_train_loss,
                    'min_test_loss': min_test_loss,
                    'seed': seed,
                    'noise_per': noise_per,
                    'snr_db': snr_db
                }
                if scenario == 'domain_shift':
                    results_per_epoch['shift'] = shift

                run_results.append(results_per_epoch)

            run_results = pd.DataFrame(run_results)
            if scenario == 'domain_shift':
                file_name = f'seed_{seed}_noise_per_{noise_per}_snr_{snr_db}_shift_{shift}_latent_{latent_dim}_' \
                            f'hidden_{hidden_dim}_samples_{len(train_loader.dataset)}.csv'
            else:
                file_name = f'seed_{seed}_noise_per_{noise_per}_snr_{snr_db}_latent_{latent_dim}_' \
                            f'hidden_{hidden_dim}_samples_{len(train_loader.dataset)}.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)
