import argparse
import re

try:
    from functools import lru_cache
except ImportError:
    from backports.functools_lru_cache import lru_cache

import mxnet as mx
import numpy as np
from mxnet import gluon
from mxnet.gluon import nn, rnn

from utils.embedding_maker import (encoding_and_padding, load_embedding,
                                   load_vocab)

MAX_LEN = 198

def spacing(sent):
    if len(sent) > MAX_LEN:
        splitted_sent = [sent[y-MAX_LEN:y] for y in range(MAX_LEN, len(sent)+MAX_LEN, MAX_LEN)]
        spaced_sent = ''.join([predictor.get_spaced_sent(ss)
                               for ss in splitted_sent])
    else:
        spaced_sent = predictor.get_spaced_sent(sent)
    return spaced_sent.strip()

# Model class
class korean_autospacing_base(gluon.HybridBlock):
    def __init__(self, n_hidden, vocab_size, embed_dim, max_seq_length,
                 **kwargs):
        super(korean_autospacing_base, self).__init__(**kwargs)
        # 입력 시퀀스 길이
        self.in_seq_len = max_seq_length
        # 출력 시퀀스 길이
        self.out_seq_len = max_seq_length
        # GRU의 hidden 개수
        self.n_hidden = n_hidden
        # 고유문자개수
        self.vocab_size = vocab_size
        # max_seq_length
        self.max_seq_length = max_seq_length
        # 임베딩 차원수
        self.embed_dim = embed_dim

        with self.name_scope():
            self.embedding = nn.Embedding(input_dim=self.vocab_size,
                                          output_dim=self.embed_dim)

            self.conv_unigram = nn.Conv2D(channels=128,
                                          kernel_size=(1, self.embed_dim))

            self.conv_bigram = nn.Conv2D(channels=256,
                                         kernel_size=(2, self.embed_dim),
                                         padding=(1, 0))

            self.conv_trigram = nn.Conv2D(channels=128,
                                          kernel_size=(3, self.embed_dim),
                                          padding=(1, 0))

            self.conv_forthgram = nn.Conv2D(channels=64,
                                            kernel_size=(4, self.embed_dim),
                                            padding=(2, 0))

            self.conv_fifthgram = nn.Conv2D(channels=32,
                                            kernel_size=(5, self.embed_dim),
                                            padding=(2, 0))

            self.bi_gru = rnn.GRU(hidden_size=self.n_hidden, layout='NTC', bidirectional=True)
            self.dense_sh = nn.Dense(100, activation='relu', flatten=False)
            self.dense = nn.Dense(1, activation='sigmoid', flatten=False)

    def hybrid_forward(self, F, inputs):
        embed = self.embedding(inputs)
        embed = F.expand_dims(embed, axis=1)
        unigram = self.conv_unigram(embed)
        bigram = self.conv_bigram(embed)
        trigram = self.conv_trigram(embed)
        forthgram = self.conv_forthgram(embed)
        fifthgram = self.conv_fifthgram(embed)

        grams = F.concat(unigram,
                         F.slice_axis(bigram,
                                      axis=2,
                                      begin=0,
                                      end=self.max_seq_length),
                         trigram,
                         F.slice_axis(forthgram,
                                      axis=2,
                                      begin=0,
                                      end=self.max_seq_length),
                         F.slice_axis(fifthgram,
                                      axis=2,
                                      begin=0,
                                      end=self.max_seq_length),
                         dim=1)

        grams = F.transpose(grams, (0, 2, 3, 1))
        grams = F.reshape(grams, (-1, self.max_seq_length, -3))
        grams = self.bi_gru(grams)
        fc1 = self.dense_sh(grams)
        return (self.dense(fc1))

# https://raw.githubusercontent.com/haven-jeon/Train_KoSpacing/master/img/kosapcing_img.png
class korean_autospacing2(gluon.HybridBlock):
    def __init__(self, n_hidden, vocab_size, embed_dim, max_seq_length,
                 **kwargs):
        super(korean_autospacing2, self).__init__(**kwargs)
        # 입력 시퀀스 길이
        self.in_seq_len = max_seq_length
        # 출력 시퀀스 길이
        self.out_seq_len = max_seq_length
        # GRU의 hidden 개수
        self.n_hidden = n_hidden
        # 고유문자개수
        self.vocab_size = vocab_size
        # max_seq_length
        self.max_seq_length = max_seq_length
        # 임베딩 차원수
        self.embed_dim = embed_dim

        with self.name_scope():
            self.embedding = nn.Embedding(input_dim=self.vocab_size,
                                          output_dim=self.embed_dim)

            self.conv_unigram = nn.Conv2D(channels=128,
                                          kernel_size=(1, self.embed_dim))

            self.conv_bigram = nn.Conv2D(channels=128,
                                         kernel_size=(2, self.embed_dim),
                                         padding=(1, 0))

            self.conv_trigram = nn.Conv2D(channels=64,
                                          kernel_size=(3, self.embed_dim),
                                          padding=(2, 0))

            self.conv_forthgram = nn.Conv2D(channels=32,
                                            kernel_size=(4, self.embed_dim),
                                            padding=(3, 0))

            self.conv_fifthgram = nn.Conv2D(channels=16,
                                            kernel_size=(5, self.embed_dim),
                                            padding=(4, 0))
            # for reverse convolution
            self.conv_rev_bigram = nn.Conv2D(channels=128,
                                             kernel_size=(2, self.embed_dim),
                                             padding=(1, 0))

            self.conv_rev_trigram = nn.Conv2D(channels=64,
                                              kernel_size=(3, self.embed_dim),
                                              padding=(2, 0))

            self.conv_rev_forthgram = nn.Conv2D(channels=32,
                                                kernel_size=(4,
                                                             self.embed_dim),
                                                padding=(3, 0))

            self.conv_rev_fifthgram = nn.Conv2D(channels=16,
                                                kernel_size=(5,
                                                             self.embed_dim),
                                                padding=(4, 0))
            self.bi_gru = rnn.GRU(hidden_size=self.n_hidden, layout='NTC', bidirectional=True)
            # self.bi_gru = rnn.BidirectionalCell(
            #     rnn.GRUCell(hidden_size=self.n_hidden),
            #     rnn.GRUCell(hidden_size=self.n_hidden))
            self.dense_sh = nn.Dense(100, activation='relu', flatten=False)
            self.dense = nn.Dense(1, activation='sigmoid', flatten=False)

    def hybrid_forward(self, F, inputs):
        embed = self.embedding(inputs)
        embed = F.expand_dims(embed, axis=1)
        rev_embed = embed.flip(axis=2)

        unigram = self.conv_unigram(embed)
        bigram = self.conv_bigram(embed)
        trigram = self.conv_trigram(embed)
        forthgram = self.conv_forthgram(embed)
        fifthgram = self.conv_fifthgram(embed)

        rev_bigram = self.conv_rev_bigram(rev_embed).flip(axis=2)
        rev_trigram = self.conv_rev_trigram(rev_embed).flip(axis=2)
        rev_forthgram = self.conv_rev_forthgram(rev_embed).flip(axis=2)
        rev_fifthgram = self.conv_rev_fifthgram(rev_embed).flip(axis=2)

        grams = F.concat(unigram,
                         F.slice_axis(bigram,
                                      axis=2,
                                      begin=0,
                                      end=self.max_seq_length),
                         F.slice_axis(rev_bigram,
                                      axis=2,
                                      begin=0,
                                      end=self.max_seq_length),
                         F.slice_axis(trigram,
                                      axis=2,
                                      begin=0,
                                      end=self.max_seq_length),
                         F.slice_axis(rev_trigram,
                                      axis=2,
                                      begin=0,
                                      end=self.max_seq_length),
                         F.slice_axis(forthgram,
                                      axis=2,
                                      begin=0,
                                      end=self.max_seq_length),
                         F.slice_axis(rev_forthgram,
                                      axis=2,
                                      begin=0,
                                      end=self.max_seq_length),
                         F.slice_axis(fifthgram,
                                      axis=2,
                                      begin=0,
                                      end=self.max_seq_length),
                         F.slice_axis(rev_fifthgram,
                                      axis=2,
                                      begin=0,
                                      end=self.max_seq_length),
                         dim=1)

        grams = F.transpose(grams, (0, 2, 3, 1))
        grams = F.reshape(grams, (-1, self.max_seq_length, -3))
        grams = self.bi_gru(grams)
        fc1 = self.dense_sh(grams)
        return (self.dense(fc1))

