import torch
import torch.nn as nn
import torch.optim as optim
import argparse
import os
from torch.optim.lr_scheduler import MultiStepLR
from tqdm import tqdm
import logging
import numpy as np
import math
import random
from torchvision.utils import save_image
from tqdm.contrib import tzip

from models.SiaStegNet import KeNet
from utils.terminal import MetricMonitor
from utils.logger import logger_info
from utils.dataset import get_train_loader, get_val_loader, get_test_loader
from utils.dirs import mkdirs
from utils.losses import ContrastiveLoss
import config as c


device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")

model_save_dir = os.path.join('checkpoints', 'SiaStegNet')
results_save_dir = os.path.join('results', 'SiaStegNet')
mkdirs(model_save_dir); mkdirs(results_save_dir)

logger_name = c.mode
logger_info(logger_name, log_path=os.path.join(results_save_dir, logger_name +'.log'))
logger = logging.getLogger(logger_name)
logger.info('#'*50)
logger.info('mode: {:s}'.format(c.mode))
logger.info('model: SiaStegNet')
logger.info('train data dir: {:s}'.format(c.train_data_dir))
logger.info('val data dir: {:s}'.format(c.val_data_dir))
logger.info('test data dir: {:s}'.format(c.test_data_dir))

margin = 1.00
alpha = 0.1

