import os
import torch
import argparse
import yaml
import time
import multiprocessing as mp
from tabulate import tabulate
from tqdm import tqdm
from torch.utils.data import DataLoader
from pathlib import Path
from torch.utils.tensorboard import SummaryWriter
from torch.cuda.amp import GradScaler, autocast
from torch.nn.parallel import DistributedDataParallel as DDP
from torch.utils.data import DistributedSampler, RandomSampler
from torch import distributed as dist
from semseg.models import *
from semseg.datasets import *
from semseg.augmentations_mm import get_train_augmentation, get_val_augmentation
from semseg.losses import get_loss
from semseg.schedulers import get_scheduler
from semseg.optimizers import get_optimizer
from semseg.utils.utils import fix_seeds, setup_cudnn, cleanup_ddp, setup_ddp, get_logger, cal_flops, print_iou
from tools.val_mm_mmkd import evaluate
import torch.nn as nn
from semseg.models.segformer.seg_block_UMDt import Seg as Seg_s
from semseg.models.segformer.seg_block_select_UMD import Seg
from semseg.models.segformer.UMD import UMD
from semseg.models.segformer.CMD import UMD as CMD

from mmkd.Prototype import PrototypeSegmentation
from mmkd.PCUMD import PCUMD
from mmkd.IFVMD import IFVMD
from mmkd.PGUMD import PGUMD