def pick_model(model_nm, n_hidden, vocab_size, embed_dim, max_seq_length):
    if model_nm.lower() == 'kospacing':
        model = korean_autospacing_base(n_hidden=n_hidden,
                                        vocab_size=vocab_size,
                                        embed_dim=embed_dim,
                                        max_seq_length=max_seq_length)
    elif model_nm.lower() == 'kospacing2':
        model = korean_autospacing2(n_hidden=n_hidden,
                                    vocab_size=vocab_size,
                                    embed_dim=embed_dim,
                                    max_seq_length=max_seq_length)
    else:
        assert False
    return model

class pred_spacing:
    def __init__(self, model, w2idx):
        self.model = model
        self.w2idx = w2idx
        self.pattern = re.compile(r'\s+')

    @lru_cache(maxsize=None)
    def get_spaced_sent(self, raw_sent):
        raw_sent_ = "«" + raw_sent + "»"
        raw_sent_ = raw_sent_.replace(' ', '^')
        sents_in = [raw_sent_,]
        mat_in = encoding_and_padding(word2idx_dic=self.w2idx,
                                      sequences=sents_in,
                                      maxlen=200,
                                      padding='post',
                                      truncating='post')
        mat_in = mx.nd.array(mat_in, ctx=mx.cpu(0))
        results = self.model(mat_in)
        mat_set = results[0, ]
        preds = np.array(
            ['1' if i > 0.5 else '0' for i in mat_set[:len(raw_sent_)]])
        return self.make_pred_sents(raw_sent_, preds)

    def make_pred_sents(self, x_sents, y_pred):
        res_sent = []
        for i, j in zip(x_sents, y_pred):
            if j == '1':
                res_sent.append(i)
                res_sent.append(' ')
            else:
                res_sent.append(i)
        subs = re.sub(self.pattern, ' ', ''.join(res_sent).replace('^', ' '))
        subs = subs.replace('«', '')
        subs = subs.replace('»', '')
        return subs

# 사전 파일 로딩
w2idx, idx2w = load_vocab('model/w2idx.dic')
# 임베딩 파일 로딩
weights = load_embedding('model/kospacing_wv.np')
vocab_size = weights.shape[0]
embed_dim = weights.shape[1]
model = pick_model('kospacing', 200, vocab_size, embed_dim, 200)

# model.collect_params().initialize(mx.init.Xavier(), ctx=mx.cpu(0))
# model.embedding.weight.set_data(weights)
model.load_parameters('model/kospacing.params', ctx=mx.cpu(0))
predictor = pred_spacing(model, w2idx)


# -*- coding: utf-8 -*-
"""
Created on Wed Dec 15 11:29:06 2021

@author: fran-
"""

#from keras import backend as K
import tensorflow as tf
import os

from tensorflow.keras.models import load_model
from tensorflow.keras.preprocessing import image
from tensorflow.keras.optimizers import Adam
from imageio import imread
import numpy as np
from matplotlib import pyplot as plt

from models.keras_ssd512 import ssd_512
from keras_loss_function.keras_ssd_loss import SSDLoss
from keras_layers.keras_layer_AnchorBoxes import AnchorBoxes
from keras_layers.keras_layer_DecodeDetections import DecodeDetections
from keras_layers.keras_layer_DecodeDetectionsFast import DecodeDetectionsFast
from keras_layers.keras_layer_L2Normalization import L2Normalization

from ssd_encoder_decoder.ssd_output_decoder import decode_detections, decode_detections_fast

from data_generator.object_detection_2d_data_generator import DataGenerator
from data_generator.object_detection_2d_photometric_ops import ConvertTo3Channels
from data_generator.object_detection_2d_geometric_ops import Resize
from data_generator.object_detection_2d_misc_utils import apply_inverse_transforms

# Set the image size.
img_height = 512
img_width = 512

# 1: Build the Keras model

tf.keras.backend.clear_session() # Clear previous models from memory.

model = ssd_512(image_size=(img_height, img_width, 3),
                n_classes=20,
                mode='inference',
                l2_regularization=0.0005,
                scales=[0.07, 0.15, 0.3, 0.45, 0.6, 0.75, 0.9, 1.05], # The scales for MS COCO are [0.04, 0.1, 0.26, 0.42, 0.58, 0.74, 0.9, 1.06]
                aspect_ratios_per_layer=[[1.0, 2.0, 0.5],
                                         [1.0, 2.0, 0.5, 3.0, 1.0/3.0],
                                         [1.0, 2.0, 0.5, 3.0, 1.0/3.0],
                                         [1.0, 2.0, 0.5, 3.0, 1.0/3.0],
                                         [1.0, 2.0, 0.5, 3.0, 1.0/3.0],
                                         [1.0, 2.0, 0.5],
                                         [1.0, 2.0, 0.5]],
               two_boxes_for_ar1=True,
               steps=[8, 16, 32, 64, 128, 256, 512],
               offsets=[0.5, 0.5, 0.5, 0.5, 0.5, 0.5, 0.5],
               clip_boxes=False,
               variances=[0.1, 0.1, 0.2, 0.2],
               normalize_coords=True,
               subtract_mean=[123, 117, 104],
               swap_channels=[2, 1, 0],
               confidence_thresh=0.5,
               iou_threshold=0.2,
               top_k=200,
               nms_max_output_size=400)

# 2: Load the trained weights into the model.

# TODO: Set the path of the trained weights.
weights_path = 'VGG_VOC0712Plus_SSD_512x512_ft_iter_160000.h5'

model.load_weights(weights_path, by_name=True)

# 3: Compile the model so that Keras won't complain the next time you load it.

adam = Adam(lr=0.001, beta_1=0.9, beta_2=0.999, epsilon=1e-08, decay=0.0)

ssd_loss = SSDLoss(neg_pos_ratio=3, alpha=1.0)

model.compile(optimizer=adam, loss=ssd_loss.compute_loss)

orig_images = [] # Store the images here.
input_images = [] # Store resized versions of the images here.

# We'll only load one image in this example.
img_path = 'person_examples/'

input_img_paths = list(map(lambda x: os.path.join(img_path, x), os.listdir(img_path)))

for im_path in input_img_paths:
    orig_images.append(imread(im_path))
    img = image.load_img(im_path, target_size=(img_height, img_width))
    img = image.img_to_array(img)
    input_images.append(img)
input_images = np.array(input_images)

y_pred = model.predict(input_images)

confidence_threshold = 0

y_pred_thresh = [y_pred[k][y_pred[k,:,1] > confidence_threshold] for k in range(y_pred.shape[0])]

np.set_printoptions(precision=2, suppress=True, linewidth=90)
print("Predicted boxes:\n")
print('   class   conf xmin   ymin   xmax   ymax')
print(y_pred_thresh)

# Display the image and draw the predicted boxes onto it.

# Set the colors for the bounding boxes
colors = plt.cm.hsv(np.linspace(0, 1, 21)).tolist()
classes = ['background',
           'aeroplane', 'bicycle', 'bird', 'boat',
           'bottle', 'bus', 'car', 'cat',
           'chair', 'cow', 'diningtable', 'dog',
           'horse', 'motorbike', 'person', 'pottedplant',
           'sheep', 'sofa', 'train', 'tvmonitor']
for i in range(len(orig_images)):

    plt.figure(figsize=(20,12))
    plt.imshow(orig_images[i])

    current_axis = plt.gca()

    for box in y_pred_thresh[i]:
        # Transform the predicted bounding boxes for the 512x512 image to the original image dimensions.
        xmin = box[-4] * orig_images[0].shape[1] / img_width
        ymin = box[-3] * orig_images[0].shape[0] / img_height
        xmax = box[-2] * orig_images[0].shape[1] / img_width
        ymax = box[-1] * orig_images[0].shape[0] / img_height
        color = colors[int(box[0])]
        label = '{}: {:.2f}'.format(classes[int(box[0])], box[1])
        current_axis.add_patch(plt.Rectangle((xmin, ymin), xmax-xmin, ymax-ymin, color=color, fill=False, linewidth=2))
        current_axis.text(xmin, ymin, label, size='x-large', color='white', bbox={'facecolor':color, 'alpha':1.0})


# -*- coding:utf-8 -*-
# @TIME     :2019/1/19 10:57
# @Author   :Fan
# @File     :ImageAugment.py
"""
Purpose: To expand the dataset for our project
Contain: Aug_scale, Aug_crop, Aug_flip, Aug_rotate
Param_Trans: store the value of parameters
"""

"""
The order in this work:
0: nose 	        1: neck             2: right shoulder   3: right elbow
4: right wrist      5: left shoulder    6: left elbow	     7: left wrist
8: right hip        9: right knee	    10: right ankle	    11: left hip
12: left knee       13: left ankle	    14: right eye	    15: left eye
16: right ear       17: left ear
"""

import cv2
import numpy as np

