# Copyright (c) 2017-present, Facebook, Inc.
# All rights reserved.
#
# This source code is licensed under the license found in the
# LICENSE file in the root directory of this source tree.
#
from __future__ import annotations
import os
import pickle
import numpy as np
import re
from scipy.ndimage import distance_transform_edt as distance
# from skimage import segmentation as skimage_seg
import torch
from torch.utils.data.sampler import Sampler
import torch.distributed as dist
import cv2
from numpy import ndarray

# import networks

# many issues with this function
def load_model(path):
    """Loads model and return it without DataParallel table."""
    if os.path.isfile(path):
        print("=> loading checkpoint '{}'".format(path))
        checkpoint = torch.load(path)

        for key in checkpoint["state_dict"]:
            print(key)

        # size of the top layer
        N = checkpoint["state_dict"]["decoder.out_conv.bias"].size()

        # build skeleton of the model
        sob = "sobel.0.weight" in checkpoint["state_dict"].keys()
        model = models.__dict__[checkpoint["arch"]](sobel=sob, out=int(N[0]))

        # deal with a dataparallel table
        def rename_key(key):
            if not "module" in key:
                return key
            return "".join(key.split(".module"))

        checkpoint["state_dict"] = {
            rename_key(key): val for key, val in checkpoint["state_dict"].items()
        }

        # load weights
        model.load_state_dict(checkpoint["state_dict"])
        print("Loaded")
    else:
        model = None
        print("=> no checkpoint found at '{}'".format(path))
    return model


def load_checkpoint(path, model, optimizer, from_ddp=False):
    """loads previous checkpoint

    Args:
        path (str): path to checkpoint
        model (model): model to restore checkpoint to
        optimizer (optimizer): torch optimizer to load optimizer state_dict to
        from_ddp (bool, optional): load DistributedDataParallel checkpoint to regular model. Defaults to False.

    Returns:
        model, optimizer, epoch_num, loss
    """
    # load checkpoint
    checkpoint = torch.load(path)
    # transfer state_dict from checkpoint to model
    model.load_state_dict(checkpoint["state_dict"])
    # transfer optimizer state_dict from checkpoint to model
    optimizer.load_state_dict(checkpoint["optimizer_state_dict"])
    # track loss
    loss = checkpoint["loss"]
    return model, optimizer, checkpoint["epoch"], loss.item()


def restore_model(logger, snapshot_path, model_num=None):
    """wrapper function to read log dir and load restore a previous checkpoint

    Args:
        logger (Logger): logger object (for info output to console)
        snapshot_path (str): path to checkpoint directory

    Returns:
        model, optimizer, start_epoch, performance
    """
    try:
        # check if there is previous progress to be restored:
        logger.info(f"Snapshot path: {snapshot_path}")
        iter_num = []
        name = "model_iter"
        if model_num:
            name = model_num
        for filename in os.listdir(snapshot_path):
            if name in filename:
                basename, extension = os.path.splitext(filename)
                iter_num.append(int(basename.split("_")[2]))
        iter_num = max(iter_num)
        for filename in os.listdir(snapshot_path):
            if name in filename and str(iter_num) in filename:
                model_checkpoint = filename
    except Exception as e:
        logger.warning(f"Error finding previous checkpoints: {e}")

    try:
        logger.info(f"Restoring model checkpoint: {model_checkpoint}")
        model, optimizer, start_epoch, performance = load_checkpoint(
            snapshot_path + "/" + model_checkpoint, model, optimizer
        )
        logger.info(f"Models restored from iteration {iter_num}")
        return model, optimizer, start_epoch, performance
    except Exception as e:
        logger.warning(f"Unable to restore model checkpoint: {e}, using new model")


def save_checkpoint(epoch, model, optimizer, loss, path):
    """Saves model as checkpoint"""
    torch.save(
        {
            "epoch": epoch,
            "state_dict": model.state_dict(),
            "optimizer_state_dict": optimizer.state_dict(),
            "loss": loss,
        },
        path,
    )


