import os
import math
from decimal import Decimal
from data import common
from torchvision import transforms
import utility

import torch
import torch.nn.utils as utils
from tqdm import tqdm
import PIL

class Trainer():
    def __init__(self, args, loader, my_model, my_loss, ckp):
        self.args = args
        self.scale = args.scale

        self.ckp = ckp
        self.loader_train = loader.loader_train
        self.loader_test = loader.loader_test
        self.model = my_model
        self.loss = my_loss
        self.optimizer = utility.make_optimizer(args, self.model)

        if self.args.load != '':
            self.optimizer.load(ckp.dir, epoch=len(ckp.log))

        self.error_last = 1e8

    def train(self):
        self.loss.step()
        epoch = self.optimizer.get_last_epoch() + 1
        lr = self.optimizer.get_lr()

        self.ckp.write_log(
            '[Epoch {}]\tLearning rate: {:.2e}'.format(epoch, Decimal(lr))
        )
        self.loss.start_log()
        self.model.train()

        timer_data, timer_model = utility.timer(), utility.timer()
        # TEMP
        self.loader_train.dataset.set_scale(0)
        for batch, (lr, hr, _,) in enumerate(self.loader_train):
            lr, hr = self.prepare(lr, hr)
            timer_data.hold()
            timer_model.tic()

            self.optimizer.zero_grad()
            sr = self.model(lr, 0)
            loss = self.loss(sr, hr)
            loss.backward()
            if self.args.gclip > 0:
                utils.clip_grad_value_(
                    self.model.parameters(),
                    self.args.gclip
                )
            self.optimizer.step()

            timer_model.hold()

            if (batch + 1) % self.args.print_every == 0:
                self.ckp.write_log('[{}/{}]\t{}\t{:.1f}+{:.1f}s'.format(
                    (batch + 1) * self.args.batch_size,
                    len(self.loader_train.dataset),
                    self.loss.display_loss(batch),
                    timer_model.release(),
                    timer_data.release()))

            timer_data.tic()

        self.loss.end_log(len(self.loader_train))
        self.error_last = self.loss.log[-1, -1]
        self.optimizer.schedule()

    def test(self):
        torch.set_grad_enabled(False)

        epoch = self.optimizer.get_last_epoch() + 1
        self.ckp.write_log('\nEvaluation:')
        self.ckp.add_log(
            torch.zeros(1, len(self.loader_test), len(self.scale))
        )
        self.model.eval()

        timer_test = utility.timer()
        if self.args.save_results: self.ckp.begin_background()
        for idx_data, d in enumerate(self.loader_test):
            for idx_scale, scale in enumerate(self.scale):
                d.dataset.set_scale(idx_scale)
                if self.args.data_test[0] != 'DemoSplit':
                    for lr, hr, filename in tqdm(d, ncols=80):
                        lr, hr = self.prepare(lr, hr)
                        sr = self.model(lr, idx_scale)
                        sr = utility.quantize(sr, self.args.rgb_range)

                        save_list = [sr]
                        self.ckp.log[-1, idx_data, idx_scale] += utility.calc_psnr(
                            sr, hr, scale, self.args.rgb_range, dataset=d
                        )
                        if self.args.save_gt:
                            save_list.extend([lr, hr])

                        if self.args.save_results:
                            self.ckp.save_results(d, filename[0], save_list, scale)

                else:
                    for lr, patch_num, filename in tqdm(d, ncols=80):
                        step = int(patch_num[0]*patch_num[1])/self.args.batch_size
                        step = math.ceil(step)
                        lr = torch.squeeze(lr, 0)
                        for i in range(step):
                            index = i*self.args.batch_size
                            if(i+self.args.batch_size<patch_num[0]*patch_num[1]):
                                lr_batch, hr = self.prepare(lr[index:index+self.args.batch_size, :, :, :], torch.zeros((1)))
                            else:
                                lr_batch, hr = self.prepare(lr[index:patch_num[0]*patch_num[1], :, :, :], torch.zeros((1)))
                            sr = self.model(lr_batch, idx_scale)
                            if(i == 0):
                                whole_sr = sr
                            else:
                                whole_sr = torch.cat((whole_sr, sr), 0)

                        target_img = torch.zeros((patch_num[0]*self.args.patch_size, patch_num[1]*self.args.patch_size, self.args.n_colors))
                        for i in range(patch_num[0]):
                            for j in range(patch_num[1]):
                                a = whole_sr[i*patch_num[1]+j, :, :, :]
                                target_img[i * self.args.patch_size:(i + 1) * self.args.patch_size, j * self.args.patch_size:(j + 1) * self.args.patch_size, :] \
                                    = torch.squeeze(a.permute(2, 3, 1, 0), 3)

                        target_img = target_img.permute(2, 0, 1)
                        target_img = utility.quantize(target_img, self.args.rgb_range)
                        target_img = target_img.cpu().clone()
                        unloader = transforms.ToPILImage()
                        target_img = PIL.ImageOps.invert(unloader(target_img))
                        target_img.save(os.path.join(self.args.dir_demo, 'result_'+filename[0]+'.png'))

                        save_list = []

                        if self.args.save_results:
                            self.ckp.save_results(d, filename[0], save_list, scale)


                self.ckp.log[-1, idx_data, idx_scale] /= len(d)
                best = self.ckp.log.max(0)
                self.ckp.write_log(
                    '[{} x{}]\tPSNR: {:.3f} (Best: {:.3f} @epoch {})'.format(
                        d.dataset.name,
                        scale,
                        self.ckp.log[-1, idx_data, idx_scale],
                        best[0][idx_data, idx_scale],
                        best[1][idx_data, idx_scale] + 1
                    )
                )

        self.ckp.write_log('Forward: {:.2f}s\n'.format(timer_test.toc()))
        self.ckp.write_log('Saving...')

        if self.args.save_results:
            self.ckp.end_background()

        if not self.args.test_only:
            self.ckp.save(self, epoch, is_best=(best[1][0, 0] + 1 == epoch))

        self.ckp.write_log(
            'Total: {:.2f}s\n'.format(timer_test.toc()), refresh=True
        )

        torch.set_grad_enabled(True)

    def prepare(self, *args):
        device = torch.device('cpu' if self.args.cpu else 'cuda')
        def _prepare(tensor):
            if self.args.precision == 'half': tensor = tensor.half()
            return tensor.to(device)

        return [_prepare(a) for a in args]

    def terminate(self):
        if self.args.test_only:
            self.test()
            return True
        else:
            epoch = self.optimizer.get_last_epoch() + 1
            return epoch >= self.args.epochs