def Aug_scale(meta, img, mask_miss, Param_Trans):
    param = np.random.rand()    # range in (0, 1)
    scale_max = Param_Trans['scale_max']
    scale_min = Param_Trans['scale_min']

    # multiplier linear shear into [scale_min, scale_max]
    multiplier = scale_min + (scale_max - scale_min) * param
    scale = multiplier * (Param_Trans['target_dist'] / meta['scale_provided'])
    img = cv2.resize(img, None, fx=scale, fy=scale,
                     interpolation=cv2.INTER_CUBIC)
    mask_miss = cv2.resize(mask_miss, None, fx=scale, fy=scale,
                     interpolation=cv2.INTER_CUBIC)

    # modify meta data
    meta['objpos'] *= scale
    meta['joint_self'][:, :2] *= scale
    if(meta['numOtherPeople'] != 0):
        meta['objpos_other'] *= scale
        meta['joint_others'][:, :, :2] *= scale
    return meta, img, mask_miss

def Aug_crop(meta, img, mask_miss, Param_Trans):
    offset_max = 2 * Param_Trans['crop_disturb']
    crop_x = int(Param_Trans['crop_size_x'])
    crop_y = int(Param_Trans['crop_size_y'])
    param_x = np.random.rand()
    param_y = np.random.rand()

    # offset has both positive and negative
    offset_x = int((param_x - 0.5) * offset_max)
    offset_y = int((param_y - 0.5) * offset_max)
    center = meta['objpos'] + np.array([offset_x, offset_y])
    center = center.astype(int)

    # padding
    pad_col = np.ones((crop_y, img.shape[1], 3), dtype=np.uint8) * 128
    pad_col_mask = np.ones((crop_y, mask_miss.shape[1]), dtype=np.uint8) * 255

    img = np.concatenate((pad_col, img, pad_col), axis=0)
    mask_miss = np.concatenate((pad_col_mask, mask_miss, pad_col_mask), axis=0)

    pad_row = np.ones((img.shape[0], crop_x, 3), dtype=np.uint8) * 128
    pad_row_mask = np.ones((mask_miss.shape[0], crop_x), dtype=np.uint8) * 255

    img = np.concatenate((pad_row, img, pad_row), axis=1)
    mask_miss = np.concatenate((pad_row_mask, mask_miss, pad_row_mask), axis=1)

    # crop
    img = img[int(center[1] + crop_y / 2):int(center[1] + crop_y * 3 / 2),
              int(center[0] + crop_x / 2):int(center[0] + crop_x * 3 / 2), :]
    mask_miss = mask_miss[int(center[1] + crop_y / 2):int(center[1] + crop_y * 3 / 2 + 1),
                          int(center[0] + crop_x / 2):int(center[0] + crop_x * 3 / 2 + 1)]

    offset_left = crop_x - (center[0] + crop_x / 2)
    offset_up = crop_y - (center[1] + crop_y / 2)
    offset = np.array([offset_left, offset_up])

    # modify meta data
    meta['objpos'] += offset
    meta['joint_self'][:, :2] += offset
    mask = np.logical_or.reduce((
        meta['joint_self'][:, 0] >= crop_x,
        meta['joint_self'][:, 0] < 0,
        meta['joint_self'][:, 1] >= crop_y,
        meta['joint_self'][:, 1] < 0
    ))

    # screen out the joint out of the cropped image
    meta['joint_self'][mask == True, 2] = 2
        # 0:visible + labeled   1:labeled   2:unlabeled
    if(meta['numOtherPeople'] != 0):
        meta['objpos_other'] += offset
        meta['joint_others'][:, :, :2] += offset
        mask_other = np.logical_or.reduce((
            meta['joint_others'][:, :, 0] >= crop_x,
            meta['joint_others'][:, :, 0] < 0,
            meta['joint_others'][:, :, 1] >= crop_y,
            meta['joint_others'][:, :, 1] < 0
        ))
        meta['joint_others'][mask_other == True, 2] = 2
    return meta, img, mask_miss

def Aug_flip(meta, img, mask_miss):
    param = np.random.rand()
    if (param <= 0.5):      # perform the flip (around y-axis)
        img_ = img.copy()
        mask_miss_ = mask_miss.copy()
        width = img.shape[1]

        # flipCode: 0:x-axis, positive value:y-axis, negative value:both
        cv2.flip(src=img_, flipCode=1, dst=img)
        cv2.flip(src=mask_miss_, flipCode=1, dst=mask_miss)

        # modify meta data
        meta['objpos'][0] = width - 1 - meta['objpos'][0]
        meta['joint_self'][:, 0] = width - 1 - meta['joint_self'][:, 0]
        meta['joint_self'] =  meta['joint_self'][[0,1,5,6,7,2,3,4,11,12,13,8,9,10,15,14,17,16]]
        if (meta['numOtherPeople'] != 0):
            meta['objpos_other'][:, 0] = width - 1 - meta['objpos_other'][:, 0]
            meta['joint_others'][:, :, 0] = width - 1 - meta['joint_others'][:, :, 0]
            for i in range(meta['numOtherPeople']):
                meta['joint_others'][i] = meta['joint_others'][i][
                    [0,1,5,6,7,2,3,4,11,12,13,8,9,10,15,14,17,16]
                ]
    return meta, img, mask_miss

def rotate_point(point, mat):
    p = np.ones((3, 1))
    p[0] = point[0]
    p[1] = point[1]
    new_point = mat.dot(p)
    point[0] = new_point[0]
    point[1] = new_point[1]
    return point