class UnifLabelSampler(Sampler):
    """Samples elements uniformely accross pseudolabels.
    Args:
        N (int): size of returned iterator.
        images_lists: dict of key (target), value (list of data with this target)
    """

    def __init__(self, N, images_lists):
        self.N = N
        self.images_lists = images_lists
        self.indexes = self.generate_indexes_epoch()

    def generate_indexes_epoch(self):
        size_per_pseudolabel = int(self.N / len(self.images_lists)) + 1
        res = np.zeros(size_per_pseudolabel * len(self.images_lists))

        for i in range(len(self.images_lists)):
            indexes = np.random.choice(
                self.images_lists[i],
                size_per_pseudolabel,
                replace=(len(self.images_lists[i]) <= size_per_pseudolabel),
            )
            res[i * size_per_pseudolabel : (i + 1) * size_per_pseudolabel] = indexes

        np.random.shuffle(res)
        return res[: self.N].astype("int")

    def __iter__(self):
        return iter(self.indexes)

    def __len__(self):
        return self.N


class AverageMeter(object):
    """Computes and stores the average and current value"""

    def __init__(self):
        self.reset()

    def reset(self):
        self.val = 0
        self.avg = 0
        self.sum = 0
        self.count = 0

    def update(self, val, n=1):
        self.val = val
        self.sum += val * n
        self.count += n
        self.avg = self.sum / self.count


def learning_rate_decay(optimizer, t, lr_0):
    for param_group in optimizer.param_groups:
        lr = lr_0 / np.sqrt(1 + lr_0 * param_group["weight_decay"] * t)
        param_group["lr"] = lr


class Logger:
    """Class to update every epoch to keep trace of the results
    Methods:
        - log() log and save
    """

    def __init__(self, path):
        self.path = path
        self.data = []

    def log(self, train_point):
        self.data.append(train_point)
        with open(os.path.join(self.path), "wb") as fp:
            pickle.dump(self.data, fp, -1)


def compute_sdf(img_gt, out_shape):
    """
    compute the signed distance map of binary mask
    input: segmentation, shape = (batch_size, x, y, z)
    output: the Signed Distance Map (SDM)
    sdf(x) = 0; x in segmentation boundary
             -inf|x-y|; x in segmentation
             +inf|x-y|; x out of segmentation
    normalize sdf to [-1,1]
    """

    img_gt = img_gt.astype(np.uint8)
    normalized_sdf = np.zeros(out_shape)

    for b in range(out_shape[0]):  # batch size
        posmask = img_gt[b].astype(np.bool)
        if posmask.any():
            negmask = ~posmask
            posdis = distance(posmask)
            negdis = distance(negmask)
            boundary = skimage_seg.find_boundaries(posmask, mode="inner").astype(
                np.uint8
            )
            sdf = (negdis - np.min(negdis)) / (np.max(negdis) - np.min(negdis)) - (
                posdis - np.min(posdis)
            ) / (np.max(posdis) - np.min(posdis))
            sdf[boundary == 1] = 0
            normalized_sdf[b] = sdf
            # assert np.min(sdf) == -1.0, print(np.min(posdis), np.max(posdis), np.min(negdis), np.max(negdis))
            # assert np.max(sdf) ==  1.0, print(np.min(posdis), np.min(negdis), np.max(posdis), np.max(negdis))

    return normalized_sdf


# set up process group for distributed computing
def distributed_setup(rank, world_size):
    os.environ["MASTER_ADDR"] = "localhost"
    os.environ["MASTER_PORT"] = "12355"
    print("setting up dist process group now")
    dist.init_process_group("nccl", rank=rank, world_size=world_size)


def load_ddp_to_nddp(state_dict):
    pattern = re.compile("module")
    for k, v in state_dict.items():
        if re.search("module", k):
            model_dict[re.sub(pattern, "", k)] = v
        else:
            model_dict = state_dict
    return model_dict

def save_osmancheckpoint(epoch, ema_model, model, optimizer, best_dice, best_iter, stu_best_dice, stu_best_iter, path):
    """Saves model as checkpoint"""
    torch.save(
        {
            "epoch": epoch,
            "ema_state_dict": ema_model.state_dict(),
            "state_dict": model.state_dict(),
            "optimizer_state_dict": optimizer.state_dict(),
            "best_dice": best_dice,
            "best_iter": best_iter,
            "stu_best_dice": stu_best_dice,
            "stu_best_iter": stu_best_iter
        },
        path,
    )

