
import json
import os
import shutil
import sys
import time
import warnings

import numpy as np
import segmentation_models_pytorch as smp
import torch
import torch.nn as nn
import torch.nn.functional as F
from loguru import logger
from PIL import Image
from pytorch_toolbelt import losses as L
from segmentation_models_pytorch.losses import DiceLoss, SoftCrossEntropyLoss
from skimage.morphology import binary_dilation, disk
# import torchcontrib
from torch.nn import CrossEntropyLoss, DataParallel
from torch.nn.modules.loss import BCELoss, BCEWithLogitsLoss
from torch.optim import Adam, AdamW
from torch.utils.data import DataLoader
from torch.utils.tensorboard import SummaryWriter
from tqdm import tqdm

import loss
from datasets.change_detection import ChangeDetection
from loss import annealing_softmax_focalloss
from models.model_zoo import get_model
from models.pointrend import point_sample
from torchgeo1.models import FarSeg_CD, FarSeg_CD_Res
from torchgeo.models import ChangeStarFarSeg
from utils.image import add_image, visualize, visualize_CD, visualize_edge
from utils.metric import F1_score
from utils.options import Options
from utils.palette import color_map

warnings.filterwarnings("ignore")
# os.environ["CUDA_VISIBLE_DEVICES"] = "1"


# def layer_ce_loss(y_pred, y_true, l_weight=1.0):
#     label = y_true.detach().cpu().numpy()
#     # 统计类别数量
#     label_num = np.bincount(label.flatten())
#     weights = label_num.sum() / (label_num * 2)
#     if len(weights) == 1:
#         weights = np.append(weights, 0.)
#     weights = torch.FloatTensor(weights).cuda()
#     loss = CrossEntropyLoss(weight=weights)(y_pred, y_true)
#     return l_weight*loss
# 加载训练参数
args = Options().parse()
# 设置实验保存文件夹
time_str = time.strftime('%Y-%m-%d-%H:%M')
args.exp_name = '{}_{}_{}_{}'.format(
    args.model, args.backbone, args.exp_name, time_str)
# 日志记录路径
logs_dir = os.path.join('exp_result2', args.exp_name, 'logs')
# 日志记录文件
logger.add("{}/val_log.log".format(logs_dir))