def Aug_rotate(meta, img, mask_miss, Param_Trans, num_joint=18):
    rotate_max = 2 * Param_Trans['rotate_degree']
    degree = (np.random.rand() - 0.5) * rotate_max
    scale = 1.0
    num_joint = num_joint

    # perform rotation without being cut off
    (img_h, img_w) = img.shape[:2]
    (mask_h, mask_w) = mask_miss.shape[:2]

    # apply the negative of the angle to rotate clockwise
    img_mat = cv2.getRotationMatrix2D(center=(img_w // 2, img_h // 2),
                                      angle=-degree, scale=scale)
    mask_mat = cv2.getRotationMatrix2D(center=(mask_w // 2, mask_h // 2),
                                       angle=-degree, scale=scale)

    img_cos = np.abs(img_mat[0, 0] / scale)
    img_sin = np.abs(img_mat[0, 1] / scale)
    mask_cos = np.abs(mask_mat[0, 0] / scale)
    mask_sin = np.abs(mask_mat[0, 1] / scale)

    # compute the new size of the image
    img_nw = int((img_h * img_sin) + (img_w * img_cos))
    img_nh = int((img_h * img_cos) + (img_w * img_sin))
    mask_nw = int((mask_h * mask_sin) + (mask_w * mask_cos))
    mask_nh = int((mask_h * mask_cos) + (mask_w * mask_sin))

    # adjust the rotation mat
    img_mat[0, 2] += (img_nw / 2) - img_w // 2
    img_mat[1, 2] += (img_nh / 2) - img_h // 2
    mask_mat[0, 2] += (mask_nw / 2) - mask_w // 2
    mask_mat[1, 2] += (mask_nh / 2) - mask_h // 2

    # do the rotate
    img_rot = cv2.warpAffine(src=img, M=img_mat, dsize=(img_nw, img_nh), flags=cv2.INTER_CUBIC,
                             borderMode=cv2.BORDER_CONSTANT, borderValue=(128, 128, 128))
    mask_rot = cv2.warpAffine(src=mask_miss, M=mask_mat, dsize=(mask_nw, mask_nh), flags=cv2.INTER_CUBIC,
                              borderMode=cv2.BORDER_CONSTANT, borderValue=(255))

    # modify meta data
    meta['objpos'] = rotate_point(meta['objpos'], img_mat)
    for i in range(num_joint):
        meta['joint_self'][i, :] = rotate_point(meta['joint_self'][i, :], img_mat)

    if (meta['numOtherPeople'] != 0):
        for j in range(meta['numOtherPeople']):
            meta['objpos_other'][j, :] = rotate_point(meta['objpos_other'][j, :], img_mat)

            for i in range(num_joint):
                meta['joint_others'][j, i, :] = rotate_point(
                    meta['joint_others'][j, i, :], img_mat)
    return meta, img_rot, mask_rot


import sys
sys.path.append('./face_modules/')
import torch
import torchvision.transforms as transforms
import torch.nn.functional as F
from face_modules.model import Backbone, Arcface, MobileFaceNet, Am_softmax, l2_norm
from network.AEI_Net import *
from face_modules.mtcnn import *
import cv2
import PIL.Image as Image
import numpy as np
import time
import copy

def swap_faces(Xs_raw, Xt_raw):
    # print ("2----")
    # detector = MTCNN()
    # device = torch.device('cpu')
    # G = AEI_Net(c_id=512)
    # G.eval()
    # G.load_state_dict(torch.load('./saved_models/G_latest.pth', map_location=device))
    # G = G.cpu()

    # arcface = Backbone(50, 0.6, 'ir_se').to(device)
    # arcface.eval()
    # arcface.load_state_dict(torch.load('./id_model/model_ir_se50.pth', map_location=device), strict=False)
    #
    # test_transform = transforms.Compose([
    #     transforms.ToTensor(),
    #     transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))
    # ])
    #
    # Xs_img = Image.fromarray(Xs_raw)
    # Xs = detector.align(Xs_img, crop_size=(256, 256))
    #
    # Xs = test_transform(Xs)
    # Xs = Xs.unsqueeze(0).cpu()
    # with torch.no_grad():
    #     embeds, Xs_feats = arcface(F.interpolate(Xs, (112, 112), mode='bilinear', align_corners=True))
    #     embeds = embeds.mean(dim=0, keepdim=True)

    mask = np.zeros([256, 256], dtype=np.float)
    for i in range(256):
        for j in range(256):
            dist = np.sqrt((i-32)**2 + (j-32)**2)/32
            dist = np.minimum(dist, 1)
            mask[i, j] = 1.-dist/10
    mask = cv2.dilate(mask, None, iterations=20)

    Xt_img = Image.fromarray(Xt_raw)

    Xt, trans_inv = detector.align(Xt_img, crop_size=(256, 256), return_trans_inv=True)

    Xt = test_transform(Xt)

    Xt = Xt.unsqueeze(0).cpu()
    with torch.no_grad():
        st = time.time()
        Yt, _ = G(Xt, embeds)
        Yt = Yt.squeeze().detach().cpu().numpy()
        st = time.time() - st
        print(f'inference time: {st} sec')
        Yt = Yt.transpose([1, 2, 0])*0.5 + 0.5
        Yt_trans_inv = cv2.warpAffine(Yt, trans_inv, (np.size(Xt_raw, 1), np.size(Xt_raw, 0)), borderMode=cv2.BORDER_TRANSPARENT)
        mask_ = cv2.warpAffine(mask, trans_inv, (np.size(Xt_raw, 1), np.size(Xt_raw, 0)), borderMode=cv2.BORDER_TRANSPARENT)
        mask_ = np.expand_dims(mask_, 2)
        Yt_trans_inv = mask_*Yt_trans_inv + (1-mask_)*(Xt_raw.astype(np.float)/255.)

        merge = Yt_trans_inv

        return merge

import dlib

if __name__ == '__main__':
    #人脸分类器
    detector_dlib = dlib.get_frontal_face_detector()
    # 获取人脸检测器
    predictor = dlib.shape_predictor(
        "./shape_predictor_68_face_landmarks.dat"
    )
    #-----------------------------------------------
    detector = MTCNN()
    device = torch.device('cuda')
    G = AEI_Net(c_id=512)
    G.eval()
    G.load_state_dict(torch.load('./saved_models/G_latest.pth', map_location=torch.device('cpu')))
    G = G.cuda()

    arcface = Backbone(50, 0.6, 'ir_se').to(device)
    arcface.eval()
    arcface.load_state_dict(torch.load('./id_model/model_ir_se50.pth', map_location=device), strict=False)

    test_transform = transforms.Compose([
        transforms.ToTensor(),
        transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))
    ])

    Xs_path = './samples/hd.jpg'
    # Xt_path = './samples/s1.jpg'

    Xs_raw = cv2.imread(Xs_path)
    Xs_raw_r = cv2.resize(Xs_raw, (256,256), interpolation = cv2.INTER_LINEAR)
    print(Xs_raw.shape)

    cv2.namedWindow('source',0)
    cv2.imshow('source', Xs_raw)

    video_capture = cv2.VideoCapture('./video/video1.mp4')
    # video_capture = cv2.VideoCapture('E:/m_cc/cv_cource/chapter_07/video/rw_4.mp4')

    loc_time = time.localtime()

    str_time = time.strftime("%Y-%m-%d_%H-%M-%S", loc_time)

    video_writer = cv2.VideoWriter("./demo/demo_{}.mp4".format(str_time), cv2.VideoWriter_fourcc(*"mp4v"), fps=15, frameSize=(2560, 720))
    idx = 0
    while True:
        ret, Xt_raw = video_capture.read()
        if ret:
            idx += 1
            if 1:
                Xt_raw_w = copy.deepcopy(Xt_raw)

                Xs = detector.align(Image.fromarray(Xs_raw), crop_size=(256, 256))
                Xt = detector.align(Image.fromarray(Xt_raw), crop_size=(256, 256))

                if Xt is None:
                    Xt_raw_w = cv2.resize(Xt_raw_w, (1280,720), interpolation = cv2.INTER_LINEAR)

                    Xt_raw_w2 = copy.deepcopy(Xt_raw_w)
                    Xt_raw_w2[0:256,0:256,:] = Xs_raw_r
                    # cv2.namedWindow('raw',0)
                    # cv2.imshow('raw', Xt_raw_w)
                    # cv2.namedWindow('fusion',0)
                    # cv2.imshow('fusion', merge)

                    img_stack = np.hstack([Xt_raw_w,Xt_raw_w2])
                    video_writer.write(img_stack)

                    cv2.namedWindow('fusion',0)
                    cv2.imshow('fusion', img_stack)
                    cv2.waitKey(1)

                    continue

                # print(Xs)
                # print(Xt)

                # Xs_raw = np.array(Xs)
                Xs = test_transform(Xs)
                Xs = Xs.unsqueeze(0).cuda()

                Xt_raw = np.array(Xt)
                Xt = test_transform(Xt)
                Xt = Xt.unsqueeze(0).cuda()
                with torch.no_grad():
                    embeds, _ = arcface(F.interpolate(Xs[:, :, 19:237, 19:237], (112, 112), mode='bilinear', align_corners=True))

                mask = np.zeros([256, 256], dtype=np.float)
                for i in range(256):
                    for j in range(256):
                        dist = np.sqrt((i-64)**2 + (j-64)**2)/64
                        dist = np.minimum(dist, 1)
                        mask[i, j] = 1.-dist/8.99
                mask = cv2.dilate(mask, None, iterations=20)

                #---------------------------------------------------------------
                gray = cv2.cvtColor(Xt_raw_w, cv2.COLOR_BGR2GRAY)
                dets = detector_dlib(gray, 1)

                #---------------------------------------------------------------

                Xt_img = Image.fromarray(Xt_raw_w)

                Xt, trans_inv = detector.align(Xt_img, crop_size=(256, 256), return_trans_inv=True)

                Xt = test_transform(Xt)

                Xt = Xt.unsqueeze(0)
                with torch.no_grad():
                    st = time.time()
                    Yt, _ = G(Xt.cuda(), embeds.cuda())
                    Yt = Yt.squeeze().detach().cpu().numpy()
                    st = time.time() - st
                    print(idx,') ',f'inference time: {st} sec')
                    Yt = Yt.transpose([1, 2, 0])*0.5 + 0.5
                    Yt_trans_inv = cv2.warpAffine(Yt, trans_inv, (np.size(Xt_raw_w, 1), np.size(Xt_raw_w, 0)), borderMode=cv2.BORDER_TRANSPARENT)
                    mask_ = cv2.warpAffine(mask, trans_inv, (np.size(Xt_raw_w, 1), np.size(Xt_raw_w, 0)), borderMode=cv2.BORDER_TRANSPARENT)
                    mask_ = np.expand_dims(mask_, 2)

                #-----------------------------------------------------------
                if False:
                    for face in dets:
                        shape = predictor(Xt_raw_w, face)  # 寻找人脸的68个标定点
                        # 遍历所有点，打印出其坐标，并圈出来
                        f_idx  = 0
                        #
                        pts_cc = []
                        for pt in shape.parts():
                            pt_pos = (pt.x, pt.y)

                            if f_idx <=26:
                                pts_cc.append([int(pt.x),int(pt.y)])

                            # if f_idx<=16:
                            #     cv2.circle(Xt_raw_w, pt_pos, 2, (255, 0, 0), -1)
                            #
                            # elif 21>=f_idx>=17:
                            #     cv2.circle(Xt_raw_w, pt_pos, 2, (0, 0, 255), -1)
                            # elif 26>=f_idx>=22:
                            #     cv2.circle(Xt_raw_w, pt_pos, 2, (0, 255, 255), -1)
                            # else :
                            #     cv2.circle(Xt_raw_w, pt_pos, 1, (0, 255, 0), -1)

                            f_idx += 1
                    #-----------------------------------------------------------
                    points_array = np.zeros((1,27,2),dtype = np.int32)
                    for k in range(len(pts_cc)):
                        if k <=16:
                            points_array[0,k,0] = pts_cc[k][0]
                            points_array[0,k,1] = pts_cc[k][1]

                    points_array[0,17] = pts_cc[26]
                    points_array[0,18] = (pts_cc[25][0],pts_cc[25][1]-27)
                    points_array[0,19] = (pts_cc[24][0],pts_cc[24][1]-39)
                    points_array[0,20] = (pts_cc[23][0],pts_cc[23][1]-42)
                    points_array[0,21] = (pts_cc[22][0],pts_cc[22][1]-43)
                    points_array[0,22] = (pts_cc[21][0],pts_cc[21][1]-43)
                    points_array[0,23] = (pts_cc[20][0],pts_cc[20][1]-42)
                    points_array[0,24] = (pts_cc[19][0],pts_cc[19][1]-39)
                    points_array[0,25] = (pts_cc[18][0],pts_cc[18][1]-27)
                    points_array[0,26] = pts_cc[17]

                    # cv2.drawContours(Xt_raw_w,points_array,-1,(255,255,0),thickness=-1)
                    mask_roi = np.zeros([Xt_raw_w.shape[0], Xt_raw_w.shape[1]], dtype=np.float)
                    cv2.drawContours(mask_roi,points_array,-1,(1.),thickness=-1)
                    mask_roi = cv2.dilate(mask_roi, None, iterations=12)
                    cv2.imshow('mask_roi',mask_roi)

                    mask_ = np.expand_dims(mask_roi, 2)*mask_

                #----------------------------------------------------------------------------------

                Yt_trans_inv = mask_*Yt_trans_inv + (1-mask_)*(Xt_raw_w.astype(np.float)/255.)
                cv2.imshow('mask',mask_)
                merge = Yt_trans_inv

                #----------------------------------------------------------------------------------

                Xt_raw_w = cv2.resize(Xt_raw_w, (1280,720), interpolation = cv2.INTER_LINEAR)

                merge = cv2.resize(merge, (1280,720), interpolation = cv2.INTER_LINEAR)
                merge = (merge*255).astype(np.uint8)
                merge[0:256,0:256,:] = Xs_raw_r
                # cv2.namedWindow('raw',0)
                # cv2.imshow('raw', Xt_raw_w)
                # cv2.namedWindow('fusion',0)
                # cv2.imshow('fusion', merge)

                img_stack = np.hstack([Xt_raw_w,merge])
                video_writer.write(img_stack)

                cv2.namedWindow('fusion',0)
                cv2.imshow('fusion', img_stack)

                if cv2.waitKey(1) == 27:
                    break
            # except:
            #     print('error~')
            #     continue
    cv2.destroyAllWindows()
    video_writer.release()