def preprocess_data(images, labels, random_crop):
    # images of shape: NxCxHxW
    if images.ndim == 5:  # 1xNxCxHxW
        images = images.squeeze(0)
        labels = labels.squeeze(0)
    h, w = images.shape[-2:]

    if random_crop:
        ch = random.randint(h * 3 // 4, h)  # h // 2      #256
        cw = random.randint(w * 3 // 4, w)  # square ch   #256

        h0 = random.randint(0, h - ch)  # 128
        w0 = random.randint(0, w - cw)  # 128
    else:
        ch, cw, h0, w0 = h, w, 0, 0

    # if args.model == 'kenet':  
    cw = cw & ~1
    inputs = [
        images[..., h0:h0 + ch, w0:w0 + cw // 2],
        images[..., h0:h0 + ch, w0 + cw // 2:w0 + cw]
    ]
    # elif args.model == 'sid':
    #     inputs = [images[..., h0:h0 + ch, w0:w0 + cw]]

    # if args.cuda:
    inputs = [x.cuda() for x in inputs]
    labels = labels.cuda()
    return inputs, labels


model = KeNet().to(device)

if c.mode == 'test':

    test_loader = get_test_loader(c.test_data_dir, c.test_batch_size)
    model.load_state_dict(torch.load(c.pre_trained_siastegnet_path))

    model.eval()

    with torch.no_grad():
        val_accuracy = []
        metric_monitor = MetricMonitor(float_precision=4)
        stream = tqdm(test_loader)

        for batch_idx, (inputs, labels) in enumerate(stream):

            inputs, labels = preprocess_data(inputs*255, labels, False)

            # inputs = inputs.to(device, dtype=torch.float)
            # labels = labels.to(device, dtype=torch.long)

            ################## forward ####################
            outputs, feats_0, feats_1 = model(*inputs)
            
            ################### loss ######################
            prediction = outputs.data.max(1)[1]

            accuracy = (
                prediction.eq(labels.data).sum() * 100.0 / (labels.size()[0])
            )
            val_accuracy.append(accuracy.item())
            metric_monitor.update("ACC", accuracy)
            stream.set_description(
                "Testing.  {metric_monitor}".format(metric_monitor=metric_monitor)
            )  

        val_acc_avg = np.mean(np.array(val_accuracy))

        logger.info('Testing, AVG_ACC: {:.4f}'.format(val_acc_avg))


else: # c.mode == 'train'

    train_loader = get_train_loader(c.train_data_dir, c.train_batch_size,)
    val_loader = get_val_loader(c.val_data_dir, c.val_batch_size)

    loss_fn1 = torch.nn.CrossEntropyLoss()
    loss_fn2 = ContrastiveLoss(margin=1.00)
    # optimizer = torch.optim.Adamax(model.parameters(), lr=1e-3, eps=1e-8, weight_decay=1e-4) # default opptimizer
    optimizer = torch.optim.Adam(model.parameters(), lr=c.lr, weight_decay=c.weight_decay)
    scheduler = torch.optim.lr_scheduler.StepLR(optimizer, step_size=c.weight_decay_step, gamma=c.gamma)

    for epoch in range(c.epochs):
        epoch += 1
        loss_history=[]
        train_accuracy = []
        ###############################################################
        #                            train                            # 
        ###############################################################
        model.train()
        metric_monitor = MetricMonitor(float_precision=4)
        stream = tqdm(train_loader)

        for batch_idx, train_batch in enumerate(stream):
            
            inputs = torch.cat((train_batch["cover"], train_batch["stego"]), 0)
            labels = torch.cat(
                (train_batch["label"][0], train_batch["label"][1]), 0
            )
            # inputs = inputs.to(device, dtype=torch.float)
            # labels = labels.to(device, dtype=torch.long)
            inputs, labels = preprocess_data(inputs*255, labels, False)

            ################## forward ####################
            outputs, feats_0, feats_1 = model(*inputs)

            ################### loss ######################
            loss = loss_fn1(outputs, labels) + \
                   alpha * loss_fn2(feats_0, feats_1, labels)

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

            loss_history.append(loss.item())
            prediction = outputs.data.max(1)[1]
            accuracy = (
                prediction.eq(labels.data).sum() * 100.0 / (labels.size()[0])
            )
            train_accuracy.append(accuracy.item())

            metric_monitor.update("Loss", loss.item())
            metric_monitor.update("ACC", accuracy)
            stream.set_description(
                "Epoch: {epoch}/{total_epochs}. Training.   {metric_monitor}".format(epoch=epoch, total_epochs=c.epochs, metric_monitor=metric_monitor)
            )

        scheduler.step()  
        train_losses_avg = np.mean(np.array(loss_history))
        train_acc_avg = np.mean(np.array(train_accuracy))  

        ###############################################################
        #                              val                            # 
        ###############################################################
        model.eval()
        if epoch % c.val_freq == 0:
            with torch.no_grad():
                loss_history=[]
                val_accuracy = []
                metric_monitor = MetricMonitor(float_precision=4)
                stream = tqdm(val_loader)

                for batch_idx, train_batch in enumerate(stream):
            
                    inputs = torch.cat((train_batch["cover"], train_batch["stego"]), 0)
                    labels = torch.cat(
                        (train_batch["label"][0], train_batch["label"][1]), 0
                    )

                    inputs, labels = preprocess_data(inputs*255, labels, False)

                    ################## forward ####################
                    outputs, feats_0, feats_1 = model(*inputs)

                    ################### loss ######################
                    loss = loss_fn1(outputs, labels) + \
                    alpha * loss_fn2(feats_0, feats_1, labels)


                    loss_history.append(loss.item())
                    prediction = outputs.data.max(1)[1]

                    accuracy = (
                        prediction.eq(labels.data).sum() * 100.0 / (labels.size()[0])
                    )
                    val_accuracy.append(accuracy.item())

                    metric_monitor.update("Loss", loss.item())
                    metric_monitor.update("ACC", accuracy)
                    
                    stream.set_description(
                        "Epoch: {epoch}/{total_epochs}. Validating.  {metric_monitor}".format(epoch=epoch, total_epochs=c.epochs, metric_monitor=metric_monitor)
                    )  

                val_losses_avg = np.mean(np.array(loss_history))
                val_acc_avg = np.mean(np.array(val_accuracy))
                lr = optimizer.state_dict()['param_groups'][0]['lr']
                logger.info('Epoch: {}/{}, Learning Rate: {:.5f} | Training, AVG_Loss {:.4f}, AVG_ACC: {:.4f} | Validating, AVG_Loss {:.4f}, AVG_ACC: {:.4f}'.format(epoch, c.epochs, lr, train_losses_avg, train_acc_avg, val_losses_avg, val_acc_avg))

        if epoch % c.save_freq == 0 and epoch >= c.strat_save_epoch:
            torch.save(model.state_dict(), os.path.join(model_save_dir, 'checkpoint_%.3i' % epoch + '.pt'))
            
        


