import torch
import copy
from torch import nn
import torch.nn.functional as F

from layers.roi_align_3d import ROIAlign3D
from backbone.resunet import ResUNet, ResUNet_FPN
from rpn.rpn_head import RpnHead
from rpn.rpn_module import RpnModule
from roi_heads.box_head import BoxHead
from roi_heads.box_module import BoxModule
from util import convert_to_roi_format, pad2factor, norm, add_rcnn_probability, get_anchors
from loss_fun.loss import Loss, Loss_FPN


@torch.jit.script
def make_rpn_windows(fs, stride, anchors):
    """
    Generating anchor boxes at each voxel on the feature map,
    the center of the anchor box on each voxel corresponds to center
    on the original input image.

    return
    windows: list of anchor boxes, [z, y, x, d, h, w]
    """
    offset = (stride - 1) / 2
    device = fs.device

    # anchors = torch.tensor(anchors, dtype=torch.float32, device=device)
    # D, H , W =(i/cfg['stride'] for i in cfg['crop_size'])

    _, _, D, H, W = fs.shape
    oz = torch.arange(offset, offset + stride * (D - 1) + 1, step=stride, dtype=torch.float32, device=device)
    oh = torch.arange(offset, offset + stride * (H - 1) + 1, step=stride, dtype=torch.float32, device=device)
    ow = torch.arange(offset, offset + stride * (W - 1) + 1, step=stride, dtype=torch.float32, device=device)
    oanchor = torch.arange(anchors.size()[0], dtype=torch.float32, device=device)

    shift_z, shift_y, shift_x, anchor_idx = torch.meshgrid(oz, oh, ow, oanchor)
    shift_z = shift_z.reshape(-1)
    shift_y = shift_y.reshape(-1)
    shift_x = shift_x.reshape(-1)
    anchor_idx = anchor_idx.reshape(-1).long()
    shift_anchors = torch.index_select(anchors, 0, anchor_idx)
    shift_anchor_z = shift_anchors[:, 0].squeeze()
    shift_anchor_y = shift_anchors[:, 1].squeeze()
    shift_anchor_x = shift_anchors[:, 2].squeeze()
    windows = torch.stack((shift_z, shift_y, shift_x, shift_anchor_z, shift_anchor_y, shift_anchor_x), dim=1)

    return windows


