
import os
import argparse
import shutil
import timeit
import torch
import torchvision
import torchvision.transforms as transforms
import numpy as np
import torch.nn as nn
import torch.backends.cudnn as cudnn
import random
import matplotlib.pyplot as plt
import matplotlib as mpl
from torch import autograd
from torchvision.utils import save_image
import csv
from tqdm import tqdm
import gc
import h5py


#############################
# Settings
#############################

parser = argparse.ArgumentParser(description='Pre-train AE for computing FID')
parser.add_argument('--root_path', type=str, default='')
parser.add_argument('--data_path', type=str, default='')
parser.add_argument('--num_workers', type=int, default=0)
parser.add_argument('--dim_bottleneck', type=int, default=512)
parser.add_argument('--epochs', type=int, default=200, metavar='N',
                    help='number of epochs to train CNNs (default: 200)')
parser.add_argument('--resume_epoch', type=int, default=0)
parser.add_argument('--batch_size_train', type=int, default=128, metavar='N',
                    help='input batch size for training')
parser.add_argument('--batch_size_valid', type=int, default=10, metavar='N',
                    help='input batch size for testing')
parser.add_argument('--base_lr', type=float, default=1e-3,
                    help='learning rate, default=1e-3')
parser.add_argument('--lr_decay_epochs', type=int, default=50) #decay lr rate every dre_lr_decay_epochs epochs
parser.add_argument('--lr_decay_factor', type=float, default=0.1)
parser.add_argument('--lambda_sparsity', type=float, default=1e-4, help='penalty for sparsity')
parser.add_argument('--weight_dacay', type=float, default=1e-4,
                    help='Weigth decay, default=1e-4')
parser.add_argument('--seed', type=int, default=2020, metavar='S',
                    help='random seed (default: 1)')
parser.add_argument('--CVMode', action='store_true', default=False,
                   help='CV mode?')
parser.add_argument('--img_size', type=int, default=128, metavar='N')
parser.add_argument('--min_label', type=float, default=0.0)
parser.add_argument('--max_label', type=float, default=90.0)
args = parser.parse_args()

wd = args.root_path
os.chdir(wd)
from models import *
from utils import IMGs_dataset, SimpleProgressBar

# some parameters in the opts
dim_bottleneck = args.dim_bottleneck
epochs = args.epochs
base_lr = args.base_lr
lr_decay_epochs = args.lr_decay_epochs
lr_decay_factor = args.lr_decay_factor
resume_epoch = args.resume_epoch
lambda_sparsity = args.lambda_sparsity


# random seed
random.seed(args.seed)
torch.manual_seed(args.seed)
torch.backends.cudnn.deterministic = True
cudnn.benchmark = False
np.random.seed(args.seed)

# directories for checkpoint, images and log files
save_models_folder = wd + '/output/eval_models'
os.makedirs(save_models_folder, exist_ok=True)
save_AE_images_in_train_folder = save_models_folder + '/AE_lambda_{}_images_in_train'.format(lambda_sparsity)
os.makedirs(save_AE_images_in_train_folder, exist_ok=True)
save_AE_images_in_valid_folder = save_models_folder + '/AE_lambda_{}_images_in_valid'.format(lambda_sparsity)
os.makedirs(save_AE_images_in_valid_folder, exist_ok=True)


###########################################################################################################
# Data
###########################################################################################################
# data loader
data_filename = args.data_path + '/RC-49_' + str(args.img_size) + 'x' + str(args.img_size) + '.h5'
hf = h5py.File(data_filename, 'r')
labels = hf['labels'][:]
labels = labels.astype(float)
images = hf['images'][:]
hf.close()
N_all = len(images)
assert len(images) == len(labels)

q1 = args.min_label
q2 = args.max_label
indx = np.where((labels>q1)*(labels<q2)==True)[0]
labels = labels[indx]
images = images[indx]
assert len(labels)==len(images)

# define training and validation sets
if args.CVMode:
    #90% Training; 10% valdation
    valid_prop = 0.1 #proportion of the validation samples
    indx_all = np.arange(len(images))
    np.random.shuffle(indx_all)
    indx_valid = indx_all[0:int(valid_prop*len(images))]
    indx_train = indx_all[int(valid_prop*len(images)):]

    trainset = IMGs_dataset(images[indx_train], labels=None, normalize=True)
    trainloader = torch.utils.data.DataLoader(trainset, batch_size=args.batch_size_train, shuffle=True, num_workers=args.num_workers)
    validset = IMGs_dataset(images[indx_valid], labels=None, normalize=True)
    validloader = torch.utils.data.DataLoader(validset, batch_size=args.batch_size_valid, shuffle=False, num_workers=args.num_workers)

else:
    trainset = IMGs_dataset(images, labels=None, normalize=True)
    trainloader = torch.utils.data.DataLoader(trainset, batch_size=args.batch_size_train, shuffle=True, num_workers=args.num_workers)


###########################################################################################################
# Necessary functions
###########################################################################################################

