from datasets.dataset import MakeTrainSet, Kodak24
from utlz import Quantization
from invertible_net import Inveritible_Decolorization
from datasets.utls import str2bool
import torch.utils.data as data
from criteria import ConsistencyLoss
from torch.optim import Adam
import os, sys, time, math
import torch, cv2
import logging
import argparse
import numpy as np
import torch.nn as nn
from termcolor import colored
from mmcv.utils import get_logger
from skimage.metrics import peak_signal_noise_ratio as compare_psnr


class Logger(object):
    def __init__(self, filename='default.log', stream=sys.stdout):
        self.terminal = stream
        self.log = open(filename, 'a')

    def write(self, message):
        self.terminal.write(message)
        self.log.write(message)

    def flush(self):
        pass


def main(args):
    beginner = args.beginner
    stride = args.stride
    device = args.device
    batch_size = args.batch_size
    epoch = args.Epoch
    lr = args.lr
    weight_decay = args.weight_decay
    model_path = "./models/"

    timestamp = time.strftime('%Y%m%d_%H%M%S', time.localtime())
    log_file = os.path.join(model_path, f'{timestamp}.log')
    logger = get_logger(name='IDN', log_file=log_file, log_level=logging.INFO)
    logger.info(f"batch size {batch_size}")

    TrainSet = MakeTrainSet()
    TrainLoader = data.DataLoader(TrainSet, batch_size=batch_size, num_workers=4, shuffle=True, pin_memory=True)
    ValidSet = Kodak24()
    ValidLoader = data.DataLoader(ValidSet, batch_size=8, num_workers=8, shuffle=False, pin_memory=True)

    logger.info("# Training Samples: " + str(TrainSet.__len__()) + "; Valid Samples: " + str(ValidSet.__len__()))

    net = Inveritible_Decolorization()
    optimizer = Adam(net.parameters(), weight_decay=weight_decay, betas=(0.5, 0.999), lr=lr)

    if args.load_checkpoint:
        logger.info("loading checkpoint...")
        net.load_state_dict(torch.load("./models/****.pth"))

    net = torch.nn.DataParallel(net, device_ids=[0, 1])
    net.cuda(device=device)
    quantize = Quantization()
    loss_cons = ConsistencyLoss(device="cuda:0", img_shape=(batch_size, 1, 128, 128), c_weight=args.c_weight)
    loss_dist = nn.L1Loss(reduction="sum")

    tmp = filter(lambda x: x.requires_grad, net.parameters())
    num = sum(map(lambda x: np.prod(x.shape), tmp))
    print('Total trainable tensors:', str(num))

    BestPSNR = 0.
    
    # use warmup to stablize the training, or it may not converage

    for i in range(beginner, epoch):
        loss_record = 0
        trainSamples = 0

        if i != beginner:
            TrainSet = MakeTrainSet()
            TrainLoader = data.DataLoader(TrainSet, batch_size=batch_size, num_workers=4, shuffle=True, pin_memory=True)

        for batch_idx, (tensor_g, tensor_c) in enumerate(TrainLoader):
            n, _, h, w = tensor_c.size()

            tensor_c = tensor_c.cuda(device=device)
            tensor_g = tensor_g.cuda(device=device)

            tensor_g.requires_grad = False
            tensor_c.requires_grad = False

            trainSamples += n

            net.train(True)

            optimizer.zero_grad()

            tensor_x = net(x=[tensor_c], rev=False)[0]
            tensor_prg = tensor_x[:, [0], :, :]
            tensor_z = tensor_x[:, 1:, :, :]

            tensor_z_ = torch.randn_like(tensor_z).to(tensor_z.device)
            tensor_g_ = quantize(tensor_prg)
            tensor_x_ = torch.cat((tensor_g_, tensor_z_), dim=1)
            tensor_y = net(x=[tensor_x_], rev=True)[0]

            loss_invt = loss_cons(gray_img=tensor_prg, ref_img=tensor_g, original_img=tensor_c,
                                  restored_img=tensor_y, loss_stage=1, s_wegiht=args.s_weight) / n
            loss_self = (tensor_z**2).sum() / (128**2 * 2 * n)
            loss_gray = loss_dist(tensor_prg, tensor_g) / n
            loss = loss_invt + args.r_weight * loss_self / 2 + args.g_weight * loss_gray

            loss.backward()
            torch.nn.utils.clip_grad_norm_(net.parameters(), 10)
            optimizer.step()

            loss_record += loss.detach().cpu() * n
            if batch_idx % 10 == 9 or batch_idx == 0:
                print('#batch: %4d; learning rate: %.4e; loss_invt: %3.4f; loss_self: %3.4f; loss_gray: %3.4f.'
                      % (batch_idx + 1, optimizer.param_groups[0]['lr'],
                         float(loss_invt), float(loss_self), float(loss_gray)))
        avg_loss = loss_record / trainSamples
        logger.info('Train: Epoch = %3d | Ave Loss = %.4f | Train Samples = %3d.' % (i + 1, avg_loss, trainSamples))

        if i % stride == (stride - 1):
            torch.save(net.module.state_dict(), os.path.join(model_path, 'net_epoch_' + str(i + 1).zfill(3) + '.pth'))

        net.eval()
        validSamples = 0

        ColorPSNR = 0
        GrayPSNR = 0

        for batch_idx, (tensor_g, tensor_c, h, w) in enumerate(ValidLoader):
            n, _, _, _ = tensor_c.size()

            tensor_g, tensor_c = tensor_g.cuda(device=device), tensor_c.cuda(device=device)
            tensor_g.requires_grad = False
            tensor_c.requires_grad = False

            validSamples += tensor_c.size(0)
            with torch.no_grad():
                tensor_pg = net(x=[tensor_c], rev=False)[0][:, [0], :, :]
                tensor_pg = quantize(tensor_pg)
                z = torch.randn(size=(n, 2, 512, 512)).to(tensor_pg.device) * 1 / 255
                tensor_x = torch.cat((tensor_pg, z), dim=1)
                tensor_pc = net(x=[tensor_x], rev=True)[0]
            tensor_pg = tensor_pg.squeeze().cpu().detach().numpy()
            tensor_c = tensor_c.squeeze().cpu().detach().numpy()
            tensor_pc = tensor_pc.squeeze().cpu().detach().numpy()
            tensor_g = tensor_g.squeeze().cpu().detach().numpy()
            for ii in range(n):
                img_colorized = tensor_pc[ii].transpose((1, 2, 0)).clip(0, 1)[:h[ii], :w[ii], :]
                img_ground_th = tensor_c[ii].transpose((1, 2, 0)).clip(0, 1)[:h[ii], :w[ii], :]
                psnr_c = round(compare_psnr(img_ground_th, img_colorized, data_range=1), 4)

                img_grayscale = tensor_pg[ii].clip(0, 1)[:h[ii], :w[ii]]
                img_grayed_th = tensor_g[ii].clip(0, 1)[:h[ii], :w[ii]]
                psnr_g = round(compare_psnr(img_grayed_th, img_grayscale, data_range=1), 4)

                print("#sample:" + str(batch_idx) + "_" + str(ii) +
                      ' Colorized => PSNR: %.4f' % (psnr_c), end=" | ")
                print('Grayscale => PSNR: %.4f' % (psnr_g))

                ColorPSNR += psnr_c
                GrayPSNR += psnr_g

        ColorPSNR = ColorPSNR / validSamples
        GrayPSNR = GrayPSNR / validSamples

        ValidPSNR = (ColorPSNR + GrayPSNR) / 2
        if BestPSNR < ValidPSNR:
            BestPSNR = ValidPSNR
            torch.save(net.module.state_dict(), os.path.join(model_path, 'ColorFlow.pth'))
        logger.info(
            'Valid Samples = %3d | ' % validSamples
            + colored('Color PSNR = %.4f', 'red') % ColorPSNR + " | "
            + colored('Best PSNR = %.4f', 'green') % BestPSNR + " | "
            + 'Gray PSNR = %.4f' % GrayPSNR
        )


