from collections import OrderedDict

import torch
import torchvision

import torch.nn.functional as F
from torch import nn, Tensor

from torchvision.ops import boxes as box_ops

from torchvision.ops import roi_align

from torchvision.models.detection import _utils as det_utils

from torch.jit.annotations import Optional, List, Dict, Tuple
import numpy as np

from TorchUtils import tensor_isin
from nucls_model.torchvision_detection_utils.transforms import \
    remove_degenerate_bboxes

@torch.jit._script_if_tracing
def global_nms(boxes, scores, iou_threshold):
    # type: (Tensor, Tensor, float) -> Tensor
    """
    Performs non-maximum suppression globally (regardless of category).

    Parameters
    ----------
    boxes : Tensor[N, 4]
        boxes where NMS will be performed. They
        are expected to be in (x1, y1, x2, y2) format
    scores : Tensor[N]
        scores for each one of the boxes
    iou_threshold : float
        discards all overlapping boxes
        with IoU > iou_threshold

    Returns
    -------
    keep : Tensor
        int64 tensor with the indices of
        the elements that have been kept by NMS, sorted
        in decreasing order of scores
    """
    if boxes.numel() == 0:
        return torch.empty((0,), dtype=torch.int64, device=boxes.device)
    else:
        keep = box_ops.nms(boxes, scores, iou_threshold)
        return keep

# noinspection LongLine,PyTypeHints
def fastrcnn_loss(
        class_logits, box_regression, labels, regression_targets,
        ignore_label=None, batched_nms=True):
    # type: (Tensor, Tensor, List[Tensor], List[Tensor]) -> Tuple[Tensor, Tensor]
    """
    Computes the loss for Faster R-CNN.

    Arguments:
        class_logits (Tensor)
        box_regression (Tensor)
        labels (list[BoxList])
        regression_targets (Tensor)
        ignore_label (int)
        batched_nms (bool)

    Returns:
        classification_loss (Tensor)
        box_loss (Tensor)
    """

    labels = torch.cat(labels, dim=0)
    regression_targets = torch.cat(regression_targets, dim=0)

    # Mohamed: Ignore certain nuclei from classification loss.
    # Do not to modify in-place since labels are used later here!!
    clkeep = labels != ignore_label
    classification_loss = F.cross_entropy(class_logits[clkeep], labels[clkeep])

    # get indices that correspond to the regression targets for
    # the corresponding ground truth labels, to be used with
    # advanced indexing
    sampled_pos_inds_subset = torch.nonzero(labels > 0).squeeze(1)
    if batched_nms:
        labels_pos = labels[sampled_pos_inds_subset]
    else:
        # Mohamed: only two classes (bckgrnd, frgrnd) for box regression
        labs = 0 + labels
        labs[labs > 1] = 1
        labels_pos = labs[sampled_pos_inds_subset]
    N, num_classes = class_logits.shape
    box_regression = box_regression.reshape(N, -1, 4)

    box_loss = det_utils.smooth_l1_loss(
        box_regression[sampled_pos_inds_subset, labels_pos],
        regression_targets[sampled_pos_inds_subset],
        beta=1 / 9,
        size_average=False,
    )
    box_loss = box_loss / labels.numel()

    return classification_loss, box_loss

def maskrcnn_inference(x, labels):
    # type: (Tensor, List[Tensor]) -> List[Tensor]
    """
    From the results of the CNN, post process the masks
    by taking the mask corresponding to the class with max
    probability (which are of fixed size and directly output
    by the CNN) and return the masks in the mask field of the BoxList.

    Arguments:
        x (Tensor): the mask logits
        labels (list[BoxList]): bounding boxes that are used as
            reference, one for ech image

    Returns:
        results (list[BoxList]): one BoxList for each image, containing
            the extra field mask
    """
    mask_prob = x.sigmoid()

    # select masks corresponding to the predicted classes
    num_masks = x.shape[0]
    boxes_per_image = [label.shape[0] for label in labels]
    labels = torch.cat(labels)
    index = torch.arange(num_masks, device=labels.device)
    mask_prob = mask_prob[index, labels][:, None]
    mask_prob = mask_prob.split(boxes_per_image, dim=0)

    return mask_prob

def project_masks_on_boxes(gt_masks, boxes, matched_idxs, M):
    # type: (Tensor, Tensor, Tensor, int) -> Tensor
    """
    Given segmentation masks and the bounding boxes corresponding
    to the location of the masks in the image, this function
    crops and resizes the masks in the position defined by the
    boxes. This prepares the masks for them to be fed to the
    loss computation as the targets.
    """
    matched_idxs = matched_idxs.to(boxes)
    rois = torch.cat([matched_idxs[:, None], boxes], dim=1)
    gt_masks = gt_masks[:, None].to(rois)
    return roi_align(gt_masks, rois, (M, M), 1.)[:, 0]

# noinspection LongLine,PyTypeHints
def maskrcnn_loss(
        mask_logits, proposals, gt_masks, gt_labels, mask_matched_idxs,
        gt_ismask=None):
    # type: (Tensor, List[Tensor], List[Tensor], List[Tensor], List[Tensor]) -> Tensor
    """
    Arguments:
        mask_logits (Tensor)
        proposals (list[BoxList])

    Return:
        mask_loss (Tensor): scalar tensor containing the loss

    Parameters
    ----------
    gt_ismask
    mask_matched_idxs
    gt_labels
    proposals
    mask_logits
    gt_masks
    """

    discretization_size = mask_logits.shape[-1]
    labels = [gt_label[idxs] for gt_label, idxs in zip(gt_labels, mask_matched_idxs)]
    mask_targets = [
        project_masks_on_boxes(m, p, i, discretization_size)
        for m, p, i in zip(gt_masks, proposals, mask_matched_idxs)
    ]

    labels = torch.cat(labels, dim=0)
    mask_targets = torch.cat(mask_targets, dim=0)

    # Mohamed: added this to disregard nuclei without masks (bboxes)
    if gt_ismask is not None:
        keep_idxs = [torch.where(t > 0)[0] for t in gt_ismask]
        keep_list = [
            tensor_isin(idxs, keep_idx)
            for idxs, keep_idx in zip(mask_matched_idxs, keep_idxs)
        ]
        keep_array = torch.cat(keep_list).type(torch.bool)
        mask_logits = mask_logits[keep_array, ...]
        mask_targets = mask_targets[keep_array, ...]
        labels = labels[keep_array]

    # torch.mean (in binary_cross_entropy_with_logits) doesn't
    # accept empty tensors, so handle it separately
    if mask_targets.numel() == 0:
        return mask_logits.sum() * 0

    mask_loss = F.binary_cross_entropy_with_logits(
        mask_logits[torch.arange(labels.shape[0], device=labels.device), labels], mask_targets
    )
    return mask_loss

def keypoints_to_heatmap(keypoints, rois, heatmap_size):
    # type: (Tensor, Tensor, int) -> Tuple[Tensor, Tensor]
    offset_x = rois[:, 0]
    offset_y = rois[:, 1]
    scale_x = heatmap_size / (rois[:, 2] - rois[:, 0])
    scale_y = heatmap_size / (rois[:, 3] - rois[:, 1])

    offset_x = offset_x[:, None]
    offset_y = offset_y[:, None]
    scale_x = scale_x[:, None]
    scale_y = scale_y[:, None]

    x = keypoints[..., 0]
    y = keypoints[..., 1]

    x_boundary_inds = x == rois[:, 2][:, None]
    y_boundary_inds = y == rois[:, 3][:, None]

    x = (x - offset_x) * scale_x
    x = x.floor().long()
    y = (y - offset_y) * scale_y
    y = y.floor().long()

    x[x_boundary_inds] = heatmap_size - 1
    y[y_boundary_inds] = heatmap_size - 1

    valid_loc = (x >= 0) & (y >= 0) & (x < heatmap_size) & (y < heatmap_size)
    vis = keypoints[..., 2] > 0
    valid = (valid_loc & vis).long()

    lin_ind = y * heatmap_size + x
    heatmaps = lin_ind * valid

    return heatmaps, valid