def main(cfg, gpu, save_dir):
    start = time.time()
    best_mIoU = 0.0
    best_epoch = 0
    num_workers = 8
    device = torch.device(cfg['DEVICE'])
    train_cfg, eval_cfg = cfg['TRAIN'], cfg['EVAL']
    dataset_cfg, model_cfg = cfg['DATASET'], cfg['MODEL']
    loss_cfg, optim_cfg, sched_cfg = cfg['LOSS'], cfg['OPTIMIZER'], cfg['SCHEDULER']
    epochs, lr = train_cfg['EPOCHS'], optim_cfg['LR']
    resume_path = cfg['MODEL']['RESUME']
    gpus = int(os.environ['WORLD_SIZE'])

    traintransform = get_train_augmentation(train_cfg['IMAGE_SIZE'], seg_fill=dataset_cfg['IGNORE_LABEL'])
    valtransform = get_val_augmentation(eval_cfg['IMAGE_SIZE'])

    trainset = eval(dataset_cfg['NAME'])(dataset_cfg['ROOT'], 'train', traintransform, dataset_cfg['MODALS'])
    valset = eval(dataset_cfg['NAME'])(dataset_cfg['ROOT'], 'val', valtransform, dataset_cfg['MODALS'])
    class_names = trainset.CLASSES

    model = Seg(model_cfg['BACKBONE'], num_classes=trainset.n_classes, pretrained=True, modals=dataset_cfg['MODALS'])

    model_s = Seg_s(model_cfg['BACKBONE'], num_classes=trainset.n_classes, pretrained=True, modals=dataset_cfg['MODALS'])
    model_path = "TEACHER_PATH"  # Replace with the actual path to the teacher model
    # model_path = "/home/dell/PythonProjects/DELIVER/anyseg/MUSES_SegB0_FEL.pth"
    model_s.load_state_dict(torch.load(model_path, map_location=torch.device("cpu")), strict=False)
    model_s = model_s.to(device)
    model_s.eval()

    resume_checkpoint = None
    # model.load_state_dict(torch.load(model_path, map_location=torch.device("cpu")),strict=False)
    model = model.to(device)

    iters_per_epoch = len(trainset) // train_cfg['BATCH_SIZE'] // gpus
    loss_fn = get_loss(loss_cfg['NAME'], trainset.ignore_label, None)
    loss_kl = nn.KLDivLoss(size_average=None, reduce=None, reduction='mean', log_target=False)

    start_epoch = 0
    optimizer = get_optimizer(model, optim_cfg['NAME'], lr, optim_cfg['WEIGHT_DECAY'])
    scheduler = get_scheduler(sched_cfg['NAME'], optimizer, int((epochs + 1) * iters_per_epoch), sched_cfg['POWER'],
                              iters_per_epoch * sched_cfg['WARMUP'], sched_cfg['WARMUP_RATIO'])

    if train_cfg['DDP']:
        sampler = DistributedSampler(trainset, dist.get_world_size(), dist.get_rank(), shuffle=True)
        sampler_val = None
        model = DDP(model, device_ids=[gpu], output_device=0, find_unused_parameters=True)
    else:
        sampler = RandomSampler(trainset)
        sampler_val = None

    if resume_checkpoint:
        start_epoch = resume_checkpoint['epoch'] - 1
        optimizer.load_state_dict(resume_checkpoint['optimizer_state_dict'])
        scheduler.load_state_dict(resume_checkpoint['scheduler_state_dict'])
        loss = resume_checkpoint['loss']
        best_mIoU = resume_checkpoint['best_miou']

    trainloader = DataLoader(trainset, batch_size=train_cfg['BATCH_SIZE'], num_workers=num_workers, drop_last=True,
                             pin_memory=False, sampler=sampler)
    valloader = DataLoader(valset, batch_size=eval_cfg['BATCH_SIZE'], num_workers=num_workers, pin_memory=False,
                           sampler=sampler_val)

    scaler = GradScaler(enabled=train_cfg['AMP'])
    if (train_cfg['DDP'] and torch.distributed.get_rank() == 0) or (not train_cfg['DDP']):
        writer = SummaryWriter(str(save_dir))
        # logger.info('================== model complexity =====================')
        # cal_flops(model, dataset_cfg['MODALS'], logger)
        logger.info('================== model structure =====================')
        logger.info(model)
        logger.info('================== training config =====================')
        logger.info(cfg)

    embed_dim = [32, 64, 160, 256] if model_cfg['BACKBONE'] == "mit_b0" else [64, 128, 320, 512]
    prototype_all = [PrototypeSegmentation(num_classes=trainset.n_classes, feature_dim=embed_dim[0])]
    prototype_all.append(PrototypeSegmentation(num_classes=trainset.n_classes, feature_dim=embed_dim[1]))
    prototype_all.append(PrototypeSegmentation(num_classes=trainset.n_classes, feature_dim=embed_dim[2]))
    prototype_all.append(PrototypeSegmentation(num_classes=trainset.n_classes, feature_dim=embed_dim[3]))


    for epoch in range(start_epoch, epochs):
        torch.cuda.empty_cache()
        model.train()
        if train_cfg['DDP']: sampler.set_epoch(epoch)
        LOSS_COMBINE = "SPC"

        train_loss = 0.0
        all_loss = 0.0
        all_loss_s = 0.0
        # all_loss_umd = 0.0
        # all_loss_cmd = 0.0
        all_loss_pcumd = 0.0
        # all_loss_ifv = 0.0
        # all_loss_pmd = 0.0
        lr = scheduler.get_lr()
        lr = sum(lr) / len(lr)
        pbar = tqdm(enumerate(trainloader), total=iters_per_epoch,
                    desc=f"Epoch: [{epoch + 1}/{epochs}] Iter: [{0}/{iters_per_epoch}] LR: {lr:.8f} All Loss: {all_loss:.8f} Our Loss: {all_loss_pcumd:.8f}.")

        for iter, (sample, lbl) in pbar:
            optimizer.zero_grad(set_to_none=True)
            sample = [x.to(device) for x in sample]
            lbl = lbl.to(device)

            with autocast(enabled=train_cfg['AMP']):
                logits, index, ms_feat = model(sample)
                with torch.no_grad():
                    logits_s, ms_feat_s = model_s(sample)

                loss = loss_fn(logits, lbl)
                # print("lbl:",lbl.shape)

                # 把logits拆分
                logits = torch.log_softmax(logits, dim=1)
                logits_s = torch.softmax(logits_s, dim=1)
                loss_s = loss_kl(logits, logits_s).clamp(min=0) * 50

                # loss_umd = UMD(index, ms_feat, ms_feat_s) * 5
                # loss_cmd = CMD(index, ms_feat, ms_feat_s) * 10

                loss_pcumd = PCUMD(index, ms_feat, ms_feat_s, lbl, prototype_all) * 100
                # loss_ifv = IFVMD(index, ms_feat, ms_feat_s, lbl, trainset.n_classes) * 5
                # loss_pmd = PGUMD(index, ms_feat, ms_feat_s, lbl, prototype_all) * 100

                loss_all = loss_s + loss + loss_pcumd

            scaler.scale(loss_all).backward()
            scaler.step(optimizer)
            scaler.update()
            scheduler.step()
            torch.cuda.synchronize()

            lr = scheduler.get_lr()
            lr = sum(lr) / len(lr)
            if lr <= 1e-8:
                lr = 1e-8  # minimum of lr
            train_loss += loss.item()
            all_loss += loss_all.item()
            all_loss_s += loss_s.item()
            # all_loss_umd += loss_umd.item()
            # all_loss_cmd += loss_cmd.item()
            all_loss_pcumd += loss_pcumd.item()
            # all_loss_ifv += loss_ifv.item()
            # all_loss_pmd += loss_pmd.item()

            pbar.set_description(
                f"Epoch: [{epoch + 1}/{epochs}] Iter: [{iter + 1}/{iters_per_epoch}] LR: {lr:.8f} All Loss: {all_loss / (iter + 1):.8f} Our Loss:{all_loss_pcumd / (iter + 1): .8f}")

        train_loss /= iter + 1
        all_loss /= iter + 1
        all_loss_s /= iter + 1
        # all_loss_umd /= iter + 1
        # all_loss_cmd /= iter + 1
        # all_loss_pumd /= iter + 1
        all_loss_pcumd /= iter+ 1


        if (train_cfg['DDP'] and torch.distributed.get_rank() == 0) or (not train_cfg['DDP']):
            writer.add_scalar('train/loss', train_loss, epoch)
            writer.add_scalar('train/all_loss', all_loss, epoch)
            writer.add_scalar('train/loss_s', all_loss_s, epoch)
            # writer.add_scalar('train/loss_umd', all_loss_umd, epoch)
            # writer.add_scalar('train/loss_cmd', all_loss_cmd, epoch)
            writer.add_scalar('train/loss_pumd', all_loss_pcumd, epoch)
            # writer.add_scalar('train/loss_ifv', all_loss_ifv, epoch)
            # writer.add_scalar('train/loss_pmd', all_loss_ifv, epoch)
        torch.cuda.empty_cache()

        if ((epoch + 1) % train_cfg['EVAL_INTERVAL'] == 0 and (epoch + 1) > train_cfg['EVAL_START']) or (
                epoch + 1) == epochs:
            if (train_cfg['DDP'] and torch.distributed.get_rank() == 0) or (not train_cfg['DDP']):
                model.eval()
                acc, macc, _, _, ious, miou = evaluate(model, valloader, device)
                writer.add_scalar('val/mIoU', miou, epoch)

                if miou > best_mIoU:
                    prev_best_ckp = save_dir / f"{model_cfg['NAME']}_{model_cfg['BACKBONE']}_{dataset_cfg['NAME']}_epoch{best_epoch}_{best_mIoU}_{LOSS_COMBINE}_checkpoint.pth"
                    prev_best = save_dir / f"{model_cfg['NAME']}_{model_cfg['BACKBONE']}_{dataset_cfg['NAME']}_epoch{best_epoch}_{best_mIoU}_{LOSS_COMBINE}.pth"
                    if os.path.isfile(prev_best): os.remove(prev_best)
                    if os.path.isfile(prev_best_ckp): os.remove(prev_best_ckp)
                    best_mIoU = miou
                    best_epoch = epoch + 1
                    cur_best_ckp = save_dir / f"{model_cfg['NAME']}_{model_cfg['BACKBONE']}_{dataset_cfg['NAME']}_epoch{best_epoch}_{best_mIoU}_{LOSS_COMBINE}_checkpoint.pth"
                    cur_best = save_dir / f"{model_cfg['NAME']}_{model_cfg['BACKBONE']}_{dataset_cfg['NAME']}_epoch{best_epoch}_{best_mIoU}_{LOSS_COMBINE}.pth"
                    torch.save(model.module.state_dict() if train_cfg['DDP'] else model.state_dict(), cur_best)
                    # ---
                    torch.save({'epoch': best_epoch,
                                'model_state_dict': model.module.state_dict() if train_cfg[
                                    'DDP'] else model.state_dict(),
                                'optimizer_state_dict': optimizer.state_dict(),
                                'loss': train_loss,
                                'scheduler_state_dict': scheduler.state_dict(),
                                'best_miou': best_mIoU,
                                }, cur_best_ckp)
                    logger.info(print_iou(epoch, ious, miou, acc, macc, class_names))
                logger.info(f"Current epoch:{epoch} mIoU: {miou} Best mIoU: {best_mIoU}")

    if (train_cfg['DDP'] and torch.distributed.get_rank() == 0) or (not train_cfg['DDP']):
        writer.close()
    pbar.close()
    end = time.gmtime(time.time() - start)

    table = [
        ['Best mIoU', f"{best_mIoU:.2f}"],
        ['Total Training Time', time.strftime("%H:%M:%S", end)]
    ]
    logger.info(tabulate(table, numalign='right'))


if __name__ == '__main__':
    parser = argparse.ArgumentParser()
    parser.add_argument('--cfg', type=str, default='configs/deliver_rgbdel.yaml', help='Configuration file to use')
    args = parser.parse_args()

    with open(args.cfg) as f:
        cfg = yaml.load(f, Loader=yaml.SafeLoader)

    fix_seeds(3407)
    setup_cudnn()
    gpu = setup_ddp()
    modals = ''.join([m[0] for m in cfg['DATASET']['MODALS']])
    model = cfg['MODEL']['BACKBONE']
    exp_name = '_'.join([cfg['DATASET']['NAME'], model, modals])
    save_dir = Path(cfg['SAVE_DIR'], exp_name)
    if os.path.isfile(cfg['MODEL']['RESUME']):
        save_dir = Path(os.path.dirname(cfg['MODEL']['RESUME']))
    os.makedirs(save_dir, exist_ok=True)
    logger = get_logger(save_dir / 'train.log')
    main(cfg, gpu, save_dir)
    cleanup_ddp()