class Trainer:
    def __init__(self, args):
        self.args = args
        # 加载数据
        trainset = ChangeDetection(
            root=args.data_root, mode="train")
        valset = ChangeDetection(root=args.data_root, mode="val")
        self.trainloader = DataLoader(trainset, batch_size=args.batch_size, shuffle=True,
                                      pin_memory=False, num_workers=4, drop_last=True)
        self.valloader = DataLoader(valset, batch_size=args.val_batch_size, shuffle=True,
                                    pin_memory=True, num_workers=4, drop_last=False)
        logger.info("train samples:{}".format(len(trainset)))
        logger.info("val samples:{}".format(len(valset)))
        # *******************************模型选择*************************************
        # self.model = get_model(args.model, args.backbone, args.pretrained,
        #                        len(trainset.CLASSES)-1, args.lightweight)
        if args.backbone == "hrnet_w30":
            self.model = FarSeg_CD(backbone=args.backbone, classes=2)
        elif args.backbone == "resnext50_32x4d":
            self.model = FarSeg_CD_Res(backbone=args.backbone, classes=2)
        self.model = DataParallel(self.model)
        # self.model =
        # *******************************模型选择*************************************
        # 加载预训练模型
        if args.pretrain_from:
            print("hello")
            self.model.load_state_dict(
                torch.load(args.pretrain_from), strict=True)

        if args.load_from:
            self.model.load_state_dict(torch.load(args.load_from), strict=True)

        # *****************************损失函数的选择********************************
        # 二分类
        weights = torch.FloatTensor([0.51890911, 13.72114046]).cuda()
        self.criterion_bin = CrossEntropyLoss(weight=weights)
        self.criterion = annealing_softmax_focalloss
        # self.criterion = DiceLoss("multiclass")
        # *****************************损失函数的选择********************************
        #
        # param = [name for name, param in self.model.named_parameters()]
        # print(param)
        # sys.exit()
        # 设置优化器
        self.optimizer = AdamW(self.model.parameters(),
                               lr=args.lr, weight_decay=args.weight_decay)
        # 调整学习率
        self.scheduler = torch.optim.lr_scheduler.CosineAnnealingWarmRestarts(
            self.optimizer,
            T_0=3,  # T_0就是初始restart的epoch数目
            T_mult=2,  # T_mult就是重启之后因子,即每个restart后，T_0 = T_0 + T_0 * T_mult
            eta_min=1e-6  # 最低学习率
        )
        # 模型并行训练
        # self.model = DataParallel(self.model).cuda()
        self.model = self.model.cuda()  # 单GPU训练
        # print(next(self.model.parameters()).device)
        self.iters = 0
        # 迭代总次数
        self.total_iters = len(self.trainloader) * args.epochs
        self.previous_best = 0.0

    def training(self):
        tbar = tqdm(self.trainloader)
        self.model.train()
        total_loss = 0.0
        loss_bin1 = 0.0
        for i, (img1, img2, mask_bin, id) in enumerate(tbar):
            img1, img2 = img1.cuda(), img2.cuda()
            mask_bin = mask_bin.cuda()
            # print(img1.shape)  # torch.Size([4, 3, 512, 512])
            # print(img2.shape)  # torch.Size([4, 3, 512, 512])
            # print(mask_bin.shape)  # torch.Size([4, 512, 512])
            # 模型输出
            out_bin = self.model(img1, img2)
            loss_bin1 = self.criterion_bin(out_bin.squeeze(), mask_bin.long())
            loss_bin2 = self.criterion(
                out_bin.squeeze(), mask_bin.long(), self.iters, self.total_iters)
            # loss_bin2 = self.criterion(out_bin.squeeze(), mask_bin.long())
            loss = loss_bin1+loss_bin2
            loss = loss.mean()
            total_loss += loss.item()
            self.optimizer.zero_grad()
            loss.backward()
            self.optimizer.step()
            self.iters += 1

            # lr = self.args.lr * (1 - self.iters / self.total_iters) ** 0.9
            # self.optimizer.param_groups[0]["lr"] = lr
            # self.optimizer.param_groups[1]["lr"] = lr * 10.0

            tbar.set_description("Train Loss: %.3f" % (total_loss / (i + 1)))
        train_writer.add_scalar("loss", total_loss, epoch)
        logger.info("Train Loss:{}".format(total_loss / (i + 1)))

    def validation(self):
        total_loss = 0.0
        tbar = tqdm(self.valloader)
        self.model.eval()
        metric_change = F1_score(num_classes=2)
        with torch.no_grad():
            for i, (img1, img2,  mask_bin, id) in enumerate(tbar):
                img1, img2 = img1.cuda(), img2.cuda()
                mask_bin = mask_bin.cuda()
                out_bin = self.model(img1, img2)
                loss_bin1 = self.criterion_bin(
                    out_bin.squeeze(), mask_bin.long())
                loss_bin2 = self.criterion(
                    out_bin.squeeze(), mask_bin.long(), self.iters, self.total_iters)
                # loss_bin2 = self.criterion(out_bin.squeeze(), mask_bin.long())
                loss = loss_bin1+loss_bin2
                loss = loss.mean()
                total_loss += loss.item()
                tbar.set_description("Val Loss: %.3f" % (total_loss / (i + 1)))
                out_bin = torch.argmax(out_bin, dim=1).cpu().numpy()
                metric_change.add_batch(out_bin, mask_bin.cpu().numpy())
            f1_change = metric_change.evaluate()
            Score_bin = 0.6*f1_change
            logger.info(
                "f1_change_Score:{}".format(f1_change))
            logger.info("Score_bin:{}".format(Score_bin))
            val_writer.add_scalar('Score', Score_bin, epoch)
        logger.info("Val Loss:{}".format(total_loss / (i + 1)))
        if self.args.load_from:
            exit(0)

        Score_bin *= 100.0
        save_model_path = checkpoint_dir
        if not os.path.exists(save_model_path):
            os.makedirs(save_model_path)
        if Score_bin >= self.previous_best:
            if self.previous_best != 0:
                model_path = "%s/%s_%s_bin_%.2f.pth" % \
                             (save_model_path, self.args.model,
                              self.args.backbone, self.previous_best)
                if os.path.exists(model_path):
                    os.remove(model_path)
            torch.save(self.model.state_dict(), "%s/%s_%s_bin_%.2f.pth" %
                       (save_model_path, self.args.model, self.args.backbone, Score_bin), _use_new_zipfile_serialization=False)
            self.previous_best = Score_bin
            # # 可视化
            for i in range(1):
                visualize_CD([img1[i].cpu().numpy(), img2[i].cpu().numpy()],
                             [mask_bin[i].cpu().numpy(), out_bin[i]],
                             imgs_dir, str(epoch)+"_"+id[i]+("_bin_%.2f" % (Score_bin)) + ".png")

            self.previous_best = Score_bin
            logger.info("epoch:{}  Score:{}".format(epoch, Score_bin))


if __name__ == "__main__":
    logger.info(args.exp_name)
    # -------------------- set directory for saving files -------------------------
    # 模型文件路径
    checkpoint_dir = os.path.join(
        'exp_result2', args.exp_name, 'checkpoints')

    # 生成图像路径
    imgs_dir = os.path.join('exp_result2', args.exp_name, 'images')
    # 创建多级目录
    if not os.path.isdir(checkpoint_dir):
        os.makedirs(checkpoint_dir)
    if not os.path.isdir(logs_dir):
        os.makedirs(logs_dir)
    if not os.path.isdir(imgs_dir):
        os.makedirs(imgs_dir)

    # 保存训练参数
    with open('{}/model_cfg.json'.format(logs_dir), 'w') as jfile:
        json.dump(vars(args), jfile, indent=2)

    # 设置可视化训练过程
    train_writer = SummaryWriter(os.path.join(
        logs_dir, 'runs', 'training'))
    val_writer = SummaryWriter(os.path.join(
        logs_dir, 'runs', 'val'))
    trainer = Trainer(args)

    if args.load_from:
        trainer.validation()

    for epoch in range(args.epochs):
        logger.info("\n")
        logger.info("\n==> Epoches %i/%i, learning rate = %.7f\t\t\t\t previous best = %.2f" %
                    (epoch, args.epochs, trainer.optimizer.param_groups[0]["lr"], trainer.previous_best))
        # 训练
        trainer.training()
        # 验证
        trainer.validation()
        # 更新学习率
        trainer.scheduler.step()
