import torch
from torch import nn
from torch.optim import lr_scheduler
from monai.utils import set_determinism
from torch import optim
from monai.data import CacheDataset, DataLoader, ThreadDataLoader, PersistentDataset
from torch.nn.modules.loss import _Loss
from monai.utils import LossReduction
from monai.losses import DiceLoss
from monai.transforms import LoadImage
from monai.optimizers.lr_scheduler import WarmupCosineSchedule


class DiceBceMultilabelLoss(_Loss):
    def __init__(
        self,
        w_dice = 0.5,
        w_bce = 0.5,
        reduction = LossReduction.MEAN,
    ):
        super().__init__(reduction=LossReduction(reduction).value)
        self.w_dice = w_dice
        self.w_bce = w_bce
        self.dice_loss = DiceLoss(sigmoid=True, smooth_nr=0.01, smooth_dr=0.01, include_background=True, batch=True, squared_pred=True)
        self.bce_loss = nn.BCEWithLogitsLoss()
    def forward(self, pred, label):
        
        loss = self.dice_loss(pred, label) * self.w_dice + self.bce_loss(pred, label) * self.w_bce
        return loss


def get_train_dataloader(train_dataset, cfg):

    if cfg.gpu_cache:
        train_dataloader = ThreadDataLoader(
            train_dataset,
            shuffle=True,
            batch_size=cfg.batch_size,
            num_workers=0,
            drop_last=True,
        )
        return train_dataloader 

    train_dataloader = DataLoader(
        train_dataset,
        shuffle=True,
        batch_size=cfg.batch_size,
        num_workers=cfg.num_workers,
        drop_last=True,
    )
    return train_dataloader


def get_val_dataloader(val_dataset, cfg):
    if cfg.val_gpu_cache:
        val_dataloader = ThreadDataLoader(
            val_dataset,
            batch_size=cfg.val_batch_size,
            num_workers=0,
        )
        return val_dataloader

    val_dataloader = DataLoader(
        val_dataset,
        batch_size=cfg.val_batch_size,
        num_workers=cfg.num_workers,
    )
    return val_dataloader

def get_train_dataset(cfg):
    train_ds = CacheDataset(
        data=cfg.data_json["train"],
        transform=cfg.train_transforms,
        cache_rate=cfg.train_cache_rate,
        num_workers=cfg.num_workers,
        copy_cache=False,
    )
    # train_ds = PersistentDataset(
    #     data=cfg.data_json["train"],
    #     transform=cfg.train_transforms,
    #     cache_dir="cache_data",
    # )
    return train_ds

def get_val_dataset(cfg):
    val_ds = CacheDataset(
        data=cfg.data_json["val"],
        transform=cfg.val_transforms,
        cache_rate=cfg.val_cache_rate,
        num_workers=cfg.num_workers,
        copy_cache=False,
    )
    return val_ds

def get_val_org_dataset(cfg):
    val_ds = CacheDataset(
        data=cfg.data_json["val"],
        transform=cfg.org_val_transforms,
        cache_rate=cfg.val_cache_rate,
        num_workers=cfg.num_workers,
        copy_cache=False,
    )
    return val_ds

def get_optimizer(model, cfg):

    params = model.parameters()
    optimizer = optim.Adam(params, lr=cfg.lr, weight_decay=cfg.weight_decay)

    return optimizer

def get_scheduler(cfg, optimizer, total_steps):

    if cfg.lr_mode == "cosine":
        scheduler = lr_scheduler.CosineAnnealingLR(
            optimizer,
            T_max=cfg.epochs * (total_steps // cfg.batch_size),
            eta_min=cfg.min_lr,
        )

    elif cfg.lr_mode == "warmup_restart":
        scheduler = lr_scheduler.CosineAnnealingWarmRestarts(
            optimizer,
            T_0=cfg.restart_epoch * (total_steps // cfg.batch_size),
            T_mult=1,
            eta_min=cfg.min_lr,
        )

    return scheduler


def create_checkpoint(model, optimizer, epoch, scheduler=None, scaler=None):
    checkpoint = {
        "model": model.state_dict(),
        "optimizer": optimizer.state_dict(),
        "epoch": epoch,
    }

    if scheduler is not None:
        checkpoint["scheduler"] = scheduler.state_dict()

    if scaler is not None:
        checkpoint["scaler"] = scaler.state_dict()
    return checkpoint