import os
import sys

import numpy as np
import tensorflow as tf
from tensorflow.python.framework import ops
from tensorflow.contrib import layers

from enum import Enum

_bin_path = os.path.join(os.path.dirname(os.path.realpath(__file__)), 'bin')
LD_LIB = os.environ['LD_LIBRARY_PATH']
if str(LD_LIB).find(_bin_path) == -1:
    os.environ['LD_LIBRARY_PATH'] += os.pathsep + _bin_path
    os.execv(sys.executable, [sys.executable] + sys.argv)
_tf_graphics_module = tf.load_op_library(os.path.join(_bin_path, 'libtfgraphics.so'))

_TsdfViaVox = _tf_graphics_module.tsdf_via_vox

_FeatFusionVox = _tf_graphics_module.feat_fusion_vox
_FeatFusionVoxGrad = _tf_graphics_module.feat_fusion_vox_grad

_BilinearInterpolation = _tf_graphics_module.bilinear_interpolation
_BilinearInterpolationGrad = _tf_graphics_module.bilinear_interpolation_grad

_ForwardProjection = _tf_graphics_module.forward_projection
_ReverseProjection = _tf_graphics_module.reverse_projection
_InterpolationWeights = _tf_graphics_module.interpolation_weights
_CameraToImage = _tf_graphics_module.camera_to_image

class FusionType(Enum):
    MAX_POOLING = 1
    AVERAGE_POOLING = 2

def tsdf_via_vox(depth, scene_info, vox_map, is_flipped):
    """
    Convert a vox into tsdf format
    :param depth: the depth info
    :param scene_info: the scene info, including voxel resolution and image viewpoint
    :param vox_map: the non-empty voxel
    :param is_flipped: whether to use the flipped signal
    :type depth: tf.Tensor
    :type scene_info: tf.Tensor
    :type vox_map: tf.Tensor
    :type is_flipped: bool
    :return: tf.Tensor
    """
    if vox_map.dtype is not tf.float32:
        vox_map = tf.cast(vox_map, tf.float32)
    return _TsdfViaVox(depth, scene_info, vox_map, is_flipped)

@ops.RegisterGradient('FeatFusionVox')
def feat_fusion_vox_grad(op, vox_grad):
    """
    Gradient operation of feat_fusion_vox
    :param op: info of forward ops
    :param vox_grad: the vox features
    :type op: tf.Operation
    :type vox_grad: tf.Tensor
    :return: tf.Tensor
    """
    img_feat, img_map, img_weights, vox_weights = op.inputs
    vox_feat = op.outputs[0]
    img_grad, _ = _FeatFusionVoxGrad(img_feat, vox_grad, vox_feat, img_map, img_weights, vox_weights,
                                     op.get_attr('fusion_type'))
    return [img_grad, None, None, None]

def feat_fusion_vox(img_feat, img_map, img_weights, vox_weights, fusion_type=FusionType.MAX_POOLING):
    """
    Fusion the image features into voxel
    :param img_feat: the image features
    :param img_map: the cast relationship between voxel and image
    :param img_weights: the weights blended for volume fusion
    :param vox_weights: the volume weights for normalization
    :param fusion_type: the fusion type to deal with multiple corresponding
    :type img_feat: tf.Tensor
    :type img_map: tf.Tensor
    :type img_weights: tf.Tensor
    :type vox_weights: tf.Tensor
    :type fusion_type: FusionType
    :return: tf.Tensor
    """
    return _FeatFusionVox(img_feat, img_map, img_weights, vox_weights,
                          0 if fusion_type == FusionType.MAX_POOLING else 1)

def feat_projection(depth, img_feat, scene_info, vox_size, fusion_type=FusionType.AVERAGE_POOLING):
    img_vox_map, _, vox_occupied = img2vol_forward_projection(depth, scene_info, vox_size)
    weights = tf.constant(1, dtype=tf.float32, shape=depth.shape)
    vox_occupied = tf.cast(vox_occupied, tf.float32)
    vox_feat = feat_fusion_vox(img_feat, img_vox_map, weights, vox_occupied, fusion_type)
    return vox_feat

def tsdf_projection(depth, scene_info, vox_size, flipped_tsdf=True):
    _, _, vox_weights = img2vol_forward_projection(depth, scene_info, vox_size)
    tsdf, _ = tsdf_via_vox(depth, scene_info, vox_weights, flipped_tsdf)
    return tsdf

def feat_fusion(depth, img_feat, scene_info, vox_size, fusion_type=FusionType.MAX_POOLING, flipped_tsdf=True):
    """
    The packaged fusion operation, including projection, fusion, tsdf
    :param depth: the depth info
    :param img_feat: the image features
    :param scene_info: the scene info, including voxel details and viewpoint
    :param vox_size: the vox size
    :param fusion_type: the fusion type, max or average
    :param flipped_tsdf: whether to flip the signal or not
    :type depth: tf.Tensor
    :type img_feat: tf.Tensor or tf.Variable
    :type scene_info: tf.Tensor
    :type vox_size: list(int)
    :type fusion_type: FusionType
    :type flipped_tsdf: bool
    :return: list(tf.Tensor)
    """
    img_map, _, vox_occupied = img2vol_forward_projection(depth, scene_info, vox_size)
    tsdf_feat, _ = tsdf_via_vox(depth, scene_info, vox_occupied, flipped_tsdf)
    weights = tf.constant(1, dtype=tf.float32, shape=depth.shape)
    vox_occupied = tf.cast(vox_occupied, tf.float32)
    vox_feat = feat_fusion_vox(img_feat, img_map, weights, vox_occupied, fusion_type)
    return vox_feat, tsdf_feat, vox_occupied, img_map