def load_osmancheckpoint(path, ema_model, model, optimizer, from_ddp=False):
    """loads previous checkpoint

    Args:
        path (str): path to checkpoint
        model (model): model to restore checkpoint to
        optimizer (optimizer): torch optimizer to load optimizer state_dict to
        from_ddp (bool, optional): load DistributedDataParallel checkpoint to regular model. Defaults to False.

    Returns:
        model, optimizer, epoch_num, loss
    """
    # load checkpoint
    checkpoint = torch.load(path)
    # transfer state_dict from checkpoint to model
    ema_model.load_state_dict(checkpoint["ema_state_dict"])
    model.load_state_dict(checkpoint["state_dict"])
    # transfer optimizer state_dict from checkpoint to model
    optimizer.load_state_dict(checkpoint["optimizer_state_dict"])
    
    best_dice, best_iter = checkpoint['best_dice'], checkpoint['best_iter']
    stu_best_dice, stu_best_iter = checkpoint['stu_best_dice'], checkpoint['stu_best_iter']
    return  checkpoint["epoch"], ema_model, model, optimizer, best_dice, best_iter, stu_best_dice, stu_best_iter

def draw_contour(
    image: ndarray,
    label: ndarray,
    color: tuple[int, int, int],
) -> ndarray:
    """Draw contour of label on image with color.

    Uses `cv2.dilate` to find contour.

    Args:
        image: ndarray of shape (h, w, 3)
        label: ndarray of shape (h, w) where non-zero values are foreground
        color: color of contour
    """
    binary = (label > 0).astype('uint8') * 255
    dilated = cv2.dilate(binary, np.ones((3, 3), dtype='uint8'), iterations=1)
    contour_mask = (dilated - binary) > 0
    image[contour_mask] = color
    return image


def make_prediction(
    image: ndarray,
    *predictions: tuple[ndarray, tuple[int, int, int]],
) -> ndarray:
    """Draw contour of predictions on image.

    Args:
        image: ndarray of shape (h, w) or (h, w, 3)
            if image is (h, w), it will be repeated to (h, w, 3)
        predictions: each prediction is a tuple of (label, color)
            label: ndarray of shape (c, h, w) or (h, w)
                where c is number of classes or 1
            color: color of contour
    """
    if image.ndim == 2:
        # expand to (h, w, 3)
        image = np.repeat(image[..., np.newaxis], 3, axis=-1)

    for prediction, color in predictions:
        if prediction.ndim == 2:
            # expand to (1, h, w)
            prediction = prediction[np.newaxis, ...]
        for class_map in prediction:
            image = draw_contour(image, class_map, color)
    return image

def draw_contour_and_save(img, pred, mask, save_path='./img/1/example.png'):
    color_gt = (255,0,0)
    color_pred_list = [(0,255,0),(0,0,255),(255,255,0),(255,0,255)]
    if len(pred.shape) == 2:
        pred = pred.unsqueeze(0)
    if len(mask.shape) == 2:
        mask = mask.unsqueeze(0)
    if len(img.shape) == 2:
        img = img.unsqueeze(0)
    if len(img) == 1:
        img = img.repeat(3,1,1)
    img = (img-img.min())/(img.max()-img.min()) * 255
    img = img.cpu().numpy().transpose(1,2,0)
    classes = len(pred)
    for i in range(classes):
        predi = pred[i].float().cpu().numpy()
        maski = mask[i].float().cpu().numpy()
        seg_res = make_prediction(img, (predi, color_pred_list[i]), (maski, color_gt))
        seg_res = cv2.cvtColor(seg_res, cv2.COLOR_RGB2BGR)
    cv2.imwrite(save_path, seg_res)
    
def draw_mask_and_save(img, pred, save_path='./img/1/example.png'):
    color_list = [(255,0,0),(0,255,0),(0,0,255),(255,255,0),(255,0,255)]
    if len(pred.shape) == 2:
        pred = pred.unsqueeze(0)
    if len(img.shape) == 2:
        img = img.unsqueeze(0)
    if len(img) == 1:
        img = img.repeat(3,1,1)
    d, h, w = img.shape
    rgb = np.zeros((h, w, 3))
    mask = np.ones((h, w, 3))
    if img.min() < -0.5:
        img = (img+1)*127.5
    elif img.max() < 1.5:
        img = img*255
    img = img.cpu().numpy().transpose(1,2,0)
    classes = len(pred)
    for i in reversed(range(classes)):
        mask[pred[i] == 1] = (0.5, 0.5, 0.5)
        rgb[pred[i] == 1] = color_list[i]
    res = (img+rgb)*mask
    res = res.astype(np.uint8)
    res = cv2.cvtColor(res, cv2.COLOR_RGB2BGR)
    cv2.imwrite(save_path, res)