if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    # For Dataset and Record
    parser.add_argument("--stride", type=int, default=10, help='the stride for saving models')
    parser.add_argument("--paired", type=str2bool, default=True, help="paired or unpaired")
    parser.add_argument("--root", type=str, default="F:/datasets/Colorization/VCIP2020_Colorization_Challenge/",
                        help="data root")
    # For Training
    parser.add_argument("--load_checkpoint", type=str2bool, default=False)
    parser.add_argument("--batch_size", type=int, default=2)
    parser.add_argument("--beginner", type=int, default=0)
    parser.add_argument('--Epoch', type=int, default=1)
    parser.add_argument("--c_weight", type=float)
    parser.add_argument("--s_weight", type=float)
    parser.add_argument("--r_weight", type=float)
    parser.add_argument("--lr", type=float, default=5e-5)
    parser.add_argument("--weight_decay", type=float, default=5e-4)
    parser.add_argument('--device', type=int, default=0)
    args = parser.parse_args()

    main(args)

from datasets.dataset import MakeTrainSet, Kodak24
from utlz import Quantization
from invertible_net import Inveritible_Decolorization
from datasets.utls import str2bool
import torch.utils.data as data
from criteria import ConsistencyLoss
from torch.optim import Adam
import os, sys, time, math
import torch, cv2
import logging
import argparse
import numpy as np
import torch.nn as nn
from termcolor import colored
from mmcv.utils import get_logger
from skimage.metrics import peak_signal_noise_ratio as compare_psnr