def details_and_fov(img_height, img_width, img_scale, vox_scale):
    vox_details = np.array([0.02 * vox_scale, 0.24], np.float32)
    camera_fov = np.array([518.8579 / img_scale, 0., img_width / (2 * img_scale),
                           0., 518.8579 / img_scale, img_height / (2 * img_scale),
                           0., 0., 1.], dtype=np.float32)
    return vox_details, camera_fov

def scene_base_info(img_scale=1., vox_scale=0., batch_size=0):
    """
    Get the basic info of scene
    :param img_scale: the image re-scale factor, base size is [480, 640]
    :param batch_size: the batch size of scene, if batch_size=0, will return the non-batch result
    :param vox_scale: the vox scale factor, comparing with 240/144/240
    :return:
    """
    if vox_scale == 0.:
        vox_scale = img_scale
    img_height, img_width = [480, 640]
    vox_details_np, camera_fov_np = details_and_fov(img_height, img_width, img_scale, vox_scale)
    vox_details = tf.constant(vox_details_np, tf.float32, [2, ], name='vox_details')
    camera_fov = tf.constant(camera_fov_np, tf.float32, [9, ], name='camera_fov')

    if batch_size != 0:
        vox_details = tf.stack([vox_details] * batch_size, axis=0)
        camera_fov = tf.stack([camera_fov] * batch_size, axis=0)

    return vox_details, camera_fov

def bilinear_interpolation(sampled_feat, interp_points, interp_weights):
    """
    Bilinear interpolation op
    :param sampled_feat:
    :param interp_points:
    :param interp_weights:
    :return:
    """
    return _BilinearInterpolation(sampled_feat, interp_points, interp_weights)

@ops.RegisterGradient('BilinearInterpolation')
def bilinear_interpolation_grad(op, interp_grad):
    """
    Bilinear interpolation gradient op
    :param op: the operation of bilinear interpolation
    :param interp_grad: the gradient from
    :type op: tf.Operation
    :type interp_grad: tf.Tensor or tf.Variable
    :return:
    """
    sampled_feat, interp_points, interp_weights = op.inputs
    sampled_size = [int(axis_size) for axis_size in sampled_feat.shape[1:-1]]
    return [_BilinearInterpolationGrad(interp_grad, interp_points, interp_weights, sampled_size), None, None]

def interpolation_flooring_2d(sampled_feat, interp_size):
    """
    Using bilinear interpolation to sample an image into higher resolution, with flooring sampling scheme
    :param sampled_feat: the feature map to be sampled
    :param interp_size: the interpolated feature map size
    :type sampled_feat: tf.Tensor or tf.Variable
    :type interp_size: list(int)
    :return: tf.Tensor
    """
    sampled_size = [int(axis_size) for axis_size in sampled_feat.shape[1:-1]]
    sampling_ratio = np.array(interp_size) / np.array(sampled_size)
    interp_map = np.mgrid[0:sampling_ratio[0]:interp_size[0], 0:sampling_ratio[1]:interp_size[1]]
    indices_map = np.stack([interp_map] * int(sampled_feat.shape[0]), axis=0)
    interp_points = tf.Constant(indices_map, dtype=tf.float32)
    interp_weights = tf.Constant(np.floor(indices_map), dtype=tf.float32)
    return _BilinearInterpolation(sampled_feat, interp_points, interp_weights)

def interpolation_average_2d(sampled_feat, interp_size):
    """
    Using bilinear interpolation to sample an image into higher resolution, with average sampling scheme
    :param sampled_feat: the feature map to be sampled
    :param interp_size: the interpolated feature map size
    :type sampled_feat: tf.Tensor or tf.Variable
    :type interp_size: list(int)
    :return: tf.Tensor
    """
    sampled_size = [int(axis_size) for axis_size in sampled_feat.shape[1:-1]]
    sampling_ratio = np.array(interp_size) / np.array(sampled_size)
    interp_map = np.mgrid[0:sampling_ratio[0]:interp_size[0], 0:sampling_ratio[1]:interp_size[1]]
    indices_map = np.stack([interp_map] * int(sampled_feat.shape[0]), axis=0)
    interp_points = tf.Constant(indices_map, dtype=tf.float32)
    interp_weights = tf.Constant(np.full(indices_map.shape, 0.25), dtype=tf.float32)
    return _BilinearInterpolation(sampled_feat, interp_points, interp_weights)

def interpolation_bilinear_2d(sampled_feat, interp_size, reuse):
    """
    Using standard bilinear interpolation to sample an image into higher resolution, width bilinear sampling scheme
    :param sampled_feat: the feature map to be sampled
    :param interp_size: the interpolated feature map size
    :param reuse: whether to reuse the weight within the operation
    :type sampled_feat: tf.Tensor or tf.Variable
    :type interp_size: list(int)
    :type reuse: bool
    :return: tf.Tensor
    """
    with tf.variable_scope('BilinearInterpolation2D', reuse=reuse):
        batch_size = int(sampled_feat.shape[0])
        channels = int(sampled_feat.shape[-1])
        sampled_size = [int(axis_size) for axis_size in sampled_feat.shape[1:-1]]

        interp_x, interp_y = np.meshgrid(np.arange(0, interp_size[1]), np.arange(0, interp_size[0]))
        interp_x = interp_x * (sampled_size[1] - 1) / (interp_size[1] - 1)
        interp_y = interp_y * (sampled_size[0] - 1) / (interp_size[0] - 1)
        interp_points = np.stack([interp_y, interp_x], axis=-1)
        interp_points = np.stack([interp_points] * batch_size, axis=0)
        interp_points = np.reshape(interp_points, [batch_size, -1, 2])
        weights_tl = (np.floor(interp_x + 1) - interp_x) * (np.floor(interp_y + 1) - interp_y)
        weights_tr = (interp_x - np.floor(interp_x)) * (np.floor(interp_y + 1) - interp_y)
        weights_bl = (np.floor(interp_x + 1) - interp_x) * (interp_y - np.floor(interp_y))
        weights_br = (interp_x - np.floor(interp_x)) * (interp_y - np.floor(interp_y))
        interp_weights = np.stack([weights_tl, weights_tr, weights_bl, weights_br], axis=-1)
        interp_weights = np.stack([interp_weights] * batch_size, axis=0)
        interp_weights = np.reshape(interp_weights, [batch_size, -1, 4])

        interp_feat = _BilinearInterpolation(sampled_feat, interp_points, interp_weights)
        return tf.reshape(interp_feat, [batch_size, ] + interp_size + [channels, ])

def img_to_cam_space(depth, cam_k):
    """
    Generate the corresponding xyz position of camera coordinate via depth map and camera focus info
    :param depth: the depth map
    :param cam_k: the camera focus information
    :type depth: tf.Tensor
    :type cam_k: tf.Tensor
    :return: tf.Tensor
    """
    depth_size = [int(item) for item in depth.shape]
    img_size = depth_size[1:-1]
    y_axis = np.arange(0, img_size[0])
    x_axis = np.arange(0, img_size[1])
    x, y = np.meshgrid(x_axis, y_axis)

    img_pos_x = tf.constant(np.stack([x] * depth_size[0], axis=0), tf.float32, depth_size[:-1] + [1, ])
    img_pos_y = tf.constant(np.stack([y] * depth_size[0], axis=0), tf.float32, depth_size[:-1] + [1, ])

    cam_k0 = tf.reshape(tf.gather(cam_k, [0, ], axis=-1), [-1, 1, 1, 1])
    cam_k2 = tf.reshape(tf.gather(cam_k, [2, ], axis=-1), [-1, 1, 1, 1])
    cam_k4 = tf.reshape(tf.gather(cam_k, [4, ], axis=-1), [-1, 1, 1, 1])
    cam_k5 = tf.reshape(tf.gather(cam_k, [5, ], axis=-1), [-1, 1, 1, 1])

    cam_pos_x = (img_pos_x - cam_k2) * depth / cam_k0
    cam_pos_y = (img_pos_y - cam_k5) * depth / cam_k4

    return tf.concat([cam_pos_x, cam_pos_y, depth], axis=-1)

def forward_projection(proj_points, camera_pose, vox_unit, vox_origin, vox_size):
    """
    Forward projection sampling scheme, cast image samples into volume space
    :param proj_points: the points set needed to project into volume space
    :param camera_pose: the transform matrix from camera space into volume space
    :param vox_unit: the vox stride for each volume grid
    :param vox_origin: the origin coordinate of the volume
    :param vox_size: the vox size
    :type proj_points: tf.Tensor or tf.Variable
    :type camera_pose: tf.Tensor or tf.Variable
    :type vox_unit: tf.Tensor or tf.Variable
    :type vox_origin: tf.Tensor or tf.Variable
    :type vox_size: list(int)
    :return: list(tf.Tensor)
    """
    return _ForwardProjection(proj_points, camera_pose, vox_unit, vox_origin, vox_size)