def adjust_learning_rate(epoch, epochs, optimizer, base_lr, lr_decay_epochs, lr_decay_factor):
    lr = base_lr #1e-4

    for i in range(epochs//lr_decay_epochs):
        if epoch >= (i+1)*lr_decay_epochs:
            lr *= lr_decay_factor

    for param_group in optimizer.param_groups:
        param_group['lr'] = lr

def train_AE():

    # define optimizer
    params = list(net_encoder.parameters()) + list(net_decoder.parameters())
    optimizer = torch.optim.Adam(params, lr = base_lr, betas=(0.5, 0.999), weight_decay=1e-4)

    # criterion
    criterion = nn.MSELoss()

    if resume_epoch>0:
        print("Loading ckpt to resume training AE >>>")
        ckpt_fullpath = save_models_folder + "/AE_checkpoint_intrain/AE_checkpoint_epoch_{}_lambda_{}.pth".format(resume_epoch, lambda_sparsity)
        checkpoint = torch.load(ckpt_fullpath)
        net_encoder.load_state_dict(checkpoint['net_encoder_state_dict'])
        net_decoder.load_state_dict(checkpoint['net_decoder_state_dict'])
        optimizer.load_state_dict(checkpoint['optimizer_state_dict'])
        torch.set_rng_state(checkpoint['rng_state'])
        gen_iterations = checkpoint['gen_iterations']
    else:
        gen_iterations = 0

    start_time = timeit.default_timer()
    for epoch in range(resume_epoch, epochs):

        adjust_learning_rate(epoch, epochs, optimizer, base_lr, lr_decay_epochs, lr_decay_factor)

        train_loss = 0

        for batch_idx, batch_real_images in enumerate(trainloader):

            net_encoder.train()
            net_decoder.train()

            batch_size_curr = batch_real_images.shape[0]

            batch_real_images = batch_real_images.type(torch.float).cuda()


            batch_features = net_encoder(batch_real_images)
            batch_recons_images = net_decoder(batch_features)

            '''
            based on https://debuggercafe.com/sparse-autoencoders-using-l1-regularization-with-pytorch/
            '''
            loss = criterion(batch_recons_images, batch_real_images) + lambda_sparsity * batch_features.mean()

            #backward pass
            optimizer.zero_grad()
            loss.backward()
            optimizer.step()

            train_loss += loss.cpu().item()

            gen_iterations += 1

            if gen_iterations % 100 == 0:
                n_row=min(10, int(np.sqrt(batch_size_curr)))
                with torch.no_grad():
                    batch_recons_images = net_decoder(net_encoder(batch_real_images[0:n_row**2]))
                    batch_recons_images = batch_recons_images.detach().cpu()
                save_image(batch_recons_images.data, save_AE_images_in_train_folder + '/{}.png'.format(gen_iterations), nrow=n_row, normalize=True)

            if gen_iterations % 20 == 0:
                print("AE+lambda{}: [step {}] [epoch {}/{}] [train loss {}] [Time {}]".format(lambda_sparsity, gen_iterations, epoch+1, epochs, train_loss/(batch_idx+1), timeit.default_timer()-start_time) )
        # end for batch_idx

        if (epoch+1) % 50 == 0:
            save_file = save_models_folder + "/AE_checkpoint_intrain/AE_checkpoint_epoch_{}_lambda_{}.pth".format(epoch+1, lambda_sparsity)
            os.makedirs(os.path.dirname(save_file), exist_ok=True)
            torch.save({
                    'gen_iterations': gen_iterations,
                    'net_encoder_state_dict': net_encoder.state_dict(),
                    'net_decoder_state_dict': net_decoder.state_dict(),
                    'optimizer_state_dict': optimizer.state_dict(),
                    'rng_state': torch.get_rng_state()
            }, save_file)
    #end for epoch

    return net_encoder, net_decoder


if args.CVMode:
    def valid_AE():
        net_encoder.eval()
        net_decoder.eval()
        with torch.no_grad():
            for batch_idx, images in enumerate(validloader):
                images = images.type(torch.float).cuda()
                features = net_encoder(images)
                recons_images = net_decoder(features)
                save_image(recons_images.data, save_AE_images_in_valid_folder + '/{}_recons.png'.format(batch_idx), nrow=10, normalize=True)
                save_image(images.data, save_AE_images_in_valid_folder + '/{}_real.png'.format(batch_idx), nrow=10, normalize=True)
        return None



###########################################################################################################
# Training and validation
###########################################################################################################

# model initialization
net_encoder = encoder(dim_bottleneck=args.dim_bottleneck).cuda()
net_decoder = decoder(dim_bottleneck=args.dim_bottleneck).cuda()
net_encoder = nn.DataParallel(net_encoder)
net_decoder = nn.DataParallel(net_decoder)

filename_ckpt = save_models_folder + '/ckpt_AE_epoch_{}_seed_{}_CVMode_{}.pth'.format(args.epochs, args.seed, args.CVMode)

# training
if not os.path.isfile(filename_ckpt):
    print("\n Begin training AE: ")
    start = timeit.default_timer()
    net_encoder, net_decoder = train_AE()
    stop = timeit.default_timer()
    print("Time elapses: {}s".format(stop - start))
    # save model
    torch.save({
    'net_encoder_state_dict': net_encoder.state_dict(),
    'net_decoder_state_dict': net_decoder.state_dict(),
    }, filename_ckpt)
else:
    print("\n Ckpt already exists")
    print("\n Loading...")
    checkpoint = torch.load(filename_ckpt)
    net_encoder.load_state_dict(checkpoint['net_encoder_state_dict'])
    net_decoder.load_state_dict(checkpoint['net_decoder_state_dict'])

if args.CVMode:
    #validation
    _ = valid_AE()