class Logger(object):
    def __init__(self, filename='default.log', stream=sys.stdout):
        self.terminal = stream
        self.log = open(filename, 'a')

    def write(self, message):
        self.terminal.write(message)
        self.log.write(message)

    def flush(self):
        pass


def main(args):
    beginner = args.beginner
    stride = args.stride
    device = args.device
    batch_size = args.batch_size
    epoch = args.Epoch
    lr = args.lr
    weight_decay = args.weight_decay
    model_path = "./models/"

    timestamp = time.strftime('%Y%m%d_%H%M%S', time.localtime())
    log_file = os.path.join(model_path, f'{timestamp}.log')
    logger = get_logger(name='IDN', log_file=log_file, log_level=logging.INFO)
    logger.info(f"batch size {batch_size}")

    TrainSet = MakeTrainSet()
    TrainLoader = data.DataLoader(TrainSet, batch_size=batch_size, num_workers=4, shuffle=True, pin_memory=True)
    ValidSet = Kodak24()
    ValidLoader = data.DataLoader(ValidSet, batch_size=8, num_workers=8, shuffle=False, pin_memory=True)

    logger.info("# Training Samples: " + str(TrainSet.__len__()) + "; Valid Samples: " + str(ValidSet.__len__()))

    net = Inveritible_Decolorization()
    optimizer = Adam(net.parameters(), weight_decay=weight_decay, betas=(0.5, 0.999), lr=lr)

    if args.load_checkpoint:
        logger.info("loading checkpoint...")
        net.load_state_dict(torch.load("./models/****.pth"))

    net = torch.nn.DataParallel(net, device_ids=[0, 1])
    net.cuda(device=device)
    quantize = Quantization()
    loss_cons = ConsistencyLoss(device="cuda:0", img_shape=(batch_size, 1, 128, 128), c_weight=args.c_weight)
    loss_dist = nn.L1Loss(reduction="sum")

    tmp = filter(lambda x: x.requires_grad, net.parameters())
    num = sum(map(lambda x: np.prod(x.shape), tmp))
    print('Total trainable tensors:', str(num))

    BestPSNR = 0.
    
    # use warmup to stablize the training, or it may not converage

    for i in range(beginner, epoch):
        loss_record = 0
        trainSamples = 0

        if i != beginner:
            TrainSet = MakeTrainSet()
            TrainLoader = data.DataLoader(TrainSet, batch_size=batch_size, num_workers=4, shuffle=True, pin_memory=True)

        for batch_idx, (tensor_g, tensor_c) in enumerate(TrainLoader):
            n, _, h, w = tensor_c.size()

            tensor_c = tensor_c.cuda(device=device)
            tensor_g = tensor_g.cuda(device=device)

            tensor_g.requires_grad = False
            tensor_c.requires_grad = False

            trainSamples += n

            net.train(True)

            optimizer.zero_grad()

            tensor_x = net(x=[tensor_c], rev=False)[0]
            tensor_prg = tensor_x[:, [0], :, :]
            tensor_z = tensor_x[:, 1:, :, :]

            tensor_z_ = torch.randn_like(tensor_z).to(tensor_z.device)
            tensor_g_ = quantize(tensor_prg)
            tensor_x_ = torch.cat((tensor_g_, tensor_z_), dim=1)
            tensor_y = net(x=[tensor_x_], rev=True)[0]

            loss_invt = loss_cons(gray_img=tensor_prg, ref_img=tensor_g, original_img=tensor_c,
                                  restored_img=tensor_y, loss_stage=1, s_wegiht=args.s_weight) / n
            loss_self = (tensor_z**2).sum() / (128**2 * 2 * n)
            loss_gray = loss_dist(tensor_prg, tensor_g) / n
            loss = loss_invt + args.r_weight * loss_self / 2 + args.g_weight * loss_gray

            loss.backward()
            torch.nn.utils.clip_grad_norm_(net.parameters(), 10)
            optimizer.step()

            loss_record += loss.detach().cpu() * n
            if batch_idx % 10 == 9 or batch_idx == 0:
                print('#batch: %4d; learning rate: %.4e; loss_invt: %3.4f; loss_self: %3.4f; loss_gray: %3.4f.'
                      % (batch_idx + 1, optimizer.param_groups[0]['lr'],
                         float(loss_invt), float(loss_self), float(loss_gray)))
        avg_loss = loss_record / trainSamples
        logger.info('Train: Epoch = %3d | Ave Loss = %.4f | Train Samples = %3d.' % (i + 1, avg_loss, trainSamples))

        if i % stride == (stride - 1):
            torch.save(net.module.state_dict(), os.path.join(model_path, 'net_epoch_' + str(i + 1).zfill(3) + '.pth'))

        net.eval()
        validSamples = 0

        ColorPSNR = 0
        GrayPSNR = 0

        for batch_idx, (tensor_g, tensor_c, h, w) in enumerate(ValidLoader):
            n, _, _, _ = tensor_c.size()

            tensor_g, tensor_c = tensor_g.cuda(device=device), tensor_c.cuda(device=device)
            tensor_g.requires_grad = False
            tensor_c.requires_grad = False

            validSamples += tensor_c.size(0)
            with torch.no_grad():
                tensor_pg = net(x=[tensor_c], rev=False)[0][:, [0], :, :]
                tensor_pg = quantize(tensor_pg)
                z = torch.randn(size=(n, 2, 512, 512)).to(tensor_pg.device) * 1 / 255
                tensor_x = torch.cat((tensor_pg, z), dim=1)
                tensor_pc = net(x=[tensor_x], rev=True)[0]
            tensor_pg = tensor_pg.squeeze().cpu().detach().numpy()
            tensor_c = tensor_c.squeeze().cpu().detach().numpy()
            tensor_pc = tensor_pc.squeeze().cpu().detach().numpy()
            tensor_g = tensor_g.squeeze().cpu().detach().numpy()
            for ii in range(n):
                img_colorized = tensor_pc[ii].transpose((1, 2, 0)).clip(0, 1)[:h[ii], :w[ii], :]
                img_ground_th = tensor_c[ii].transpose((1, 2, 0)).clip(0, 1)[:h[ii], :w[ii], :]
                psnr_c = round(compare_psnr(img_ground_th, img_colorized, data_range=1), 4)

                img_grayscale = tensor_pg[ii].clip(0, 1)[:h[ii], :w[ii]]
                img_grayed_th = tensor_g[ii].clip(0, 1)[:h[ii], :w[ii]]
                psnr_g = round(compare_psnr(img_grayed_th, img_grayscale, data_range=1), 4)

                print("#sample:" + str(batch_idx) + "_" + str(ii) +
                      ' Colorized => PSNR: %.4f' % (psnr_c), end=" | ")
                print('Grayscale => PSNR: %.4f' % (psnr_g))

                ColorPSNR += psnr_c
                GrayPSNR += psnr_g

        ColorPSNR = ColorPSNR / validSamples
        GrayPSNR = GrayPSNR / validSamples

        ValidPSNR = (ColorPSNR + GrayPSNR) / 2
        if BestPSNR < ValidPSNR:
            BestPSNR = ValidPSNR
            torch.save(net.module.state_dict(), os.path.join(model_path, 'ColorFlow.pth'))
        logger.info(
            'Valid Samples = %3d | ' % validSamples
            + colored('Color PSNR = %.4f', 'red') % ColorPSNR + " | "
            + colored('Best PSNR = %.4f', 'green') % BestPSNR + " | "
            + 'Gray PSNR = %.4f' % GrayPSNR
        )


if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    # For Dataset and Record
    parser.add_argument("--stride", type=int, default=10, help='the stride for saving models')
    parser.add_argument("--paired", type=str2bool, default=True, help="paired or unpaired")
    parser.add_argument("--root", type=str, default="F:/datasets/Colorization/VCIP2020_Colorization_Challenge/",
                        help="data root")
    # For Training
    parser.add_argument("--load_checkpoint", type=str2bool, default=False)
    parser.add_argument("--batch_size", type=int, default=2)
    parser.add_argument("--beginner", type=int, default=0)
    parser.add_argument('--Epoch', type=int, default=1)
    parser.add_argument("--c_weight", type=float)
    parser.add_argument("--s_weight", type=float)
    parser.add_argument("--r_weight", type=float)
    parser.add_argument("--lr", type=float, default=5e-5)
    parser.add_argument("--weight_decay", type=float, default=5e-4)
    parser.add_argument('--device', type=int, default=0)
    args = parser.parse_args()

    main(args)

# use warmup to stablize the training, or it may not converage # For Dataset and Record # For Training