def img2vol_forward_projection(depth, scene_info, vox_size):
    vox_origin, cam_pose, vox_unit, cam_k = tf.split(scene_info, [3, 16, 2, 9], axis=-1)
    img_proj_pos = img_to_cam_space(depth, cam_k)
    proj_points = tf.reshape(img_proj_pos, [int(depth.shape[0]), -1, 3])
    point_vox_map, vox_occupied = forward_projection(proj_points, cam_pose, vox_unit, vox_origin, vox_size)
    point_vox_map = tf.reshape(point_vox_map, [int(depth.shape[0]), int(depth.shape[1]), int(depth.shape[2]), -1])
    return point_vox_map, img_proj_pos, vox_occupied

def reverse_projection(cam_pose, vox_unit, vox_origin, vox_size):
    """
    Reverse projection sampling scheme, cast volume samples into image space
    :param cam_pose: the transform matrix from camera space into volume space
    :param vox_unit: the vox stride for each volume grid
    :param vox_origin: the origin coordinate of the volume
    :param vox_size: the vox size
    :type cam_pose: tf.Tensor or tf.Variable
    :type vox_unit: tf.Tensor or tf.Variable
    :type vox_origin: tf.Tensor or tf.Variable
    :type vox_size: list(int)
    :return: list(tf.Tensor)
    """
    cam_pose = tf.reshape(cam_pose, shape=[-1, 4, 4])
    inv_camera_pose = tf.matrix_inverse(cam_pose)
    return _ReverseProjection(inv_camera_pose, vox_unit, vox_origin, vox_size)

def interpolation_weights(interp_points, vox_proj_pos, img_proj_pos, max_distance=0.02):
    """
    Using spatial distance to blend weights
    :param interp_points: the interpolated positions of image space
    :param vox_proj_pos: the spatial position of interpolated volumes
    :param img_proj_pos: the spatial position of sampled images
    :param max_distance: max distance for the valid points
    :type interp_points: tf.Tensor or tf.Variable
    :type img_proj_pos: tf.Tensor or tf.Variable
    :type max_distance: float
    :return: tf.Tensor
    """
    return _InterpolationWeights(interp_points, vox_proj_pos, img_proj_pos, max_distance)

def camera_to_image(vox_proj_pos, scene_info, scale):
    """
    Convert camera space x,y into image space pixel idx
    :param vox_proj_pos: the camera space position
    :param scene_info: the scene information
    :param scale: the scale factor, to scale up the feature maps
    :type vox_proj_pos: tf.Tensor or tf.Variable
    :type scene_info: tf.Tensor or tf.Variable
    :type scale: float
    :return: tf.Tensor
    """
    return _CameraToImage(vox_proj_pos, scene_info, scale)

def view_volume_reverse_projection(depth, img_feat, scene_info, vox_size, multi_scale=False, test_mode=False):
    """
    View volume projection -- using reverse projection
    :param depth: the depth map
    :param img_feat: the 2d features map
    :param scene_info: the scene information, including volume and transform matrix
    :param vox_size: the projected volume size
    :param multi_scale: using multiple scale features maps, based on the volume and image relationship
    :param test_mode: if in test mode, the function will return more intermediary results.
    :type depth: tf.Tensor or tf.Variable
    :type img_feat: tf.Tensor or tf.Variable
    :type scene_info: tf.Tensor or tf.Variable
    :type vox_size: list(int)
    :type multi_scale: bool
    :type test_mode: bool
    :return: tf.Tensor
    """
    if multi_scale:
        raise NotImplementedError
    vox_size = [int(size) for size in vox_size]
    batch_size = int(depth.shape[0])
    img_scale = int(int(depth.shape[-2]) / int(img_feat.shape[-2]))
    depth_batches, img_feat_batches, scene_info_batches = [tf.split(item, batch_size, axis=0)
                                                           for item in [depth, img_feat, scene_info]]
    feat_3d_batches = []
    intermediate = [[], [], [], []]
    for depth_batch, img_feat_batch, scene_info_batch in zip(depth_batches, img_feat_batches, scene_info_batches):
        _, img_proj_pos, vox_occupied = img2vol_forward_projection(depth_batch, scene_info_batch, vox_size)
        img_proj_pos = layers.avg_pool2d(img_proj_pos, img_scale, img_scale)

        vox_origin, cam_pose, vox_unit, _ = tf.split(scene_info_batch, [3, 16, 2, 9], axis=-1)
        vox_proj_pos = reverse_projection(cam_pose, vox_unit, vox_origin, vox_size)

        vox_occupied_indices = tf.where(tf.reshape(vox_occupied, shape=vox_occupied.shape[:-1]) > 0)
        vox_proj_pos = tf.expand_dims(tf.gather_nd(vox_proj_pos, vox_occupied_indices), axis=0)
        intermediate[0].append(vox_proj_pos)

        interp_points = camera_to_image(vox_proj_pos, scene_info_batch, img_scale)
        intermediate[1].append(interp_points)
        interp_weights, interp_distance = interpolation_weights(interp_points, vox_proj_pos, img_proj_pos,
                                                                max_distance=0.04 * img_scale)
        intermediate[2].append(interp_weights)
        intermediate[3].append(interp_distance)
        interp_feat = bilinear_interpolation(img_feat_batch, interp_points, interp_weights)
        interp_feat = tf.reshape(interp_feat, shape=[-1, int(interp_feat.shape[-1])])

        feat_3d_shape = tf.constant([1, ] + vox_size + [int(img_feat.shape[-1]), ], dtype=tf.int64)
        feat_3d_batch = tf.scatter_nd(vox_occupied_indices, interp_feat, feat_3d_shape)
        feat_3d_batches.append(feat_3d_batch)
    feat_3d = tf.concat(feat_3d_batches, axis=0)
    if test_mode:
        return feat_3d, intermediate
    else:
        return feat_3d


import numpy
import pytest
import io

import cellprofiler_core.image
import cellprofiler_core.measurement

import cellprofiler_core.modules.injectimage
import cellprofiler_core.object
import cellprofiler_core.pipeline
import cellprofiler_core.workspace
from cellprofiler_core.constants.measurement import COLTYPE_FLOAT, COLTYPE_INTEGER
from cellprofiler_core.utilities.core import modules as cpmodules

IMAGE_NAME = "image"
OUTPUT_IMAGE_F = "outputimage%d"

N_CHANNELS = 4

import plugins.measureimageintensitymultichannel as mimc

@pytest.fixture(scope="function")
def image():
    return cellprofiler_core.image.Image()

@pytest.fixture(scope="function")
def objects(image):
    objects = cellprofiler_core.object.Objects()

    objects.parent_image = image

    return objects

@pytest.fixture(scope="function")
def measurements():
    return cellprofiler_core.measurement.Measurements()

@pytest.fixture(scope="function")
def module():
    module = mimc.MeasureImageIntensityMultiChannel()

    module.images_list.value = "image"

    module.objects_list.value = "objects"

    module.nchannels.value = N_CHANNELS

    return module

@pytest.fixture(scope="function")
def workspace(image, measurements, module, objects):
    image_set_list = cellprofiler_core.image.ImageSetList()

    image_set = image_set_list.get_image_set(0)

    image_set.add("image", image)

    object_set = cellprofiler_core.object.ObjectSet()

    object_set.add_objects(objects, "objects")

    return cellprofiler_core.workspace.Workspace(
        cellprofiler_core.pipeline.Pipeline(),
        module,
        image_set,
        object_set,
        measurements,
        image_set_list,
    )

def test_init():
    x = mimc.MeasureImageIntensityMultiChannel()

def test_zeros(image, measurements, module, workspace):
    """Test operation on a completely-masked image"""
    image.pixel_data = numpy.zeros((10, 10, N_CHANNELS))

    image.mask = numpy.zeros((10, 10), bool)

    module.run(workspace)

    for i in range(N_CHANNELS):
        assert (
            measurements.get_current_measurement(
                "Image", f"Intensity_TotalArea_image_c{i+1}"
            )
            == 0
        )

    assert len(measurements.get_object_names()) == 1

    assert measurements.get_object_names()[0] == "Image"

    columns = module.get_measurement_columns(workspace.pipeline)

    features = measurements.get_feature_names("Image")

    assert len(columns) == len(features)

    for column in columns:
        assert column[1] in features