class NoduleNet(nn.Module):
    def __init__(self, cfg, mode='train'):
        super(NoduleNet, self).__init__()
        self.mode = mode
        self.cfg = cfg

        self.stride = torch.tensor(cfg.data_process_stride, dtype=torch.float32)
        # self.anchors = torch.tensor(get_anchors(cfg.MODEL.ANCHOR.BASES, cfg.MODEL.ANCHOR.ASPECT_RATIOS),
        #                             dtype=torch.float32).cuda()
        self.anchors = [torch.tensor(get_anchors(cfg.model_anchor_bases, cfg.model_abchor_aspect_ratios)[i],
                                    dtype=torch.float32).cuda() for i in range(3)]
        if cfg.MODEL.BACKBONE.FPN:
            self.feature_net = ResUNet_FPN(cfg, 1, 128)
            self.rpn_heads = nn.ModuleList([RpnHead(len(cfg.model_anchor_bases[0]), in_channels=64),
                              RpnHead(len(cfg.model_anchor_bases[1]), in_channels=64),
                              RpnHead(len(cfg.model_anchor_bases[2]), in_channels=128)])
        else:
            self.feature_net = ResUNet(cfg, 1, 128)
            self.rpn_head = RpnHead(len(cfg.model_anchor_bases), in_channels=128)
        self.box_head = BoxHead(cfg.model_roi_box_head_num_class, cfg.model_roi_box_head_pooler_resolution,
                                in_channels=128)
        self.rpn = RpnModule(cfg, self.mode)
        self.roi_box = BoxModule(cfg)
        self.roi_crop = ROIAlign3D(cfg.model_roi_box_head_pooler_resolution,
                                   spatial_scale=1 / cfg.data_process_stride,
                                   sampling_ratio=cfg.model_roi_box_head_sample_ratio)
        self.use_rcnn = False
        if cfg.model_backbone_fpn:
            self.loss = Loss_FPN(self.rpn.loss, self.roi_box.loss, self.use_rcnn)
        else:
            self.loss = Loss(self.rpn.loss, self.roi_box.loss, self.use_rcnn)

    def forward_depoly(self, inputs):
        inputs = pad2factor(inputs)
        inputs = norm(inputs)
        features, feat_4 = self.feature_net(inputs)
        if self.cfg.model_backbone_fpn:
            self.rpn_logits, self.rpn_deltas = [], []
            self.rpn_windows = [make_rpn_windows(features[0], self.stride * 2, self.anchors[0]),
                                make_rpn_windows(features[1], self.stride, self.anchors[1]),
                                make_rpn_windows(features[2], self.stride // 2, self.anchors[2])]
            for i in range(3):
                rpn_logits, rpn_deltas = self.rpn_heads[i](features[i])
                self.rpn_logits.append(rpn_logits)
                self.rpn_deltas.append(rpn_deltas)
        else:
            self.rpn_window = make_rpn_windows(features[-1], self.stride, self.anchors)
            self.rpn_logits, self.rpn_deltas = self.rpn_head(features[-1])
        if self.cfg.model_backbone_fpn:
            self.rpn_proposals = [
                self.rpn.rpn_nms(inputs, self.rpn_windows[i], self.rpn_logits[i], self.rpn_deltas[i])
                for i in range(3)]
        else:
            self.rpn_proposals = self.rpn.rpn_nms(inputs, self.rpn_window, self.rpn_logits, self.rpn_deltas)
        # self.rpn_proposal_rois = convert_to_roi_format(self.rpn_proposals)
        # self.rcnn_logits, self.rcnn_deltas = self.box_head(
        #     self.roi_crop(feat_4.float(), self.rpn_proposal_rois.float()).half())
        # self.rpn_proposals = add_rcnn_probability(self.rpn_proposals, self.rcnn_logits)
        return self.rpn_proposals

    def forward(self, inputs, truth_bboxes):
        features, feat_4 = self.feature_net(inputs)
        if self.cfg.model_backbone_fpn:
            self.rpn_logits, self.rpn_deltas = [], []
            self.rpn_windows = [make_rpn_windows(features[0], self.stride * 2, self.anchors[0]),
                                make_rpn_windows(features[1], self.stride, self.anchors[1]),
                                make_rpn_windows(features[2], self.stride // 2, self.anchors[2])]
            for i in range(3):
                rpn_logits, rpn_deltas = self.rpn_heads[i](features[i])
                self.rpn_logits.append(rpn_logits)
                self.rpn_deltas.append(rpn_deltas)
        else:
            self.rpn_window = make_rpn_windows(features[-1], self.stride, self.anchors)
            self.rpn_logits, self.rpn_deltas = self.rpn_head(features[-1])
        if self.mode in ['train', 'valid']:
            if self.cfg.model_backbone_fpn:
                self.rpn_labels, self.rpn_label_assigns, self.rpn_label_weights, self.rpn_targets, self.rpn_target_weights = \
                    [], [], [], [], []
                for i in range(3):
                    rpn_labels, rpn_label_assigns, rpn_label_weights, rpn_targets, rpn_target_weights = \
                        self.rpn.make_rpn_target(self.rpn_windows[i], truth_bboxes, self.cfg.DATA.DATA_PROCESS.FPN_DIAMETER_RANGE[i])
                    self.rpn_labels.append(rpn_labels)
                    self.rpn_label_assigns.append(rpn_label_assigns)
                    self.rpn_label_weights.append(rpn_label_weights)
                    self.rpn_targets.append(rpn_targets)
                    self.rpn_target_weights.append(rpn_target_weights)
            else:
                self.rpn_labels, self.rpn_label_assigns, self.rpn_label_weights, self.rpn_targets, self.rpn_target_weights = \
                    self.rpn.make_rpn_target(self.rpn_window, truth_bboxes)

            if self.use_rcnn:
                self.rpn_proposals = self.rpn.rpn_nms(inputs, self.rpn_window, self.rpn_logits, self.rpn_deltas)
                self.rpn_proposals, self.rcnn_labels, self.rcnn_targets = self.roi_box.make_box_target(
                    self.rpn_proposals, truth_bboxes)
                if len(self.rpn_proposals) > 0:
                    self.rpn_proposal_rois = convert_to_roi_format(self.rpn_proposals)
                    self.rcnn_logits, self.rcnn_deltas = self.box_head(self.roi_crop(feat_4, self.rpn_proposal_rois))
            else:
                self.rcnn_logits, self.rcnn_deltas, self.rcnn_labels, self.rcnn_targets = None, None, None, None

            return self.rpn_logits, self.rpn_deltas, \
                   self.rpn_labels, self.rpn_label_weights, self.rpn_targets, self.rpn_target_weights, \
                   self.rcnn_logits, self.rcnn_deltas, self.rcnn_labels, self.rcnn_targets

        elif self.mode in ['test']:
            if self.cfg.model_backbone_fpn:
                self.rpn_proposals = [
                    self.rpn.rpn_nms(inputs, self.rpn_windows[i], self.rpn_logits[i], self.rpn_deltas[i])
                    for i in range(3)]
            else:
                self.rpn_proposals = self.rpn.rpn_nms(inputs, self.rpn_window, self.rpn_logits, self.rpn_deltas)

            if self.use_rcnn:
                if len(self.rpn_proposals) > 0:
                    self.ensemble_proposals = copy.deepcopy(self.rpn_proposals)
                    self.rpn_proposal_rois = convert_to_roi_format(self.rpn_proposals)
                    self.rcnn_logits, self.rcnn_deltas = self.box_head(self.roi_crop(feat_4, self.rpn_proposal_rois))
                    self.fpr_res = self.roi_box.get_probability(inputs, self.rpn_proposals, self.rcnn_logits,
                                                                self.rcnn_deltas)
                    self.ensemble_proposals[:, 1] = (self.ensemble_proposals[:, 1] + self.fpr_res[:, 0]) / 2
                    return self.rpn_proposals, self.fpr_res, self.ensemble_proposals
                else:
                    raise RuntimeError("The nms pre score threshold {} is too high, change for a lower one.".format(
                        self.pre_nms_score_threshold))
            else:
                return self.rpn_proposals
        else:
            raise ValueError('rpn_nms(): invalid mode = %s?' % self.mode)

    def set_mode(self, mode):
        assert mode in ['train', 'valid', 'test']
        self.mode = mode
        # self.rpn.mode = self.mode
        if mode in ['train']:
            self.train()
        else:
            self.eval()
