import numpy as np
import tqdm
from losses.dsm import dsm_score_estimation
import torch.nn.functional as F
import logging
import torch
import os
import shutil
import tensorboardX
import torch.optim as optim
from torchvision.datasets import MNIST, CIFAR10, FashionMNIST
import torchvision.transforms as transforms
from torch.utils.data import DataLoader, Subset
from datasets.celeba import CelebA
from models.refinenet_dilated_baseline import RefineNetDilated
from torchvision.utils import save_image, make_grid

__all__ = ['BaselineRunner']


class BaselineRunner():
    def __init__(self, args, config):
        self.args = args
        self.config = config

    def get_optimizer(self, parameters):
        if self.config.optim.optimizer == 'Adam':
            return optim.Adam(parameters, lr=self.config.optim.lr, weight_decay=self.config.optim.weight_decay,
                              betas=(self.config.optim.beta1, 0.999), amsgrad=self.config.optim.amsgrad)
        elif self.config.optim.optimizer == 'RMSProp':
            return optim.RMSprop(parameters, lr=self.config.optim.lr, weight_decay=self.config.optim.weight_decay)
        elif self.config.optim.optimizer == 'SGD':
            return optim.SGD(parameters, lr=self.config.optim.lr, momentum=0.9)
        else:
            raise NotImplementedError('Optimizer {} not understood.'.format(self.config.optim.optimizer))

    def logit_transform(self, image, lam=1e-6):
        image = lam + (1 - 2 * lam) * image
        return torch.log(image) - torch.log1p(-image)

    def train(self):
        if self.config.data.random_flip is False:
            tran_transform = test_transform = transforms.Compose([
                transforms.Resize(self.config.data.image_size),
                transforms.ToTensor()
            ])
        else:
            tran_transform = transforms.Compose([
                transforms.Resize(self.config.data.image_size),
                transforms.RandomHorizontalFlip(p=0.5),
                transforms.ToTensor()
            ])
            test_transform = transforms.Compose([
                transforms.Resize(self.config.data.image_size),
                transforms.ToTensor()
            ])

        if self.config.data.dataset == 'CIFAR10':
            dataset = CIFAR10(os.path.join(self.args.run, 'datasets', 'cifar10'), train=True, download=True,
                              transform=tran_transform)
            test_dataset = CIFAR10(os.path.join(self.args.run, 'datasets', 'cifar10_test'), train=False, download=True,
                                   transform=test_transform)
        elif self.config.data.dataset == 'MNIST':
            dataset = MNIST(os.path.join(self.args.run, 'datasets', 'mnist'), train=True, download=True,
                            transform=tran_transform)
            test_dataset = MNIST(os.path.join(self.args.run, 'datasets', 'mnist_test'), train=False, download=True,
                                 transform=test_transform)
        elif self.config.data.dataset == 'FashionMNIST':
            dataset = FashionMNIST(os.path.join(self.args.run, 'datasets', 'fashionmnist'), train=True, download=True,
                            transform=tran_transform)
            test_dataset = FashionMNIST(os.path.join(self.args.run, 'datasets', 'fashionmnist_test'), train=False, download=True,
                                 transform=test_transform)

        elif self.config.data.dataset == 'CELEBA':
            if self.config.data.random_flip:
                dataset = CelebA(root=os.path.join(self.args.run, 'datasets', 'celeba'), split='train',
                                 transform=transforms.Compose([
                                     transforms.CenterCrop(140),
                                     transforms.Resize(self.config.data.image_size),
                                     transforms.RandomHorizontalFlip(),
                                     transforms.ToTensor(),
                                 ]), download=True)
            else:
                dataset = CelebA(root=os.path.join(self.args.run, 'datasets', 'celeba'), split='train',
                                 transform=transforms.Compose([
                                     transforms.CenterCrop(140),
                                     transforms.Resize(self.config.data.image_size),
                                     transforms.ToTensor(),
                                 ]), download=True)

            test_dataset = CelebA(root=os.path.join(self.args.run, 'datasets', 'celeba_test'), split='test',
                                  transform=transforms.Compose([
                                      transforms.CenterCrop(140),
                                      transforms.Resize(self.config.data.image_size),
                                      transforms.ToTensor(),
                                  ]), download=True)


        dataloader = DataLoader(dataset, batch_size=self.config.training.batch_size, shuffle=True, num_workers=4)
        test_loader = DataLoader(test_dataset, batch_size=self.config.training.batch_size, shuffle=True,
                                 num_workers=4, drop_last=True)

        test_iter = iter(test_loader)
        self.config.input_dim = self.config.data.image_size ** 2 * self.config.data.channels

        tb_path = os.path.join(self.args.run, 'tensorboard', self.args.doc)
        if os.path.exists(tb_path):
            shutil.rmtree(tb_path)

        tb_logger = tensorboardX.SummaryWriter(log_dir=tb_path)
        score = RefineNetDilated(self.config).to(self.config.device)

        score = torch.nn.DataParallel(score)

        optimizer = self.get_optimizer(score.parameters())

        if self.args.resume_training:
            states = torch.load(os.path.join(self.args.log, 'checkpoint.pth'))
            score.load_state_dict(states[0])
            optimizer.load_state_dict(states[1])

        step = 0

        for epoch in range(self.config.training.n_epochs):
            for i, (X, y) in enumerate(dataloader):
                step += 1

                score.train()
                X = X.to(self.config.device)
                X = X / 256. * 255. + torch.rand_like(X) / 256.


                # flip X with probability p
                if not hasattr(self.config.data, 'flip_p'): 
                    self.config.data.flip_p = 0.0

                if self.config.data.flip_p > 0.0:
                    p = torch.rand(X.shape[0], device=self.config.device).view(-1,1,1,1).repeat(1, self.config.data.channels, self.config.data.image_size, self.config.data.image_size)
                    flip_p = self.config.data.flip_p
                    X = (p > flip_p) * X + ~(p > flip_p) * (1-X)

                    
                if self.config.data.logit_transform:
                    X = self.logit_transform(X)

                loss = dsm_score_estimation(score, X, sigma=0.01)

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

                tb_logger.add_scalar('loss', loss, global_step=step)
                logging.info("step: {}, loss: {}".format(step, loss.item()))

                if step >= self.config.training.n_iters:
                    return 0

                if step % 100 == 0:
                    score.eval()
                    try:
                        test_X, test_y = next(test_iter)
                    except StopIteration:
                        test_iter = iter(test_loader)
                        test_X, test_y = next(test_iter)

                    test_X = test_X.to(self.config.device)
                    test_X = test_X / 256. * 255. + torch.rand_like(test_X) / 256.
                    if self.config.data.logit_transform:
                        test_X = self.logit_transform(test_X)

                    with torch.no_grad():
                        test_dsm_loss = dsm_score_estimation(score, test_X, sigma=0.01)

                    tb_logger.add_scalar('test_dsm_loss', test_dsm_loss, global_step=step)

                if step % self.config.training.snapshot_freq == 0:
                    states = [
                        score.state_dict(),
                        optimizer.state_dict(),
                    ]
                    torch.save(states, os.path.join(self.args.log, 'checkpoint_{}.pth'.format(step)))
                    torch.save(states, os.path.join(self.args.log, 'checkpoint.pth'))

    def Langevin_dynamics(self, x_mod, scorenet, n_steps=1000, step_lr=0.00002):
        images = []
        sigmas = np.exp(np.linspace(np.log(self.config.model.sigma_begin), np.log(self.config.model.sigma_end),
                                    self.config.model.num_classes))
        
        with torch.no_grad():
            for c, sigma in tqdm.tqdm(enumerate(sigmas), total=len(sigmas), desc='Langevin dynamics sampling'):
                labels = torch.ones(x_mod.shape[0], device=x_mod.device) * c
                labels = labels.long()
                step_size = step_lr * (sigma / sigmas[-1]) ** 2
                for s in range(n_steps // self.config.model.num_classes):
                    images.append(torch.clamp(x_mod, 0.0, 1.0).to('cpu'))
                    noise = torch.randn_like(x_mod) * np.sqrt(step_size * 2)
                    grad = scorenet(x_mod)
                    x_mod = x_mod + step_size * grad + noise
                    #print("modulus of grad components: mean {}, max {}".format(grad.abs().mean(), grad.abs().max()))

            return images

    def test(self, test_iter=1000, step_lr=0.00002):
        states = torch.load(os.path.join(self.args.log, 'checkpoint.pth'), map_location=self.config.device)
        score = RefineNetDilated(self.config).to(self.config.device)
        score = torch.nn.DataParallel(score)

        score.load_state_dict(states[0])

        if not os.path.exists(self.args.image_folder):
            os.makedirs(self.args.image_folder)

        score.eval()
        
        grid_size = 10

        test_transform = transforms.Compose([
            transforms.Resize(self.config.data.image_size),
            transforms.ToTensor()
        ])

        if self.config.data.dataset == 'CIFAR10':
            test_dataset = CIFAR10(os.path.join(self.args.run, 'datasets', 'cifar10_test'), train=False, download=True,
                                   transform=test_transform)
        elif self.config.data.dataset == 'MNIST':
            test_dataset = MNIST(os.path.join(self.args.run, 'datasets', 'mnist_test'), train=False, download=True,
                                 transform=test_transform)
        elif self.config.data.dataset == 'FashionMNIST':
            test_dataset = FashionMNIST(os.path.join(self.args.run, 'datasets', 'fashionmnist_test'), train=False, download=True,
                                 transform=test_transform)

        elif self.config.data.dataset == 'CELEBA':
            test_dataset = CelebA(root=os.path.join(self.args.run, 'datasets', 'celeba'), split='test',
                                  transform=transforms.Compose([
                                      transforms.CenterCrop(140),
                                      transforms.Resize(self.config.data.image_size),
                                      transforms.ToTensor(),
                                  ]), download=True)
            
        test_loader = DataLoader(test_dataset, batch_size=grid_size**2, shuffle=True,
                                 num_workers=4, drop_last=True)

        for i, (X, y) in enumerate(test_loader):
            image_from_dataset = X.to(self.config.device)
            image_grid = make_grid(X, nrow=grid_size)
            save_image(image_grid, os.path.join(self.args.image_folder, 'image_original.png'))
            image_grid = make_grid(1-X, nrow=grid_size)
            save_image(image_grid, os.path.join(self.args.image_folder, 'image_flip.png'))
            break


        # initialize the sample as random noise
        samples = torch.rand(grid_size ** 2, self.config.data.channels, self.config.data.image_size, self.config.data.image_size, device=self.config.device)
        all_samples = self.Langevin_dynamics(samples, score, n_steps=test_iter, step_lr=step_lr)

        image_grid = make_grid(samples, nrow=grid_size)
        save_image(image_grid, os.path.join(self.args.image_folder, 'image_{}.png'.format(0)))

        for i, sample in enumerate(tqdm.tqdm(all_samples, total=len(all_samples), desc='saving images')):
            sample = sample.view(grid_size ** 2, self.config.data.channels, self.config.data.image_size,
                                    self.config.data.image_size)

            if self.config.data.logit_transform:
                sample = torch.sigmoid(sample)

            if i == len(all_samples) - 1:
                image_grid = make_grid(sample, nrow=grid_size)
                save_image(image_grid, os.path.join(self.args.image_folder, 'image_{}.png'.format(i)))


        # initialize the sample as the original image
        all_samples = self.Langevin_dynamics(image_from_dataset.clone().detach(), score, n_steps=test_iter, step_lr=step_lr)
        
        for i, sample in enumerate(tqdm.tqdm(all_samples, total=len(all_samples), desc='saving images')):
            sample = sample.view(grid_size ** 2, self.config.data.channels, self.config.data.image_size,
                                    self.config.data.image_size)

            if self.config.data.logit_transform:
                sample = torch.sigmoid(sample)

            if i == len(all_samples) - 1:
                image_grid = make_grid(sample, nrow=grid_size)
                save_image(image_grid, os.path.join(self.args.image_folder, 'image_original_{}.png'.format(i)), nrow=grid_size)


        # initialize the sample as flipped image
        if self.config.data.flip_p > 0.0:
            flip_sample =  1 - image_from_dataset
            all_samples = self.Langevin_dynamics(flip_sample.clone().detach(), score, n_steps=test_iter, step_lr=step_lr)

            for i, sample in enumerate(tqdm.tqdm(all_samples, total=len(all_samples), desc='saving images')):
                sample = sample.view(grid_size ** 2, self.config.data.channels, self.config.data.image_size,
                                        self.config.data.image_size)

                if self.config.data.logit_transform:
                    sample = torch.sigmoid(sample)

                if i == len(all_samples) - 1:
                    image_grid = make_grid(sample, nrow=grid_size)
                    save_image(image_grid, os.path.join(self.args.image_folder, 'image_flip_{}.png'.format(i)), nrow=grid_size)