def test_image(image, measurements, module, workspace):
    """Test operation on a single unmasked image"""
    numpy.random.seed(0)

    pixels = (
        numpy.random.uniform(size=(10, 10, N_CHANNELS)).astype(numpy.float32) * 0.99
    )

    pixels[0:2, 0:2] = 1

    image.pixel_data = pixels

    module.run(workspace)

    for c in range(N_CHANNELS):
        assert (
            measurements.get_current_measurement(
                "Image", f"Intensity_TotalArea_image_c{c+1}"
            )
            == 100
        )

        assert measurements.get_current_measurement(
            "Image", f"Intensity_TotalIntensity_image_c{c+1}"
        ) == numpy.sum(pixels[:, :, c])

        assert (
            measurements.get_current_measurement(
                "Image", f"Intensity_MeanIntensity_image_c{c+1}"
            )
            == numpy.sum(pixels[:, :, c]) / 100.0
        )

        assert measurements.get_current_image_measurement(
            f"Intensity_MinIntensity_image_c{c+1}"
        ) == numpy.min(pixels[:, :, c])

        assert measurements.get_current_image_measurement(
            f"Intensity_MaxIntensity_image_c{c+1}"
        ) == numpy.max(pixels[:, :, c])

        assert (
            measurements.get_current_image_measurement(
                f"Intensity_PercentMaximal_image_c{c+1}"
            )
            == 4.0
        )

def test_image_and_mask(image, measurements, module, workspace):
    """Test operation on a masked image"""
    numpy.random.seed(0)

    pixels = (
        numpy.random.uniform(size=(10, 10, N_CHANNELS)).astype(numpy.float32) * 0.99
    )

    pixels[1:3, 1:3] = 1

    mask = numpy.zeros((10, 10), bool)

    mask[1:9, 1:9] = True

    image.pixel_data = pixels

    image.mask = mask

    module.run(workspace)

    for c in range(N_CHANNELS):
        assert (
            measurements.get_current_measurement(
                "Image", f"Intensity_TotalArea_image_c{c+1}"
            )
            == 64
        )

        assert measurements.get_current_measurement(
            "Image", f"Intensity_TotalIntensity_image_c{c+1}"
        ) == numpy.sum(pixels[1:9, 1:9, c])

        assert (
            measurements.get_current_measurement(
                "Image", f"Intensity_MeanIntensity_image_c{c+1}"
            )
            == numpy.sum(pixels[1:9, 1:9, c]) / 64.0
        )

        assert (
            measurements.get_current_measurement(
                "Image", f"Intensity_PercentMaximal_image_c{c+1}"
            )
            == 400.0 / 64.0
        )

def test_image_and_objects(image, measurements, module, objects, workspace):
    """Test operation on an image masked by objects"""
    numpy.random.seed(0)

    pixels = (
        numpy.random.uniform(size=(10, 10, N_CHANNELS)).astype(numpy.float32) * 0.99
    )

    pixels[1:3, 1:3] = 1

    image.pixel_data = pixels

    labels = numpy.zeros((10, 10), int)

    labels[1:9, 1:5] = 1

    labels[1:9, 5:9] = 2

    objects.segmented = labels

    module.wants_objects.value = True

    module.run(workspace)

    for c in range(N_CHANNELS):
        assert (
            measurements.get_current_measurement(
                "Image", f"Intensity_TotalArea_image_objects_c{c+1}"
            )
            == 64
        )

        assert measurements.get_current_measurement(
            "Image", f"Intensity_TotalIntensity_image_objects_c{c+1}"
        ) == numpy.sum(pixels[1:9, 1:9, c])

        assert (
            measurements.get_current_measurement(
                "Image", f"Intensity_MeanIntensity_image_objects_c{c+1}"
            )
            == numpy.sum(pixels[1:9, 1:9, c]) / 64.0
        )

        assert (
            measurements.get_current_measurement(
                "Image", f"Intensity_PercentMaximal_image_objects_c{c+1}"
            )
            == 400.0 / 64.0
        )

    assert len(measurements.get_object_names()) == 1

    assert measurements.get_object_names()[0] == "Image"

    columns = module.get_measurement_columns(workspace.pipeline)

    features = measurements.get_feature_names("Image")

    assert len(columns) == len(features)

    for column in columns:
        assert column[1] in features

def test_image_and_objects_and_mask(image, measurements, module, objects, workspace):
    """Test operation on an image masked by objects and a mask"""
    numpy.random.seed(0)

    pixels = numpy.random.uniform(size=(10, 10, N_CHANNELS)).astype(numpy.float32)

    mask = numpy.zeros((10, 10), bool)

    mask[1:9, :9] = True

    image.pixel_data = pixels

    image.mask = mask

    labels = numpy.zeros((10, 10), int)

    labels[1:9, 1:5] = 1

    labels[1:9, 5:] = 2

    objects.segmented = labels

    module.wants_objects.value = True

    module.run(workspace)

    for c in range(N_CHANNELS):
        assert (
            measurements.get_current_measurement(
                "Image", f"Intensity_TotalArea_image_objects_c{c+1}"
            )
            == 64
        )

        assert measurements.get_current_measurement(
            "Image", f"Intensity_TotalIntensity_image_objects_c{c+1}"
        ) == numpy.sum(pixels[1:9, 1:9, c])

        assert (
            measurements.get_current_measurement(
                "Image", f"Intensity_MeanIntensity_image_objects_c{c+1}"
            )
            == numpy.sum(pixels[1:9, 1:9, c]) / 64.0
        )

def test_get_measurement_columns_whole_image_mode(module):
    image_names = ["image%d" % i for i in range(3)]

    module.wants_objects.value = False

    expected_suffixes = []

    for image_name in image_names:
        im = module.images_list.value[-1]

        module.images_list.value.append(image_name)

        expected_suffixes.append(image_name)

    columns = module.get_measurement_columns(None)

    assert all([column[0] == "Image" for column in columns])

    for expected_suffix in expected_suffixes:
        for feature, coltype in (
            (
                mimc.F_TOTAL_INTENSITY,
                COLTYPE_FLOAT,
            ),
            (
                mimc.F_MEAN_INTENSITY,
                COLTYPE_FLOAT,
            ),
            (
                mimc.F_MIN_INTENSITY,
                COLTYPE_FLOAT,
            ),
            (
                mimc.F_MAX_INTENSITY,
                COLTYPE_FLOAT,
            ),
            (
                mimc.F_TOTAL_AREA,
                COLTYPE_INTEGER,
            ),
            (
                mimc.F_PERCENT_MAXIMAL,
                COLTYPE_FLOAT,
            ),
            (
                mimc.F_MAD_INTENSITY,
                COLTYPE_FLOAT,
            ),
            (
                mimc.F_LOWER_QUARTILE,
                COLTYPE_FLOAT,
            ),
            (
                mimc.F_UPPER_QUARTILE,
                COLTYPE_FLOAT,
            ),
        ):
            # feature names are now formatting strings
            feature_name = feature % expected_suffix

            assert any(
                [
                    (column[1] == f"{feature_name}_c{c+1}" and column[2] == coltype)
                    for c in range(N_CHANNELS)
                    for column in columns
                ]
            )

def test_get_measurement_columns_object_mode(module):
    image_names = ["image%d" % i for i in range(3)]

    object_names = ["object%d" % i for i in range(3)]

    module.wants_objects.value = True

    expected_suffixes = []

    for image_name in image_names:
        module.images_list.value.append(image_name)

        for object_name in object_names:
            module.objects_list.value.append(object_name)

            expected_suffixes.append("%s_%s" % (image_name, object_name))

    columns = module.get_measurement_columns(None)

    assert all([column[0] == "Image" for column in columns])

    for expected_suffix in expected_suffixes:
        for feature, coltype in (
            (
                mimc.F_TOTAL_INTENSITY,
                COLTYPE_FLOAT,
            ),
            (
                mimc.F_MEAN_INTENSITY,
                COLTYPE_FLOAT,
            ),
            (
                mimc.F_MIN_INTENSITY,
                COLTYPE_FLOAT,
            ),
            (
                mimc.F_MAX_INTENSITY,
                COLTYPE_FLOAT,
            ),
            (
                mimc.F_TOTAL_AREA,
                COLTYPE_INTEGER,
            ),
            (
                mimc.F_PERCENT_MAXIMAL,
                COLTYPE_FLOAT,
            ),
            (
                mimc.F_MAD_INTENSITY,
                COLTYPE_FLOAT,
            ),
            (
                mimc.F_LOWER_QUARTILE,
                COLTYPE_FLOAT,
            ),
            (
                mimc.F_UPPER_QUARTILE,
                COLTYPE_FLOAT,
            ),
        ):
            # feature names are now formatting strings
            feature_name = feature % expected_suffix
            assert any(
                [
                    (column[1] == f"{feature_name}_c{c+1}" and column[2] == coltype)
                    for c in range(N_CHANNELS)
                    for column in columns
                ]
            )