# noinspection LongLine
def _onnx_heatmaps_to_keypoints(maps, maps_i, roi_map_width, roi_map_height,
                                widths_i, heights_i, offset_x_i, offset_y_i):
    num_keypoints = torch.scalar_tensor(maps.size(1), dtype=torch.int64)

    width_correction = widths_i / roi_map_width
    height_correction = heights_i / roi_map_height

    roi_map = F.interpolate(
        maps_i[:, None], size=(int(roi_map_height), int(roi_map_width)), mode='bicubic', align_corners=False)[:, 0]

    w = torch.scalar_tensor(roi_map.size(2), dtype=torch.int64)
    pos = roi_map.reshape(num_keypoints, -1).argmax(dim=1)

    x_int = (pos % w)
    y_int = ((pos - x_int) // w)

    x = (torch.tensor(0.5, dtype=torch.float32) + x_int.to(dtype=torch.float32)) * \
        width_correction.to(dtype=torch.float32)
    y = (torch.tensor(0.5, dtype=torch.float32) + y_int.to(dtype=torch.float32)) * \
        height_correction.to(dtype=torch.float32)

    xy_preds_i_0 = x + offset_x_i.to(dtype=torch.float32)
    xy_preds_i_1 = y + offset_y_i.to(dtype=torch.float32)
    xy_preds_i_2 = torch.ones((xy_preds_i_1.shape), dtype=torch.float32)  # noqa
    xy_preds_i = torch.stack([xy_preds_i_0.to(dtype=torch.float32),
                              xy_preds_i_1.to(dtype=torch.float32),
                              xy_preds_i_2.to(dtype=torch.float32)], 0)

    # mTODO: simplify when indexing without rank will be supported by ONNX
    base = num_keypoints * num_keypoints + num_keypoints + 1
    ind = torch.arange(num_keypoints)
    ind = ind.to(dtype=torch.int64) * base
    end_scores_i = roi_map.index_select(1, y_int.to(dtype=torch.int64)) \
        .index_select(2, x_int.to(dtype=torch.int64)).view(-1).index_select(0, ind.to(dtype=torch.int64))

    return xy_preds_i, end_scores_i

# noinspection LongLine
@torch.jit._script_if_tracing
def _onnx_heatmaps_to_keypoints_loop(maps, rois, widths_ceil, heights_ceil,
                                     widths, heights, offset_x, offset_y, num_keypoints):
    xy_preds = torch.zeros((0, 3, int(num_keypoints)), dtype=torch.float32, device=maps.device)
    end_scores = torch.zeros((0, int(num_keypoints)), dtype=torch.float32, device=maps.device)

    for i in range(int(rois.size(0))):
        xy_preds_i, end_scores_i = _onnx_heatmaps_to_keypoints(maps, maps[i],
                                                               widths_ceil[i], heights_ceil[i],
                                                               widths[i], heights[i],
                                                               offset_x[i], offset_y[i])
        xy_preds = torch.cat((xy_preds.to(dtype=torch.float32),
                              xy_preds_i.unsqueeze(0).to(dtype=torch.float32)), 0)
        end_scores = torch.cat((end_scores.to(dtype=torch.float32),
                                end_scores_i.to(dtype=torch.float32).unsqueeze(0)), 0)
    return xy_preds, end_scores

# noinspection LongLine
def heatmaps_to_keypoints(maps, rois):
    """Extract predicted keypoint locations from heatmaps. Output has shape
    (#rois, 4, #keypoints) with the 4 rows corresponding to (x, y, logit, prob)
    for each keypoint.
    """
    # This function converts a discrete image coordinate in a HEATMAP_SIZE x
    # HEATMAP_SIZE image to a continuous keypoint coordinate. We maintain
    # consistency with keypoints_to_heatmap_labels by using the conversion from
    # Heckbert 1990: c = d + 0.5, where d is a discrete coordinate and c is a
    # continuous coordinate.
    offset_x = rois[:, 0]
    offset_y = rois[:, 1]

    widths = rois[:, 2] - rois[:, 0]
    heights = rois[:, 3] - rois[:, 1]
    widths = widths.clamp(min=1)
    heights = heights.clamp(min=1)
    widths_ceil = widths.ceil()
    heights_ceil = heights.ceil()

    num_keypoints = maps.shape[1]

    if torchvision._is_tracing():
        xy_preds, end_scores = _onnx_heatmaps_to_keypoints_loop(maps, rois,
                                                                widths_ceil, heights_ceil, widths, heights,
                                                                offset_x, offset_y,
                                                                torch.scalar_tensor(num_keypoints, dtype=torch.int64))
        return xy_preds.permute(0, 2, 1), end_scores

    xy_preds = torch.zeros((len(rois), 3, num_keypoints), dtype=torch.float32, device=maps.device)
    end_scores = torch.zeros((len(rois), num_keypoints), dtype=torch.float32, device=maps.device)
    for i in range(len(rois)):
        roi_map_width = int(widths_ceil[i].item())
        roi_map_height = int(heights_ceil[i].item())
        width_correction = widths[i] / roi_map_width
        height_correction = heights[i] / roi_map_height
        roi_map = F.interpolate(
            maps[i][:, None], size=(roi_map_height, roi_map_width), mode='bicubic', align_corners=False)[:, 0]
        # roi_map_probs = scores_to_probs(roi_map.copy())
        w = roi_map.shape[2]
        pos = roi_map.reshape(num_keypoints, -1).argmax(dim=1)

        x_int = pos % w
        y_int = (pos - x_int) // w
        # assert (roi_map_probs[k, y_int, x_int] ==
        #         roi_map_probs[k, :, :].max())
        x = (x_int.float() + 0.5) * width_correction
        y = (y_int.float() + 0.5) * height_correction
        xy_preds[i, 0, :] = x + offset_x[i]
        xy_preds[i, 1, :] = y + offset_y[i]
        xy_preds[i, 2, :] = 1
        end_scores[i, :] = roi_map[torch.arange(num_keypoints), y_int, x_int]

    return xy_preds.permute(0, 2, 1), end_scores

# noinspection LongLine
def keypointrcnn_loss(keypoint_logits, proposals, gt_keypoints, keypoint_matched_idxs):
    # type: (Tensor, List[Tensor], List[Tensor], List[Tensor]) -> Tensor
    N, K, H, W = keypoint_logits.shape
    assert H == W
    discretization_size = H
    heatmaps = []
    valid = []
    for proposals_per_image, gt_kp_in_image, midx in zip(proposals, gt_keypoints, keypoint_matched_idxs):
        kp = gt_kp_in_image[midx]
        heatmaps_per_image, valid_per_image = keypoints_to_heatmap(
            kp, proposals_per_image, discretization_size
        )
        heatmaps.append(heatmaps_per_image.view(-1))
        valid.append(valid_per_image.view(-1))

    keypoint_targets = torch.cat(heatmaps, dim=0)
    valid = torch.cat(valid, dim=0).to(dtype=torch.uint8)
    valid = torch.nonzero(valid).squeeze(1)

    # torch.mean (in binary_cross_entropy_with_logits) does'nt
    # accept empty tensors, so handle it sepaartely
    if keypoint_targets.numel() == 0 or len(valid) == 0:
        return keypoint_logits.sum() * 0

    keypoint_logits = keypoint_logits.view(N * K, H * W)

    keypoint_loss = F.cross_entropy(keypoint_logits[valid], keypoint_targets[valid])
    return keypoint_loss

# noinspection LongLine
def keypointrcnn_inference(x, boxes):
    # type: (Tensor, List[Tensor]) -> Tuple[List[Tensor], List[Tensor]]
    kp_probs = []
    kp_scores = []

    boxes_per_image = [box.size(0) for box in boxes]  # noqa
    x2 = x.split(boxes_per_image, dim=0)

    for xx, bb in zip(x2, boxes):
        kp_prob, scores = heatmaps_to_keypoints(xx, bb)
        kp_probs.append(kp_prob)
        kp_scores.append(scores)

    return kp_probs, kp_scores

# noinspection LongLine,DuplicatedCode
def _onnx_expand_boxes(boxes, scale):
    # type: (Tensor, float) -> Tensor
    w_half = (boxes[:, 2] - boxes[:, 0]) * .5
    h_half = (boxes[:, 3] - boxes[:, 1]) * .5
    x_c = (boxes[:, 2] + boxes[:, 0]) * .5
    y_c = (boxes[:, 3] + boxes[:, 1]) * .5

    w_half = w_half.to(dtype=torch.float32) * scale
    h_half = h_half.to(dtype=torch.float32) * scale

    boxes_exp0 = x_c - w_half
    boxes_exp1 = y_c - h_half
    boxes_exp2 = x_c + w_half
    boxes_exp3 = y_c + h_half
    boxes_exp = torch.stack((boxes_exp0, boxes_exp1, boxes_exp2, boxes_exp3), 1)
    return boxes_exp

# the next two functions should be merged inside Masker
# but are kept here for the moment while we need them
# temporarily for paste_mask_in_image
# noinspection DuplicatedCode
def expand_boxes(boxes, scale):
    # type: (Tensor, float) -> Tensor
    if torchvision._is_tracing():
        return _onnx_expand_boxes(boxes, scale)
    w_half = (boxes[:, 2] - boxes[:, 0]) * .5
    h_half = (boxes[:, 3] - boxes[:, 1]) * .5
    x_c = (boxes[:, 2] + boxes[:, 0]) * .5
    y_c = (boxes[:, 3] + boxes[:, 1]) * .5

    w_half *= scale
    h_half *= scale

    boxes_exp = torch.zeros_like(boxes)
    boxes_exp[:, 0] = x_c - w_half
    boxes_exp[:, 2] = x_c + w_half
    boxes_exp[:, 1] = y_c - h_half
    boxes_exp[:, 3] = y_c + h_half
    return boxes_exp

# noinspection LongLine
@torch.jit.unused
def expand_masks_tracing_scale(M, padding):
    # type: (int, int) -> float
    return torch.tensor(M + 2 * padding).to(torch.float32) / torch.tensor(M).to(torch.float32)

# noinspection LongLine
def expand_masks(mask, padding):
    # type: (Tensor, int) -> Tuple[Tensor, float]
    M = mask.shape[-1]
    if torch._C._get_tracing_state():  # could not import is_tracing(), not sure why
        scale = expand_masks_tracing_scale(M, padding)
    else:
        scale = float(M + 2 * padding) / M
    padded_mask = torch.nn.functional.pad(mask, (padding,) * 4)
    return padded_mask, scale

# noinspection LongLine
def paste_mask_in_image(mask, box, im_h, im_w, im_mask=None, ocode=None):
    # type: (Tensor, Tensor, int, int) -> Tensor
    TO_REMOVE = 1
    w = int(box[2] - box[0] + TO_REMOVE)
    h = int(box[3] - box[1] + TO_REMOVE)
    w = max(w, 1)
    h = max(h, 1)

    # Set shape to [batchxCxHxW]
    mask = mask.expand((1, 1, -1, -1))

    # Resize mask
    mask = F.interpolate(mask, size=(h, w), mode='bilinear', align_corners=False)
    mask = mask[0][0]

    x_0 = max(box[0], 0)
    x_1 = min(box[2] + 1, im_w)
    y_0 = max(box[1], 0)
    y_1 = min(box[3] + 1, im_h)

    if im_mask is None:
        # sparse mask (only one object per channel)
        im_mask = torch.zeros((im_h, im_w), dtype=mask.dtype, device=mask.device)
        im_mask[y_0:y_1, x_0:x_1] = mask[
            (y_0 - box[1]):(y_1 - box[1]), (x_0 - box[0]):(x_1 - box[0])
        ]
    else:
        # Mohamed: dense mask (just one channel, where code represents object)
        # IMPORTANT NOTE: this means we threshold probabilities.
        patch = im_mask[y_0:y_1, x_0:x_1]
        omask = mask[
            (y_0 - box[1]):(y_1 - box[1]), (x_0 - box[0]):(x_1 - box[0])
        ]
        patch[omask > 0.5] = ocode

    return im_mask

# noinspection LongLine
def _onnx_paste_mask_in_image(mask, box, im_h, im_w):
    one = torch.ones(1, dtype=torch.int64)
    zero = torch.zeros(1, dtype=torch.int64)

    w = (box[2] - box[0] + one)
    h = (box[3] - box[1] + one)
    w = torch.max(torch.cat((w, one)))
    h = torch.max(torch.cat((h, one)))

    # Set shape to [batchxCxHxW]
    mask = mask.expand((1, 1, mask.size(0), mask.size(1)))

    # Resize mask
    mask = F.interpolate(mask, size=(int(h), int(w)), mode='bilinear', align_corners=False)
    mask = mask[0][0]

    x_0 = torch.max(torch.cat((box[0].unsqueeze(0), zero)))
    x_1 = torch.min(torch.cat((box[2].unsqueeze(0) + one, im_w.unsqueeze(0))))
    y_0 = torch.max(torch.cat((box[1].unsqueeze(0), zero)))
    y_1 = torch.min(torch.cat((box[3].unsqueeze(0) + one, im_h.unsqueeze(0))))

    unpaded_im_mask = mask[(y_0 - box[1]):(y_1 - box[1]),
                           (x_0 - box[0]):(x_1 - box[0])]

    # mTODO : replace below with a dynamic padding when support is added in ONNX

    # pad y
    zeros_y0 = torch.zeros(y_0, unpaded_im_mask.size(1))
    zeros_y1 = torch.zeros(im_h - y_1, unpaded_im_mask.size(1))
    concat_0 = torch.cat((zeros_y0,
                          unpaded_im_mask.to(dtype=torch.float32),
                          zeros_y1), 0)[0:im_h, :]
    # pad x
    zeros_x0 = torch.zeros(concat_0.size(0), x_0)
    zeros_x1 = torch.zeros(concat_0.size(0), im_w - x_1)
    im_mask = torch.cat((zeros_x0,
                         concat_0,
                         zeros_x1), 1)[:, :im_w]
    return im_mask

@torch.jit._script_if_tracing
def _onnx_paste_masks_in_image_loop(masks, boxes, im_h, im_w):
    res_append = torch.zeros(0, im_h, im_w)
    for i in range(masks.size(0)):
        mask_res = _onnx_paste_mask_in_image(masks[i][0], boxes[i], im_h, im_w)
        mask_res = mask_res.unsqueeze(0)
        res_append = torch.cat((res_append, mask_res))
    return res_append

# noinspection LongLine
def paste_masks_in_image(masks, boxes, img_shape, padding=1):
    # type: (Tensor, Tensor, Tuple[int, int], int) -> Tensor
    masks, scale = expand_masks(masks, padding=padding)
    boxes = expand_boxes(boxes, scale).to(dtype=torch.int64)
    im_h, im_w = img_shape

    if torchvision._is_tracing():
        return _onnx_paste_masks_in_image_loop(masks, boxes,
                                               torch.scalar_tensor(im_h, dtype=torch.int64),
                                               torch.scalar_tensor(im_w, dtype=torch.int64))[:, None]
    res = [
        paste_mask_in_image(m[0], b, im_h, im_w)
        for m, b in zip(masks, boxes)
    ]
    if len(res) > 0:
        ret = torch.stack(res, dim=0)[:, None]
    else:
        ret = masks.new_empty((0, 1, im_h, im_w))
    return ret

# Mohamed: dense mask (code represents object)
def paste_and_densify_masks_in_image(masks, boxes, img_shape, padding=1):
    # type: (Tensor, Tensor, Tuple[int, int], int) -> Tensor

    masks, scale = expand_masks(masks, padding=padding)
    boxes = expand_boxes(boxes, scale).to(dtype=torch.int64)
    im_h, im_w = img_shape

    if torchvision._is_tracing():
        raise NotImplementedError(
            "I didn't support tracing yet. See paste_masks_in_image()."
        )

    # IMPORTANT NOTE: in the dense mask, each object id encoded by a code
    # that corresponds to (1 + idx), where idx is the index of the object
    # in the labels/bbox tensors. For example, where im_mask == 5 corresponds
    # to the nucleus whose label is in the 4th index in the labels tensor

    # Find the order in which to overlay objects in the mask, from big to small
    # This prevents a big object from covering a small one. Noe, however, that
    # 1. This is still a "lossy" process, but it saves up on speed and makes
    #    the memory requirement less dependent on the no of objects
    # 2. We maintain the object codes' correspondence to (1 + idx)
    areas = (boxes[:, 2] - boxes[:, 0]) * (boxes[:, 3] - boxes[:, 1])
    idxs = torch.argsort(areas, descending=True)

    im_mask = torch.zeros((im_h, im_w), dtype=torch.int64, device=masks.device)
    for idx in idxs:
        im_mask = paste_mask_in_image(
            mask=masks[idx][0], box=boxes[idx], im_h=im_h, im_w=im_w,
            im_mask=im_mask, ocode=idx + 1)

    return im_mask

# noinspection LongLine
class RoIHeads(torch.nn.Module):
    __annotations__ = {
        'box_coder': det_utils.BoxCoder,
        'proposal_matcher': det_utils.Matcher,
        'fg_bg_sampler': det_utils.BalancedPositiveNegativeSampler,
    }

    def __init__(self,
                 box_roi_pool,
                 box_head,
                 box_predictor,
                 # Faster R-CNN training
                 fg_iou_thresh, bg_iou_thresh,
                 batch_size_per_image, positive_fraction,
                 bbox_reg_weights,
                 # Faster R-CNN inference
                 score_thresh,
                 nms_thresh,
                 detections_per_img,
                 # Mask
                 mask_roi_pool=None,
                 mask_head=None,
                 mask_predictor=None,
                 keypoint_roi_pool=None,
                 keypoint_head=None,
                 keypoint_predictor=None,
                 # added by Mohamed
                 batched_nms=True,  # inference nms independently per class?
                 indep_classif_boxes=False,
                 classification_bbox_size=None,  # float
                 cconvhead=None,  # extra conv layers for classification
                 sattention_head=None,  # nuclei are aware of each other
                 ignore_label: int = None,  # label to ignore in classif. loss
                 ):
        super(RoIHeads, self).__init__()

        if indep_classif_boxes:
            assert classification_bbox_size is not None

        self.box_similarity = box_ops.box_iou
        # assign ground-truth boxes for each proposal
        self.proposal_matcher = det_utils.Matcher(
            fg_iou_thresh,
            bg_iou_thresh,
            allow_low_quality_matches=False)

        self.fg_bg_sampler = det_utils.BalancedPositiveNegativeSampler(
            batch_size_per_image,
            positive_fraction)

        if bbox_reg_weights is None:
            bbox_reg_weights = (10., 10., 5., 5.)
        self.box_coder = det_utils.BoxCoder(bbox_reg_weights)

        self.box_roi_pool = box_roi_pool
        self.box_head = box_head
        self.box_predictor = box_predictor

        self.score_thresh = score_thresh
        self.nms_thresh = nms_thresh
        self.detections_per_img = detections_per_img

        self.mask_roi_pool = mask_roi_pool
        self.mask_head = mask_head
        self.mask_predictor = mask_predictor

        self.keypoint_roi_pool = keypoint_roi_pool
        self.keypoint_head = keypoint_head
        self.keypoint_predictor = keypoint_predictor

        # added by Mohamed
        self.batched_nms = batched_nms
        self.indep_classif_boxes = indep_classif_boxes
        self.classification_bbox_size = classification_bbox_size
        self.halfclbox = self.classification_bbox_size / 2 if \
            classification_bbox_size is not None else None
        self.cconvhead = cconvhead
        self.sattention_head = sattention_head
        self.ignore_label = ignore_label

    def has_mask(self):
        if self.mask_roi_pool is None:
            return False
        if self.mask_head is None:
            return False
        if self.mask_predictor is None:
            return False
        return True

    def has_keypoint(self):
        if self.keypoint_roi_pool is None:
            return False
        if self.keypoint_head is None:
            return False
        if self.keypoint_predictor is None:
            return False
        return True

    def assign_targets_to_proposals(self, proposals, gt_boxes, gt_labels):
        # type: (List[Tensor], List[Tensor], List[Tensor]) -> Tuple[List[Tensor], List[Tensor]]
        matched_idxs = []
        labels = []
        for proposals_in_image, gt_boxes_in_image, gt_labels_in_image in zip(proposals, gt_boxes, gt_labels):

            if gt_boxes_in_image.numel() == 0:
                # Background image
                device = proposals_in_image.device
                clamped_matched_idxs_in_image = torch.zeros(
                    (proposals_in_image.shape[0],), dtype=torch.int64, device=device
                )
                labels_in_image = torch.zeros(
                    (proposals_in_image.shape[0],), dtype=torch.int64, device=device
                )
            else:
                #  set to self.box_similarity when https://github.com/pytorch/pytorch/issues/27495 lands
                match_quality_matrix = box_ops.box_iou(gt_boxes_in_image, proposals_in_image)
                matched_idxs_in_image = self.proposal_matcher(match_quality_matrix)

                clamped_matched_idxs_in_image = matched_idxs_in_image.clamp(min=0)

                labels_in_image = gt_labels_in_image[clamped_matched_idxs_in_image]
                labels_in_image = labels_in_image.to(dtype=torch.int64)

                # Label background (below the low threshold)
                bg_inds = matched_idxs_in_image == self.proposal_matcher.BELOW_LOW_THRESHOLD
                labels_in_image[bg_inds] = 0

                # Label ignore proposals (between low and high thresholds)
                ignore_inds = matched_idxs_in_image == self.proposal_matcher.BETWEEN_THRESHOLDS
                labels_in_image[ignore_inds] = -1  # -1 is ignored by sampler

            matched_idxs.append(clamped_matched_idxs_in_image)
            labels.append(labels_in_image)
        return matched_idxs, labels

    def subsample(self, labels):
        # type: (List[Tensor]) -> List[Tensor]
        sampled_pos_inds, sampled_neg_inds = self.fg_bg_sampler(labels)
        sampled_inds = []
        for img_idx, (pos_inds_img, neg_inds_img) in enumerate(
            zip(sampled_pos_inds, sampled_neg_inds)
        ):
            img_sampled_inds = torch.nonzero(pos_inds_img | neg_inds_img).squeeze(1)
            sampled_inds.append(img_sampled_inds)
        return sampled_inds

    # noinspection PyMethodMayBeStatic
    def add_gt_proposals(self, proposals, gt_boxes):
        # type: (List[Tensor], List[Tensor]) -> List[Tensor]
        proposals = [
            torch.cat((proposal, gt_box))
            for proposal, gt_box in zip(proposals, gt_boxes)
        ]

        return proposals

    def check_targets(self, targets):
        # type: (Optional[List[Dict[str, Tensor]]]) -> None
        assert targets is not None
        assert all(["boxes" in t for t in targets])
        assert all(["labels" in t for t in targets])
        if self.has_mask():
            assert all(["masks" in t for t in targets])

    def select_training_samples(self,
                                proposals,  # type: List[Tensor]
                                targets     # type: Optional[List[Dict[str, Tensor]]]
                                ):
        # type: (...) -> Tuple[List[Tensor], List[Tensor], List[Tensor], List[Tensor]]
        self.check_targets(targets)
        assert targets is not None
        dtype = proposals[0].dtype
        device = proposals[0].device

        gt_boxes = [t["boxes"].to(dtype) for t in targets]
        gt_labels = [t["labels"] for t in targets]

        # append ground-truth bboxes to propos
        proposals = self.add_gt_proposals(proposals, gt_boxes)

        # get matching gt indices for each proposal
        matched_idxs, labels = self.assign_targets_to_proposals(proposals, gt_boxes, gt_labels)
        # sample a fixed proportion of positive-negative proposals
        sampled_inds = self.subsample(labels)
        matched_gt_boxes = []
        num_images = len(proposals)
        for img_id in range(num_images):
            img_sampled_inds = sampled_inds[img_id]
            proposals[img_id] = proposals[img_id][img_sampled_inds]
            labels[img_id] = labels[img_id][img_sampled_inds]
            matched_idxs[img_id] = matched_idxs[img_id][img_sampled_inds]

            gt_boxes_in_image = gt_boxes[img_id]
            if gt_boxes_in_image.numel() == 0:
                gt_boxes_in_image = torch.zeros((1, 4), dtype=dtype, device=device)
            matched_gt_boxes.append(gt_boxes_in_image[matched_idxs[img_id]])

        regression_targets = self.box_coder.encode(matched_gt_boxes, proposals)
        return proposals, matched_idxs, labels, regression_targets

    def postprocess_detections_globalnms(self,
                               class_logits,    # type: Tensor
                               box_regression,  # type: Tensor
                               proposals,       # type: List[Tensor]
                               image_shapes     # type: List[Tuple[int, int]]
                               ):
        # type: (...) -> Tuple[List[Tensor], List[Tensor], List[Tensor]]
        device = class_logits.device
        hdim = class_logits.shape[0]
        num_classes = class_logits.shape[-1]

        boxes_per_image = [boxes_in_image.shape[0] for boxes_in_image in proposals]
        pred_boxes = self.box_coder.decode(box_regression, proposals)

        pred_scores = F.softmax(class_logits, -1)

        pred_boxes_list = pred_boxes.split(boxes_per_image, 0)
        pred_scores_list = pred_scores.split(boxes_per_image, 0)

        all_boxes = []
        all_objectness = []  # probability this is an object
        all_labels = []
        all_probabs = []  # probabilities for each class
        all_keptidxs = []
        for boxes, scores, image_shape in zip(pred_boxes_list, pred_scores_list, image_shapes):

            boxes = box_ops.clip_boxes_to_image(boxes, image_shape)

            # remove predictions with the background label
            boxes = boxes[:, 1, :]
            scores = scores[:, 1:]

            # we define the "objectness" as the sum of non-background scores
            objectness = scores.sum(1)

            # indices of detections that are kept nms and postprocessing
            keptidxs = torch.arange(hdim, device=device)

            # Mohamed: prevent (set as zero prob.) ignore_label prediction
            if self.ignore_label is not None:
                ignore = [j == self.ignore_label for j in range(1, num_classes)]
                scores[:, ignore] = 0.

            # remove low scoring boxes (i.e. low "objectness").
            inds = torch.nonzero(objectness > self.score_thresh).squeeze(1)
            boxes, scores, objectness, keptidxs = boxes[inds], scores[inds], objectness[inds], keptidxs[inds]

            # remove empty boxes
            keep = box_ops.remove_small_boxes(boxes, min_size=1e-2)
            boxes, scores, objectness, keptidxs = boxes[keep], scores[keep], objectness[keep], keptidxs[keep]

            # global nms, regardless of class
            keep = global_nms(boxes=boxes, scores=objectness, iou_threshold=self.nms_thresh)

            # keep only topk scoring predictions
            keep = keep[:self.detections_per_img]
            boxes, scores, objectness, keptidxs = boxes[keep], scores[keep], objectness[keep], keptidxs[keep]

            if boxes.shape[0] > 0:
                # make sure probabilities add to 1, while preserving the
                # ignore_label columns as zero, so we divide by total instead
                # of using softmax
                scores = scores / scores.sum(1)[:, None]

                # label is argmax, just like ordinary classification tasks. Keep
                # in mind that we got rid of the background channel, so add 1
                labels = torch.argmax(scores, dim=1) + 1
            else:
                # empty tensor; nothing was left after filtering out junk
                labels = scores[:, 0].type(torch.int64)

            all_boxes.append(boxes)
            all_objectness.append(objectness)
            all_labels.append(labels)
            all_probabs.append(scores)
            all_keptidxs.append(keptidxs)

        return all_boxes, all_objectness, all_labels, all_probabs, all_keptidxs

    def postprocess_detections(self,
                               class_logits,    # type: Tensor
                               box_regression,  # type: Tensor
                               proposals,       # type: List[Tensor]
                               image_shapes     # type: List[Tuple[int, int]]
                               ):
        # type: (...) -> Tuple[List[Tensor], List[Tensor], List[Tensor]]
        device = class_logits.device
        num_classes = class_logits.shape[-1]

        boxes_per_image = [boxes_in_image.shape[0] for boxes_in_image in proposals]
        pred_boxes = self.box_coder.decode(box_regression, proposals)

        pred_scores = F.softmax(class_logits, -1)

        pred_boxes_list = pred_boxes.split(boxes_per_image, 0)
        pred_scores_list = pred_scores.split(boxes_per_image, 0)

        all_boxes = []
        all_scores = []
        all_labels = []
        for boxes, scores, image_shape in zip(pred_boxes_list, pred_scores_list, image_shapes):
            boxes = box_ops.clip_boxes_to_image(boxes, image_shape)

            # create labels for each prediction
            labels = torch.arange(num_classes, device=device)
            labels = labels.view(1, -1).expand_as(scores)

            # remove predictions with the background label
            boxes = boxes[:, 1:]
            scores = scores[:, 1:]
            labels = labels[:, 1:]

            # Mohamed: remove self.ignore_label predictions
            if self.ignore_label is not None:
                keep = [j != self.ignore_label for j in range(1, num_classes)]
                boxes = boxes[:, keep]
                scores = scores[:, keep]
                labels = labels[:, keep]

            # batch everything, by making every class prediction be a separate instance
            boxes = boxes.reshape(-1, 4)
            scores = scores.reshape(-1)
            labels = labels.reshape(-1)

            # remove low scoring boxes
            inds = torch.nonzero(scores > self.score_thresh).squeeze(1)
            boxes, scores, labels = boxes[inds], scores[inds], labels[inds]

            # remove empty boxes
            keep = box_ops.remove_small_boxes(boxes, min_size=1e-2)
            boxes, scores, labels = boxes[keep], scores[keep], labels[keep]

            # independent nms per class
            keep = box_ops.batched_nms(
                boxes=boxes, scores=scores, idxs=labels,
                iou_threshold=self.nms_thresh)

            # keep only topk scoring predictions
            keep = keep[:self.detections_per_img]
            boxes, scores, labels = boxes[keep], scores[keep], labels[keep]

            all_boxes.append(boxes)
            all_scores.append(scores)
            all_labels.append(labels)

        return all_boxes, all_scores, all_labels

    def get_classification_proposals(self, proposals, image_shapes):

        classification_proposals = []
        for pno, prop in enumerate(proposals):
            xmins, ymins, xmaxs, ymaxs = torch.chunk(prop, 4, dim=1)
            xs = xmins + (xmaxs - xmins) / 2
            ys = ymins + (ymaxs - ymins) / 2
            boxes = torch.cat(
                [xs - self.halfclbox, ys - self.halfclbox,
                 xs + self.halfclbox, ys + self.halfclbox],
                dim=1)

            # adjust boxes
            dim0, dim1 = image_shapes[pno]
            boxes, _ = remove_degenerate_bboxes(
                boxes=boxes, dim0=dim0, dim1=dim1, min_boxside=0)

            classification_proposals.append(boxes)

        return classification_proposals

    def forward(self,
                features,  # type: Dict[str, Tensor]
                proposals,  # type: List[Tensor]
                image_shapes,  # type: List[Tuple[int, int]]
                targets=None,  # type: Optional[List[Dict[str, Tensor]]]
                _just_return_probabs=False,  # type: bool
                _cprobabs=None,  # type: List[Tensor]
                _return_prepr=False,  # type: bool
                ):
        # type: (...) -> Tuple[List[Dict[str, Tensor]], Dict[str, Tensor]]
        """
        Arguments:
            features (List[Tensor])
            proposals (List[Tensor[N, 4]])
            image_shapes (List[Tuple[H, W]])
            targets (List[Dict])
            _just_return_probabs (bool): If true, just returns logits for each
                of the proposals without postprocessing. This is only applies
                if not self.training.
            _cprobabs (List[Tensor]): If given, these "past logits" (obtained
                from variations of the proposals) would be aggregated to
                the class_logits as a form of test-time augmentation. This
                obviously only applies if not self.training
            _return_prepr (bool): also return intermediate representations?
        """
        if targets is not None:
            for t in targets:
                # mTODO: https://github.com/pytorch/pytorch/issues/26731
                floating_point_types = (torch.float, torch.double, torch.half)
                assert t["boxes"].dtype in floating_point_types, 'target boxes must of float type'
                assert t["labels"].dtype == torch.int64, 'target labels must of int64 type'
                if self.has_keypoint():
                    assert t["keypoints"].dtype == torch.float32, 'target keypoints must of float type'

        if self.training:
            proposals, matched_idxs, labels, regression_targets = self.select_training_samples(proposals, targets)
        else:
            labels = None
            regression_targets = None
            matched_idxs = None

        # Added by Mohamed: Classification proposals use an
        # independent (possibly wider) region beyond object boundary
        cproposals = None if not self.indep_classif_boxes else \
            self.get_classification_proposals(proposals=proposals, image_shapes=image_shapes)

        # Mohamed: Independent, extra convolutions for classification
        cfeatures = None if self.cconvhead is None else OrderedDict(
            {k: self.cconvhead(v) for k, v in features.items()}
        )

        # roi pooling for the object bbox
        box_features = self.box_roi_pool(features, proposals, image_shapes)
        if (cproposals is not None) or (cfeatures is not None):
            cbox_features = self.box_roi_pool(
                cfeatures or features, cproposals or proposals, image_shapes)
        else:
            cbox_features = None

        # flatten
        box_features = self.box_head(box_features)
        cbox_features = self.box_head(cbox_features) \
            if cbox_features is not None else None

        # Mohamed: per-fov nuclei attentive to each other
        if self.sattention_head is not None:

            # this only applies for classification
            if cbox_features is None:
                cbox_features = 0. + box_features

            # get the start index for each group of proposals from one FOV
            fov_start_idxs = np.cumsum(
                [0] + [j.shape[0] for j in proposals[:-1]]).tolist()

            # pass through self-attention head
            cbox_features = self.sattention_head(
                cbox_features, fov_start_idxs=fov_start_idxs)

        # pass through fully-connected layers
        if cbox_features is None:
            class_logits, box_regression = self.box_predictor(box_features)
        else:
            _, box_regression = self.box_predictor(box_features, get_scores=False)
            class_logits, _ = self.box_predictor(cbox_features, get_deltas=False)

        result = torch.jit.annotate(List[Dict[str, torch.Tensor]], [])
        losses = {}
        if self.training:
            assert labels is not None and regression_targets is not None
            loss_classifier, loss_box_reg = fastrcnn_loss(
                class_logits=class_logits,
                box_regression=box_regression,
                labels=labels,
                regression_targets=regression_targets,
                ignore_label=self.ignore_label,
                batched_nms=self.batched_nms,
            )
            losses = {
                "loss_classifier": loss_classifier,
                "loss_box_reg": loss_box_reg
            }
        else:
            # Mohamed: maybe just return class probabs without postprocessing
            # Note that this INCLUDES the background class
            if _just_return_probabs:
                return F.softmax(class_logits, -1)

            # Mohamed: "past logits" (obtained from variations of proposals)
            # would be aggregated to the class_logits as a form of test-time
            # augmentation. This obviously only applies if not self.training
            if _cprobabs is not None:
                class_logits = F.softmax(class_logits, -1)
                class_logits += _cprobabs

            if self.batched_nms:
                boxes, scores, labels = self.postprocess_detections(class_logits, box_regression, proposals, image_shapes)
                probabs = [None] * len(boxes)
                keptidxs = [None] * len(boxes)
            else:
                boxes, scores, labels, probabs, keptidxs = self.postprocess_detections_globalnms(class_logits, box_regression, proposals, image_shapes)
            num_images = len(boxes)

            if _return_prepr:
                bperim = [boxes_in_image.shape[0] for boxes_in_image in proposals]
                bfeats = box_features.split(bperim, 0)
                cbfeats = None if cbox_features is None else cbox_features.split(bperim, 0)
                clogits = class_logits.split(bperim, 0)

            for i in range(num_images):
                out = {
                    "boxes": boxes[i],
                    "labels": labels[i],
                    "scores": scores[i],
                    "probabs": probabs[i],
                }
                if _return_prepr:
                    out.update({
                        "box_features": bfeats[i][keptidxs[i], :],
                        "cbox_features": None if cbfeats is None else cbfeats[i][keptidxs[i], :],
                        "clogits": clogits[i][keptidxs[i], :],
                    })
                result.append(out)

        if self.has_mask():
            mask_proposals = [p["boxes"] for p in result]
            if self.training:
                assert matched_idxs is not None
                # during training, only focus on positive boxes
                num_images = len(proposals)
                mask_proposals = []
                pos_matched_idxs = []
                for img_id in range(num_images):
                    pos = torch.nonzero(labels[img_id] > 0).squeeze(1)
                    mask_proposals.append(proposals[img_id][pos])
                    pos_matched_idxs.append(matched_idxs[img_id][pos])
            else:
                pos_matched_idxs = None

            if self.mask_roi_pool is not None:
                mask_features = self.mask_roi_pool(features, mask_proposals, image_shapes)
                mask_features = self.mask_head(mask_features)
                mask_logits = self.mask_predictor(mask_features)
            else:
                # noinspection PyUnusedLocal
                mask_logits = torch.tensor(0)
                raise Exception("Expected mask_roi_pool to be not None")

            loss_mask = {}
            if self.training:
                assert targets is not None
                assert pos_matched_idxs is not None
                assert mask_logits is not None

                # Mohamed: if not self.batched_nms, just one mask per object
                gt_labels = [0 + t["labels"] for t in targets]
                if not self.batched_nms:
                    for labs in gt_labels:
                        labs[labs > 1] = 1

                gt_masks = [t["masks"] for t in targets]
                # Mohamed: added to ignore bboxes (false masks)
                #  so that they don't contribute to the mask loss
                if 'ismask' in targets[0]:
                    gt_ismask = [t['ismask'] for t in targets]
                else:
                    gt_ismask = None
                rcnn_loss_mask = maskrcnn_loss(
                    mask_logits=mask_logits, proposals=mask_proposals,
                    gt_masks=gt_masks, gt_labels=gt_labels,
                    mask_matched_idxs=pos_matched_idxs,
                    gt_ismask=gt_ismask)
                loss_mask = {
                    "loss_mask": rcnn_loss_mask
                }
            else:
                # Mohamed: if not self.batched_nms, just one mask per object
                mask_labels = [0 + r["labels"] for r in result]
                if not self.batched_nms:
                    for labs in mask_labels:
                        labs[labs > 1] = 1
                masks_probs = maskrcnn_inference(mask_logits, mask_labels)
                for mask_prob, r in zip(masks_probs, result):
                    r["masks"] = mask_prob

            losses.update(loss_mask)

        # keep none checks in if conditional so torchscript will conditionally
        # compile each branch
        if self.keypoint_roi_pool is not None and self.keypoint_head is not None \
                and self.keypoint_predictor is not None:
            keypoint_proposals = [p["boxes"] for p in result]
            if self.training:
                # during training, only focus on positive boxes
                num_images = len(proposals)
                keypoint_proposals = []
                pos_matched_idxs = []
                assert matched_idxs is not None
                for img_id in range(num_images):
                    pos = torch.nonzero(labels[img_id] > 0).squeeze(1)
                    keypoint_proposals.append(proposals[img_id][pos])
                    pos_matched_idxs.append(matched_idxs[img_id][pos])
            else:
                pos_matched_idxs = None

            keypoint_features = self.keypoint_roi_pool(features, keypoint_proposals, image_shapes)
            keypoint_features = self.keypoint_head(keypoint_features)
            keypoint_logits = self.keypoint_predictor(keypoint_features)

            loss_keypoint = {}
            if self.training:
                assert targets is not None
                assert pos_matched_idxs is not None

                gt_keypoints = [t["keypoints"] for t in targets]
                rcnn_loss_keypoint = keypointrcnn_loss(
                    keypoint_logits, keypoint_proposals,
                    gt_keypoints, pos_matched_idxs)
                loss_keypoint = {
                    "loss_keypoint": rcnn_loss_keypoint
                }
            else:
                assert keypoint_logits is not None
                assert keypoint_proposals is not None

                keypoints_probs, kp_scores = keypointrcnn_inference(keypoint_logits, keypoint_proposals)
                for keypoint_prob, kps, r in zip(keypoints_probs, kp_scores, result):
                    r["keypoints"] = keypoint_prob
                    r["keypoints_scores"] = kps

            losses.update(loss_keypoint)

        return result, losses


"""A module containing the core message-passing functions for belief propagation"""

import functools
import inspect
from dataclasses import dataclass
from typing import Any, Callable, Dict, Hashable, Mapping, Optional, Tuple

import jax
import jax.numpy as jnp
import numpy as np
from jax.scipy.special import logsumexp

from pgmax.factor import FAC_TO_VAR_UPDATES

from . import bp_state as bpstate
from . import bp_utils
from .bp_state import BPArrays, BPState, Evidence, FToVMessages, LogPotentials

@dataclass(frozen=True, eq=False)
class BeliefPropagation:
    """Belief propagation functions.

    Arguments:
        init: Function to create log_potentials, ftov_msgs and evidence.
            Args:
                log_potentials_updates: Optional dictionary containing log_potentials updates.
                ftov_msgs_updates: Optional dictionary containing ftov_msgs updates.
                evidence_updates: Optional dictionary containing evidence updates.
            Returns:
                A BPArrays with the log_potentials, ftov_msgs and evidence.

        update: Function to update log_potentials, ftov_msgs and evidence.
            Args:
                bp_arrays: Optional arrays of log_potentials, ftov_msgs, evidence.
                log_potentials_updates: Optional dictionary containing log_potentials updates.
                ftov_msgs_updates: Optional dictionary containing ftov_msgs updates.
                evidence_updates: Optional dictionary containing evidence updates.
            Returns:
                A BPArrays with the updated log_potentials, ftov_msgs and evidence.

        run_bp: Function to run belief propagation for num_iters with a damping_factor.
            Args:
                bp_arrays: Initial arrays of log_potentials, ftov_msgs, evidence.
                num_iters: Number of belief propagation iterations.
                damping: The damping factor to use for message updates between one timestep and the next.
            Returns:
                A BPArrays containing the updated ftov_msgs.

        get_bp_state: Function to reconstruct the BPState from a BPArrays.
            Args:
                bp_arrays: A BPArrays containing log_potentials, ftov_msgs, evidence.
            Returns:
                The reconstructed BPState

        get_beliefs: Function to calculate beliefs from a BPArrays.
            Args:
                bp_arrays: A BPArrays containing log_potentials, ftov_msgs, evidence.
            Returns:
                beliefs: Beliefs returned by belief propagation.
    """

    init: Callable
    update: Callable
    run_bp: Callable
    to_bp_state: Callable
    get_beliefs: Callable

def BP(bp_state: BPState, temperature: float = 0.0) -> BeliefPropagation:
    """Function for generating belief propagation functions.

    Args:
        bp_state: Belief propagation state.
        temperature: Temperature for loopy belief propagation.
            1.0 corresponds to sum-product, 0.0 corresponds to max-product.

    Returns:
        Belief propagation functions.
    """
    wiring = bp_state.fg_state.wiring
    edges_num_states = np.concatenate(
        [wiring[factor_type].edges_num_states for factor_type in FAC_TO_VAR_UPDATES]
    )
    max_msg_size = int(np.max(edges_num_states))

    var_states_for_edges = np.concatenate(
        [wiring[factor_type].var_states_for_edges for factor_type in FAC_TO_VAR_UPDATES]
    )

    # Inference argumnets per factor type
    inference_arguments: Dict[type, Mapping] = {}
    for factor_type in FAC_TO_VAR_UPDATES:
        this_inference_arguments = inspect.getfullargspec(
            FAC_TO_VAR_UPDATES[factor_type]
        ).args
        this_inference_arguments.remove("vtof_msgs")
        this_inference_arguments.remove("log_potentials")
        this_inference_arguments.remove("temperature")
        this_inference_arguments = {
            key: getattr(wiring[factor_type], key) for key in this_inference_arguments
        }
        inference_arguments[factor_type] = this_inference_arguments

    factor_type_to_msgs_range = bp_state.fg_state.factor_type_to_msgs_range
    factor_type_to_potentials_range = bp_state.fg_state.factor_type_to_potentials_range

    def update(
        bp_arrays: Optional[BPArrays] = None,
        log_potentials_updates: Optional[Dict[Any, jnp.ndarray]] = None,
        ftov_msgs_updates: Optional[Dict[Any, jnp.ndarray]] = None,
        evidence_updates: Optional[Dict[Any, jnp.ndarray]] = None,
    ) -> BPArrays:
        """Function to update belief propagation log_potentials, ftov_msgs, evidence.

        Args:
            bp_arrays: Optional arrays of log_potentials, ftov_msgs, evidence.
            log_potentials_updates: Optional dictionary containing log_potentials updates.
            ftov_msgs_updates: Optional dictionary containing ftov_msgs updates.
            evidence_updates: Optional dictionary containing evidence updates.

        Returns:
            A BPArrays with the updated log_potentials, ftov_msgs and evidence.
        """
        if bp_arrays is not None:
            log_potentials = bp_arrays.log_potentials
            evidence = bp_arrays.evidence
            ftov_msgs = bp_arrays.ftov_msgs
        else:
            log_potentials = jax.device_put(bp_state.log_potentials.value)
            ftov_msgs = bp_state.ftov_msgs.value
            evidence = bp_state.evidence.value

        if log_potentials_updates is not None:
            log_potentials = bpstate.update_log_potentials(
                log_potentials, log_potentials_updates, bp_state.fg_state
            )

        if ftov_msgs_updates is not None:
            ftov_msgs = bpstate.update_ftov_msgs(
                ftov_msgs, ftov_msgs_updates, bp_state.fg_state
            )

        if evidence_updates is not None:
            evidence = bpstate.update_evidence(
                evidence, evidence_updates, bp_state.fg_state
            )

        return BPArrays(
            log_potentials=log_potentials, ftov_msgs=ftov_msgs, evidence=evidence
        )

    def run_bp(
        bp_arrays: BPArrays,
        num_iters: int,
        damping: float = 0.5,
    ) -> BPArrays:
        """Function to run belief propagation for num_iters with a damping_factor.

        Args:
            bp_arrays: Initial arrays of log_potentials, ftov_msgs, evidence.
            num_iters: Number of belief propagation iterations.
            damping: The damping factor to use for message updates between one timestep and the next.

        Returns:
            A BPArrays containing the updated ftov_msgs.
        """
        log_potentials = bp_arrays.log_potentials
        evidence = bp_arrays.evidence
        ftov_msgs = bp_arrays.ftov_msgs

        # Normalize the messages to ensure the maximum value is 0.
        ftov_msgs = normalize_and_clip_msgs(ftov_msgs, edges_num_states, max_msg_size)

        @jax.checkpoint
        def update(msgs: jnp.ndarray, _) -> Tuple[jnp.ndarray, None]:
            # Compute new variable to factor messages by message passing
            vtof_msgs = pass_var_to_fac_messages(
                msgs,
                evidence,
                var_states_for_edges,
            )
            ftov_msgs = jnp.zeros_like(vtof_msgs)
            for factor_type in FAC_TO_VAR_UPDATES:
                msgs_start, msgs_end = factor_type_to_msgs_range[factor_type]
                potentials_start, potentials_end = factor_type_to_potentials_range[
                    factor_type
                ]
                ftov_msgs_type = FAC_TO_VAR_UPDATES[factor_type](
                    vtof_msgs=vtof_msgs[msgs_start:msgs_end],
                    log_potentials=log_potentials[potentials_start:potentials_end],
                    temperature=temperature,
                    **inference_arguments[factor_type],
                )
                ftov_msgs = ftov_msgs.at[msgs_start:msgs_end].set(ftov_msgs_type)

            # Use the results of message passing to perform damping and
            # update the factor to variable messages
            delta_msgs = ftov_msgs - msgs
            msgs = msgs + (1 - damping) * delta_msgs
            # Normalize and clip these damped, updated messages before returning them.
            msgs = normalize_and_clip_msgs(msgs, edges_num_states, max_msg_size)
            return msgs, None

        ftov_msgs, _ = jax.lax.scan(update, ftov_msgs, None, num_iters)

        return BPArrays(
            log_potentials=log_potentials, ftov_msgs=ftov_msgs, evidence=evidence
        )

    def to_bp_state(bp_arrays: BPArrays) -> BPState:
        """Function to reconstruct the BPState from a BPArrays

        Args:
            bp_arrays: A BPArrays containing log_potentials, ftov_msgs, evidence.

        Returns:
            The reconstructed BPState
        """
        return BPState(
            log_potentials=LogPotentials(
                fg_state=bp_state.fg_state, value=bp_arrays.log_potentials
            ),
            ftov_msgs=FToVMessages(
                fg_state=bp_state.fg_state,
                value=bp_arrays.ftov_msgs,
            ),
            evidence=Evidence(fg_state=bp_state.fg_state, value=bp_arrays.evidence),
        )

    def unflatten_beliefs(flat_beliefs, variable_groups) -> Dict[Hashable, Any]:
        """Function that returns unflattened beliefs from the flat beliefs

        Args:
            flat_beliefs: Flattened array of beliefs
            variable_groups: All the variable groups in the FactorGraph.
        """
        beliefs = {}
        start = 0
        for variable_group in variable_groups:
            num_states = variable_group.num_states
            assert isinstance(num_states, np.ndarray)
            length = num_states.sum()

            beliefs[variable_group] = variable_group.unflatten(
                flat_beliefs[start : start + length]
            )
            start += length
        return beliefs

    @jax.jit
    def get_beliefs(bp_arrays: BPArrays) -> Dict[Hashable, Any]:
        """Function to calculate beliefs from a BPArrays

        Args:
            bp_arrays: A BPArrays containing log_potentials, ftov_msgs, evidence.

        Returns:
            beliefs: Beliefs returned by belief propagation.
        """

        flat_beliefs = (
            jax.device_put(bp_arrays.evidence)
            .at[jax.device_put(var_states_for_edges)]
            .add(bp_arrays.ftov_msgs)
        )
        return unflatten_beliefs(flat_beliefs, bp_state.fg_state.variable_groups)

    bp = BeliefPropagation(
        init=functools.partial(update, None),
        update=update,
        run_bp=run_bp,
        to_bp_state=to_bp_state,
        get_beliefs=get_beliefs,
    )
    return bp

@jax.jit
def pass_var_to_fac_messages(
    ftov_msgs: jnp.array,
    evidence: jnp.array,
    var_states_for_edges: jnp.array,
) -> jnp.array:
    """Passes messages from Variables to Factors.

    The update works by first summing the evidence and neighboring factor to variable messages for
    each variable. Next, it subtracts messages from the correct elements of this sum to yield the
    correct updated messages.

    Args:
        ftov_msgs: Array of shape (num_edge_state,). This holds all the flattened factor to variable
            messages.
        evidence: Array of shape (num_var_states,) representing the flattened evidence for each variable
        var_states_for_edges: Array of shape (num_edge_states,)
            Global variable state indices for each edge state
    Returns:
        Array of shape (num_edge_state,). This holds all the flattened variable to factor messages.
    """
    var_sums_arr = evidence.at[var_states_for_edges].add(ftov_msgs)
    vtof_msgs = var_sums_arr[var_states_for_edges] - ftov_msgs
    return vtof_msgs

@functools.partial(jax.jit, static_argnames=("max_msg_size"))
def normalize_and_clip_msgs(
    msgs: jnp.ndarray,
    edges_num_states: jnp.ndarray,
    max_msg_size: int,
) -> jnp.ndarray:
    """Performs normalization and clipping of flattened messages

    Normalization is done by subtracting the maximum value of every message from every element of every message,
    clipping is done to keep every message value in the range [-1000, 0].

    Args:
        msgs: Array of shape (num_edge_state,). This holds all the flattened factor to variable messages.
        edges_num_states: Array of shape (num_edges,). Number of states for the variables connected to each edge
        max_msg_size: the max of edges_num_states

    Returns:
        Array of shape (num_edge_state,). This holds all the flattened factor to variable messages
            after normalization and clipping
    """
    msgs = msgs - jnp.repeat(
        bp_utils.segment_max_opt(msgs, edges_num_states, max_msg_size),
        edges_num_states,
        total_repeat_length=msgs.shape[0],
    )
    # Clip message values to be always greater than -1000
    msgs = jnp.clip(msgs, -1000, None)
    return msgs

@jax.jit
def decode_map_states(beliefs: Dict[Hashable, Any]) -> Any:
    """Function to decode MAP states given the calculated beliefs.

    Args:
        beliefs: An array or a PyTree container containing beliefs for different variables.

    Returns:
        An array or a PyTree container containing the MAP states for different variables.
    """
    return jax.tree_util.tree_map(lambda x: jnp.argmax(x, axis=-1), beliefs)

@jax.jit
def get_marginals(beliefs: Dict[Hashable, Any]) -> Any:
    """Function to get marginal probabilities given the calculated beliefs.

    Args:
        beliefs: An array or a PyTree container containing beliefs for different variables.

    Returns:
        An array or a PyTree container containing the marginal probabilities different variables.
    """
    return jax.tree_util.tree_map(
        lambda x: jnp.exp(x - logsumexp(x, axis=-1, keepdims=True)), beliefs
    )


# coding: utf-8
# Copyright (c) Pymatgen Development Team.
# Distributed under the terms of the MIT License.

"""
This module provides core classes needed by all define electronic structure,
such as the Spin, Orbital, etc.
"""

from monty.json import MSONable
from enum import Enum, unique
import numpy as np

__author__ = "Shyue Ping Ong"
__copyright__ = "Copyright 2011, The Materials Project"
__version__ = "1.0"
__maintainer__ = "Shyue Ping Ong"
__email__ = "shyuep@gmail.com"
__status__ = "Production"
__date__ = "Sep 23, 2011"

@unique
class Spin(Enum):
    """
    Enum type for Spin.  Only up and down.
    Usage: Spin.up, Spin.down.
    """
    up, down = (1, -1)

    def __int__(self):
        return self.value

    def __float__(self):
        return float(self.value)

    def __str__(self):
        return str(self.value)

@unique
class OrbitalType(Enum):
    """
    Enum type for orbital type. Indices are basically the azimuthal quantum
    number, l.
    """

    s = 0
    p = 1
    d = 2
    f = 3

    def __str__(self):
        return self.name

@unique
class Orbital(Enum):
    """
    Enum type for specific orbitals. The indices are basically the order in
    which the orbitals are reported in VASP and has no special meaning.
    """

    s = 0
    py = 1
    pz = 2
    px = 3
    dxy = 4
    dyz = 5
    dz2 = 6
    dxz = 7
    dx2 = 8
    f_3 = 9
    f_2 = 10
    f_1 = 11
    f0 = 12
    f1 = 13
    f2 = 14
    f3 = 15

    def __int__(self):
        return self.value

    def __str__(self):
        return self.name

    @property
    def orbital_type(self):
        """
        Returns OrbitalType of an orbital.
        """
        return OrbitalType[self.name[0]]

class Magmom(MSONable):
    """
    New class in active development. Use with caution, feedback is
    appreciated.

    Class to handle magnetic moments. Defines the magnetic moment of a
    site or species relative to a spin quantization axis. Designed for
    use in electronic structure calculations.

    * For the general case, Magmom can be specified by a vector,
      e.g. m = Magmom([1.0, 1.0, 2.0]), and subscripts will work as
      expected, e.g. m[0] gives 1.0

    * For collinear calculations, Magmom can assumed to be scalar-like,
      e.g. m = Magmom(5.0) will work as expected, e.g. float(m) gives 5.0

    Both of these cases should be safe and shouldn't give any surprises,
    but more advanced functionality is available if required.

    There also exist useful static methods for lists of magmoms:

    * Magmom.are_collinear(magmoms) - if true, a collinear electronic
      structure calculation can be safely initialized, with float(Magmom)
      giving the expected scalar magnetic moment value

    * Magmom.get_consistent_set_and_saxis(magmoms) - for non-collinear
      electronic structure calculations, a global, consistent spin axis
      has to be used. This method returns a list of Magmoms which all
      share a common spin axis, along with the global spin axis.

    All methods that take lists of magmoms will accept magmoms either as
    Magmom objects or as scalars/lists and will automatically convert to
    a Magmom representation internally.

    The following methods are also particularly useful in the context of
    VASP calculations:

    * Magmom.get_xyz_magmom_with_001_saxis()
    * Magmom.get_00t_magmom_with_xyz_saxis()

    See VASP documentation for more information:

    https://cms.mpi.univie.ac.at/wiki/index.php/SAXIS
    """

    def __init__(self, moment, saxis=(0, 0, 1)):
        """
        :param moment: magnetic moment, supplied as float or list/np.ndarray
        :param saxis: spin axis, supplied as list/np.ndarray, parameter will
            be converted to unit vector (default is [0, 0, 1])
        :return: Magmom object
        """
        # to init from another Magmom instance
        if isinstance(moment, Magmom):
            saxis = moment.saxis
            moment = moment.moment

        moment = np.array(moment, dtype='d')
        if moment.ndim == 0:
            moment = moment * [0, 0, 1]

        self.moment = moment

        saxis = np.array(saxis, dtype='d')

        self.saxis = saxis / np.linalg.norm(saxis)

    @classmethod
    def from_global_moment_and_saxis(cls, global_moment, saxis):
        """
        Convenience method to initialize Magmom from a given global
        magnetic moment, i.e. magnetic moment with saxis=(0,0,1), and
        provided saxis.

        Method is useful if you do not know the components of your
        magnetic moment in frame of your desired saxis.

        :param global_moment:
        :param saxis: desired saxis
        :return:
        """
        magmom = Magmom(global_moment)
        return cls(magmom.get_moment(saxis=saxis), saxis=saxis)

    def _get_transformation_matrix(self, saxis):

        saxis = saxis / np.linalg.norm(saxis)

        alpha = np.arctan2(saxis[1], saxis[0])
        beta = np.arctan2(np.sqrt(saxis[0] ** 2 + saxis[1] ** 2), saxis[2])

        cos_a = np.cos(alpha)
        cos_b = np.cos(beta)
        sin_a = np.sin(alpha)
        sin_b = np.sin(beta)

        m = [[cos_b * cos_a, -sin_a, sin_b * cos_a],
             [cos_b * sin_a, cos_a, sin_b * sin_a],
             [-sin_b, 0, cos_b]]

        return m

    def _get_transformation_matrix_inv(self, saxis):

        saxis = saxis / np.linalg.norm(saxis)

        alpha = np.arctan2(saxis[1], saxis[0])
        beta = np.arctan2(np.sqrt(saxis[0] ** 2 + saxis[1] ** 2), saxis[2])

        cos_a = np.cos(alpha)
        cos_b = np.cos(beta)
        sin_a = np.sin(alpha)
        sin_b = np.sin(beta)

        m = [[cos_b * cos_a, cos_b * sin_a, -sin_b],
             [-sin_a, cos_a, 0],
             [sin_b * cos_a, sin_b * sin_a, cos_b]]

        return m

    def get_moment(self, saxis=(0, 0, 1)):
        """
        Get magnetic moment relative to a given spin quantization axis.
        If no axis is provided, moment will be given relative to the
        Magmom's internal spin quantization axis, i.e. equivalent to
        Magmom.moment

        :param axis: (list/numpy array) spin quantization axis
        :return: np.ndarray of length 3
        """

        # transform back to moment with spin axis [0, 0, 1]
        m_inv = self._get_transformation_matrix_inv(self.saxis)
        moment = np.matmul(self.moment, m_inv)

        # transform to new saxis
        m = self._get_transformation_matrix(saxis)
        moment = np.matmul(moment, m)

        # round small values to zero
        moment[np.abs(moment) < 1e-8] = 0

        return moment

    @property
    def global_moment(self):
        """
        Get the magnetic moment defined in an arbitrary global reference frame.

        :return: np.ndarray of length 3
        """
        return self.get_moment()

    @property
    def projection(self):
        """
        Projects moment along spin quantisation axis. Useful for obtaining
        collinear approximation for slightly non-collinear magmoms.

        :return: float
        """
        return np.dot(self.moment, self.saxis)

    def get_xyz_magmom_with_001_saxis(self):
        """
        Returns a Magmom in the default setting of saxis = [0, 0, 1] and
        the magnetic moment rotated as required.

        :return: Magmom
        """
        return Magmom(self.get_moment())

    def get_00t_magmom_with_xyz_saxis(self):
        """
        For internal implementation reasons, in non-collinear calculations
        VASP prefers:

        MAGMOM = 0 0 total_magnetic_moment
        SAXIS = x y z

        to an equivalent:

        MAGMOM = x y z
        SAXIS = 0 0 1

        This method returns a Magmom object with magnetic moment [0, 0, t],
        where t is the total magnetic moment, and saxis rotated as required.

        A consistent direction of saxis is applied such that t might be positive
        or negative depending on the direction of the initial moment. This is useful
        in the case of collinear structures, rather than constraining assuming
        t is always positive.

        :return: Magmom
        """
        # reference direction gives sign of moment
        # entirely arbitrary, there will always be a pathological case
        # where a consistent sign is not possible if the magnetic moments
        # are aligned along the reference direction, but in practice this
        # is unlikely to happen
        ref_direction = np.array([1.01, 1.02, 1.03])
        t = abs(self)
        if t != 0:
            new_saxis = self.moment / np.linalg.norm(self.moment)
            if np.dot(ref_direction, new_saxis) < 0:
                t = -t
                new_saxis = -new_saxis
            return Magmom([0, 0, t], saxis=new_saxis)
        else:
            return Magmom(self)

    @staticmethod
    def have_consistent_saxis(magmoms):
        """
        This method checks that all Magmom objects in a list have a
        consistent spin quantization axis. To write MAGMOM tags to a
        VASP INCAR, a global SAXIS value for all magmoms has to be used.
        If saxis are inconsistent, can create consistent set with:
        Magmom.get_consistent_set(magmoms)

        :param magmoms: list of magmoms (Magmoms, scalars or vectors)
        :return: bool
        """
        magmoms = [Magmom(magmom) for magmom in magmoms]
        ref_saxis = magmoms[0].saxis
        match_ref = [magmom.saxis == ref_saxis for magmom in magmoms]
        if np.all(match_ref):
            return True
        else:
            return False

    @staticmethod
    def get_consistent_set_and_saxis(magmoms, saxis=None):
        """
        Method to ensure a list of magmoms use the same spin axis.
        Returns a tuple of a list of Magmoms and their global spin axis.

        :param magmoms: list of magmoms (Magmoms, scalars or vectors)
        :param saxis: can provide a specific global spin axis
        :return: (list of Magmoms, global spin axis) tuple
        """
        magmoms = [Magmom(magmom) for magmom in magmoms]
        if saxis is None:
            saxis = Magmom.get_suggested_saxis(magmoms)
        else:
            saxis = saxis / np.linalg.norm(saxis)
        magmoms = [magmom.get_moment(saxis=saxis) for magmom in magmoms]
        return (magmoms, saxis)

    @staticmethod
    def get_suggested_saxis(magmoms):
        """
        This method returns a suggested spin axis for a set of magmoms,
        taking the largest magnetic moment as the reference. For calculations
        with collinear spins, this would give a sensible saxis for a ncl
        calculation.

        :param magmoms: list of magmoms (Magmoms, scalars or vectors)
        :return: np.ndarray of length 3
        """
        # heuristic, will pick largest magmom as reference
        # useful for creating collinear approximations of
        # e.g. slightly canted magnetic structures
        # for fully collinear structures, will return expected
        # result

        magmoms = [Magmom(magmom) for magmom in magmoms]
        # filter only non-zero magmoms
        magmoms = [magmom for magmom in magmoms if abs(magmom)]
        magmoms.sort(reverse=True)
        if len(magmoms) > 0:
            return magmoms[0].get_00t_magmom_with_xyz_saxis().saxis
        else:
            return np.array([0, 0, 1], dtype="d")

    @staticmethod
    def are_collinear(magmoms):
        """
        Method checks to see if a set of magnetic moments are collinear
        with each other.
        :param magmoms: list of magmoms (Magmoms, scalars or vectors)
        :return: bool
        """
        magmoms = [Magmom(magmom) for magmom in magmoms]
        if not Magmom.have_consistent_saxis(magmoms):
            magmoms = Magmom.get_consistent_set(magmoms)

        # convert to numpy array for convenience
        magmoms = np.array([list(magmom) for magmom in magmoms])
        magmoms = magmoms[np.any(magmoms, axis=1)]  # remove zero magmoms
        if len(magmoms) == 0:
            return True

        # use first moment as reference to compare against
        ref_magmom = magmoms[0]
        # magnitude of cross products != 0 if non-collinear with reference
        num_ncl = np.count_nonzero(np.linalg.norm(np.cross(ref_magmom, magmoms), axis=1))
        if num_ncl > 0:
            return False
        else:
            return True

    @classmethod
    def from_moment_relative_to_crystal_axes(cls, moment, lattice):
        """
        Obtaining a Magmom object from a magnetic moment provided
        relative to crystal axes.

        Used for obtaining moments from magCIF file.
        :param magmom: list of floats specifying vector magmom
        :param lattice: Lattice
        :return: Magmom
        """
        # get matrix representing unit lattice vectors
        unit_m = lattice.matrix / np.linalg.norm(lattice.matrix, axis=1)[:, None]
        moment = np.matmul(list(moment), unit_m)
        # round small values to zero
        moment[np.abs(moment) < 1e-8] = 0
        return cls(moment)

    def get_moment_relative_to_crystal_axes(self, lattice):
        """
        If scalar magmoms, moments will be given arbitrarily along z.
        Used for writing moments to magCIF file.

        :param magmom: Magmom
        :param lattice: Lattice
        :return: vector as list of floats
        """
        # get matrix representing unit lattice vectors
        unit_m = lattice.matrix / np.linalg.norm(lattice.matrix, axis=1)[:, None]
        # note np.matmul() requires numpy version >= 1.10
        moment = np.matmul(self.global_moment, np.linalg.inv(unit_m))
        # round small values to zero
        moment[np.abs(moment) < 1e-8] = 0
        return moment

    def __getitem__(self, key):
        return self.moment[key]

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

    def __abs__(self):
        return np.linalg.norm(self.moment)

    def __eq__(self, other):
        """
        Equal if 'global' magnetic moments are the same, saxis can differ.
        """
        other = Magmom(other)
        return np.allclose(self.global_moment, other.global_moment)

    def __ne__(self, other):
        return not self.__eq__(other)

    def __lt__(self, other):
        return abs(self) < abs(other)

    def __neg__(self):
        return Magmom(-self.moment, saxis=self.saxis)

    def __hash__(self):
        return (tuple(self.moment) + tuple(self.saxis)).__hash__()

    def __float__(self):
        """
        Returns magnitude of magnetic moment with a sign with respect to
        an arbitrary direction.

        Should give unsurprising output if Magmom is treated like a
        scalar or if a set of Magmoms describes a collinear structure.

        Implemented this way rather than simpler abs(self) so that
        moments will have a consistent sign in case of e.g.
        antiferromagnetic collinear structures without additional
        user intervention.

        However, should be used with caution for non-collinear
        structures and might give non-sensical results except in the case
        of only slightly non-collinear structures (e.g. small canting).

        This approach is also used to obtain "diff" VolumetricDensity
        in pymatgen.io.vasp.outputs.VolumetricDensity when processing
        Chgcars from SOC calculations.
        """
        return float(self.get_00t_magmom_with_xyz_saxis()[2])

    def __str__(self):
        return str(float(self))

    def __repr__(self):
        if np.allclose(self.saxis, (0, 0, 1)):
            return 'Magnetic moment {0}'.format(self.moment)
        else:
            return 'Magnetic moment {0} (spin axis = {1})'.format(self.moment,
                                                                  self.saxis)


"""
NAME:
	plotting_functions.py

AUTHOR:
	Bronwyn Reichardt Chu
	Swinburne
	2021

EMAIL:
	<breichardtchu@swin.edu.au>

PURPOSE:
	To hold useful plotting functions
	Written on MacOS Mojave 10.14.5, with Python 3.7

FUNCTIONS INLCUDED:
    get_rc_params
    chen_et_al_2010
    murray_et_al_2011
    davies_et_al_2019
    kim_et_al_2020
    fitting_function
    running_mean
    lower_quantile
    upper_quantile
    binned_median_quantile_log
    binned_median_quantile_lin
    pearson_correlation
    read_in_create_wcs
    plot_continuum_contours

MODIFICATION HISTORY:
		v.1.0 - first created January 2021

"""
import numpy as np
import scipy.stats as stats

from astropy.io import fits
from astropy.wcs import WCS

#===============================================================================
# DEFINE PLOTTING PARAMETERS
#===============================================================================

def get_rc_params():
    """
    Define the rcParams that will be used in all the plots.

    Returns
    -------
    rc_params dictionary object
    """

    rc_params = {
        "text.usetex": False,
        "axes.facecolor": 'white',

        #"figure.dpi": 125,
        #"legend.fontsize": 12,
        "legend.frameon": False,
        #"legend.markerscale": 1.0,

        "axes.labelsize": 'large',

        "xtick.direction": 'in',
        "xtick.labelsize": 'medium',
        "xtick.minor.visible": True,
        "xtick.top": True,
        "xtick.major.width": 1,

        "ytick.direction": 'in',
        "ytick.labelsize": 'medium',
        "ytick.minor.visible": True,
        "ytick.right": True,
        "ytick.major.width": 1,
    }

    return rc_params

#===============================================================================
# RELATIONS FROM OTHER PAPERS
#===============================================================================

def chen_et_al_2010(sfr_surface_density_min, sfr_surface_density_max, scale_factor=1):
    """
    The trendline from Chen et al. (2010) where v_out is proportional to (SFR surface density)^0.1
    (Energy driven winds - SNe feedback)

    Parameters
    ----------
    sfr_surface_density_min : float
        The minimum value of the SFR surface density

    sfr_surface_density_max : float
        The maximum value of the SFR surface density

    scale_factor : float
        The number by which to scale the trend, can be used to bring the trend
        into the range of the data on the plot (Default = 1)

    Returns
    -------
    sfr_surface_density : :obj:'~numpy.ndarray'
        vector of SFR surface densities

    v_out : :obj:'~numpy.ndarray'
        vector of outflow velocities following the trend
    """
    #create a vector for sfr surface density
    sfr_surface_density = np.linspace(sfr_surface_density_min, sfr_surface_density_max+4, num=1000)

    #use the relationship to predict the v_out
    v_out = scale_factor*sfr_surface_density**0.1

    return sfr_surface_density, v_out

def murray_et_al_2011(sfr_surface_density_min, sfr_surface_density_max, scale_factor=1):
    """
    The trendline from Murray et al. (2011) where v_out is proportional to (SFR surface density)^2
    (Momentum driven winds - radiative feedback from young stars)

    Parameters
    ----------
    sfr_surface_density_min : float
        The minimum value of the SFR surface density

    sfr_surface_density_max : float
        The maximum value of the SFR surface density

    scale_factor : float
        The number by which to scale the trend, can be used to bring the trend
        into the range of the data on the plot (Default = 1)

    Returns
    -------
    sfr_surface_density : :obj:'~numpy.ndarray'
        vector of SFR surface densities

    v_out : :obj:'~numpy.ndarray'
        vector of outflow velocities following the trend
    """
    #create a vector for sfr surface density
    sfr_surface_density = np.linspace(sfr_surface_density_min, sfr_surface_density_max+4, num=1000)

    #use the relationship to predict the v_out
    v_out = scale_factor*sfr_surface_density**2

    return sfr_surface_density, v_out

def davies_et_al_2019(sfr_surface_density_min, sfr_surface_density_max):
    """
    The trendline from Davies et al. (2019) where the flow velocity dispersion
    is proportional to SFR surface density.

    Parameters
    ----------
    sfr_surface_density_min : float
        The minimum value of the SFR surface density

    sfr_surface_density_max : float
        The maximum value of the SFR surface density

    Returns
    -------
    sfr_surface_density : :obj:'~numpy.ndarray'
        vector of SFR surface densities

    vel_disp : :obj:'~numpy.ndarray'
        vector of outflow velocity dispersions following the trend
    """
    #create a vector for sfr surface density
    sfr_surface_density = np.linspace(sfr_surface_density_min, sfr_surface_density_max+4, num=1000)

    #use the relationship to predict the v_out
    vel_disp = 241*sfr_surface_density**0.3

    return sfr_surface_density, vel_disp

def kim_et_al_2020(sfr_surface_density_min, sfr_surface_density_max, scale_factor=1):
    """
    The trendline from Kim et al. (2020) where mass the loading factor is proportional
    to (SFR surface density)^-0.44

    Parameters
    ----------
    sfr_surface_density_min : float
        The minimum value of the SFR surface density

    sfr_surface_density_max : float
        The maximum value of the SFR surface density

    scale_factor : float
        The number by which to scale the trend, can be used to bring the trend
        into the range of the data on the plot (Default = 1)

    Returns
    -------
    sfr_surface_density : :obj:'~numpy.ndarray'
        vector of SFR surface densities

    mlf : :obj:'~numpy.ndarray'
        vector of mass loading factors following the trend
    """
    #create a vector for sfr surface density
    sfr_surface_density = np.linspace(sfr_surface_density_min, sfr_surface_density_max+4, num=1000)

    #use the relationship to predict the v_out
    mlf = scale_factor*sfr_surface_density**-0.44

    return sfr_surface_density, mlf

#===============================================================================
# USEFUL LITTLE FUNCTIONS
#===============================================================================

def fitting_function(x, a, b):
    """
    My fitting function to be fit to the v_out to sfr surface density data

    Parameters
    ----------
    x : (vector)
        the SFR surface density

    a, b : (int)
        constants to be fit

    Returns
    -------
    y : (vector)
        the outflow velocity
    """
    return a*(x**b)

def running_mean(x, N):
    """
    Calculates the running mean

    Parameters
    ----------
    x : :obj:'~numpy.ndarray'
        data
    N : integer
        bin size
    """

    cumsum = np.cumsum(np.insert(x, 0, 0))
    return (cumsum[N:] - cumsum[:-N]) / float(N)

def lower_quantile(x):
    """
    Calculate the lower quantile of x (data :obj:'~numpy.nd:obj:'~numpy.ndarray'')
    """
    return np.nanquantile(x, 0.33)

def upper_quantile(x):
    """
    Calculate the upper quantile of x (data :obj:'~numpy.ndarray')
    """
    return np.nanquantile(x, 0.66)

def binned_median_quantile_log(x, y, num_bins, weights=None, min_bin=None, max_bin=None):
    """
    Calculate the median, upper and lower quantile for an array of data in
    logarithmically increasing bins

    Parameters
    ----------
    x : :obj:'~numpy.ndarray'
        x-axis logarithmic data

    y : :obj:'~numpy.ndarray'
        y-axis data

    num_bins : integer
        the number of bins to divide the data into

    weights : :obj:'~numpy.ndarray'
        array to multiply x by, usually the error (Default = None)

    min_bin : float
        starting value of the first bin (Default = None)

    max_bin : float
        ending value of the last bin (Default = None)

    Returns
    -------
    logspace : :obj:'~numpy.ndarray'
        the logarithmic array of bin edges in x

    bin_center : :obj:'~numpy.ndarray'
        values indicating the centres of the bins in x

    bin_avg : :obj:'~numpy.ndarray'
        values of the median of the bins in y

    lower_quantile : :obj:'~numpy.ndarray'
        values for the lower quantile of each bin in y

    upper_quantile : :obj:'~numpy.ndarray'
        values for the upper quantile of each bin in y

    bin_stdev : :obj:'~numpy.ndarray'
        values for the standard deviation of each bin in y
    """
    if min_bin == None:
        min_bin = np.nanmin(x)

    if max_bin == None:
        max_bin = np.nanmax(x)

    #create the logspace - these are the bin edges
    logspace = np.logspace(np.log10(min_bin), np.log10(max_bin), num=num_bins+1)

    #calculate the average
    bin_avg = np.zeros(len(logspace)-1)
    upper_quantile = np.zeros(len(logspace)-1)
    lower_quantile = np.zeros(len(logspace)-1)
    bin_stdev = np.zeros(len(logspace)-1)

    for i in range(0, len(logspace)-1):
        left_bound = logspace[i]
        right_bound = logspace[i+1]
        items_in_bin = y[(x>left_bound)&(x<=right_bound)]
        print('Number of items in bin '+str(i)+': '+str(items_in_bin.shape))
        #calculate the median of the bin
        if weights == None:
            bin_avg[i] = np.nanmedian(items_in_bin)
        else:
            weights_in_bin = weights[0][(x>left_bound)&(x<=right_bound)]
            weights_in_bin = 1.0 - weights_in_bin/items_in_bin
            bin_avg[i] = np.average(items_in_bin, weights=weights_in_bin)

        #calculate the quartiles of the bin
        if items_in_bin.shape[0] < 10:
            upper_quantile[i] = np.nanquantile(items_in_bin, 0.80)
            lower_quantile[i] = np.nanquantile(items_in_bin, 0.20)
        else:
            upper_quantile[i] = np.nanquantile(items_in_bin, 0.66)
            lower_quantile[i] = np.nanquantile(items_in_bin, 0.33)

        #calculate the standard deviation of the bin
        bin_stdev[i] = np.nanstd(items_in_bin)

    #calculate the bin center for plotting
    bin_center = np.zeros(len(logspace)-1)
    for i in range(0, len(logspace)-1):
        bin_center[i] = np.nanmean([logspace[i],logspace[i+1]])

    return logspace, bin_center, bin_avg, lower_quantile, upper_quantile, bin_stdev

def binned_median_quantile_lin(x, y, num_bins, weights=None, min_bin=None, max_bin=None):
    """
    Calculate the median, upper and lower quantile for an array of data in
    linearly increasing bins

    Parameters
    ----------
    x : :obj:'~numpy.ndarray'
        x-axis linear data

    y : :obj:'~numpy.ndarray'
        y-axis data

    num_bins : integer
        the number of bins to divide the data into

    weights : :obj:'~numpy.ndarray'
        array to multiply x by, usually the error (Default = None)

    min_bin : float
        starting value of the first bin (Default = None)

    max_bin : float
        ending value of the last bin (Default = None)

    Returns
    -------
    linspace : :obj:'~numpy.ndarray'
        the array of linear bin edges in x

    bin_center : :obj:'~numpy.ndarray'
        values indicating the centres of the bins in x

    bin_avg : :obj:'~numpy.ndarray'
        values of the median of the bins in y

    lower_quantile : :obj:'~numpy.ndarray'
        values for the lower quantile of each bin in y

    upper_quantile : :obj:'~numpy.ndarray'
        values for the upper quantile of each bin in y

    bin_stdev : :obj:'~numpy.ndarray'
        values for the standard deviation of each bin in y
    """
    if min_bin == None:
        min_bin = np.nanmin(x)
    if max_bin == None:
        max_bin = np.nanmax(x)

    #create the logspace - these are the bin edges
    linspace = np.linspace(min_bin, max_bin, num=num_bins+1)

    #calculate the average
    bin_avg = np.zeros(len(linspace)-1)
    upper_quantile = np.zeros(len(linspace)-1)
    lower_quantile = np.zeros(len(linspace)-1)
    bin_stdev = np.zeros(len(linspace)-1)

    for i in range(0, len(linspace)-1):
        left_bound = linspace[i]
        right_bound = linspace[i+1]
        items_in_bin = y[(x>left_bound)&(x<=right_bound)]
        print('Number of items in bin '+str(i)+': '+str(items_in_bin.shape))
        if weights == None:
            bin_avg[i] = np.nanmedian(items_in_bin)
        else:
            weights_in_bin = weights[0][(x>left_bound)&(x<=right_bound)]
            weights_in_bin = 1.0 - weights_in_bin/items_in_bin
            bin_avg[i] = np.average(items_in_bin, weights=weights_in_bin)

        if items_in_bin.shape[0] < 10:
            upper_quantile[i] = np.nanquantile(items_in_bin, 0.80)
            lower_quantile[i] = np.nanquantile(items_in_bin, 0.20)
        else:
            upper_quantile[i] = np.nanquantile(items_in_bin, 0.66)
            lower_quantile[i] = np.nanquantile(items_in_bin, 0.33)

        #calculate the standard deviation of the bin
        bin_stdev[i] = np.nanstd(items_in_bin)

    #calculate the bin center for plotting
    bin_center = np.zeros(len(linspace)-1)
    for i in range(0, len(linspace)-1):
        bin_center[i] = np.nanmean([linspace[i],linspace[i+1]])

    return linspace, bin_center, bin_avg, lower_quantile, upper_quantile, bin_stdev

def pearson_correlation(x, y):
    """
    Calculate the Pearson correlation coefficient and p-value

    Parameters
    ----------
    x : :obj:'~numpy.ndarray'
        Input array - x values

    y : :obj:'~numpy.ndarray'
        Input array - y values

    Returns
    -------
    r : float
        Pearson's correlation coefficient

    p_value : float
        Two-tailed p-value
    """
    r, p_value = stats.pearsonr(x, y)

    return r, p_value

def read_in_create_wcs(fits_file, index=0, shift=None):
    """
    Reads in the fits file and creates the wcs

    Parameters
    ----------
    fits_file : string
        the filepath for the fits file to read in

    index : int
        the index of the extension to be loaded (default is 0)

    shift : list or None
        how to alter the header if the wcs is going to be wrong.
        e.g. ['CRPIX2', 32.0] will change the header value of CRPIX2 to 32.0

    Returns
    -------
    fits_data : :obj:'~numpy.ndarray'
        the fits data as a numpy array

    fits_wcs : astropy WCS object
        the world coordinate system for the fits file
    """
    #read the data in from fits
    with fits.open(fits_file) as hdu:
        hdu.info()
        fits_data = hdu[index].data
        fits_header = hdu[index].header
    hdu.close()

    #shift the header
    if shift:
        fits_header[shift[0]] = shift[1]

    #create the WCS
    fits_wcs = WCS(fits_header)

    return fits_data, fits_header, fits_wcs

def plot_continuum_contours(lamdas, xx, yy, data, z, ax):
    """
    Plots the continuum contours, using the rest wavelengths between 4600 and 4800 to define the continuum.

    Parameters
    ----------
    lamdas : :obj:'~numpy.ndarray'
        wavelength vector (1D)

    xx : :obj:'~numpy.ndarray'
        x coordinate array (2D)

    yy : :obj:'~numpy.ndarray'
        y coordinate array (2D)

    data : :obj:'~numpy.ndarray'
        data array (3D)

    z : float
        redshift of the galaxy

    ax : matplotlib axis instance
        axis for matplotlib to draw on

    Returns
    -------
    cont_contours : matplotlib.contour.QuadContourSet instance

    """
    #create a mask for the continuum
    cont_mask = (lamdas>4600*(1+z))&(lamdas<4800*(1+z))

    #find the median of the continuum
    cont_median = np.median(data[cont_mask,:,:], axis=0)

    #create the contours
    cont_contours = ax.contour(xx, yy, cont_median, colors='black', linewidths=0.7, alpha=0.7, levels=(0.2,0.3,0.4,0.7,1.0,2.0,4.0))

    return cont_contours