"""
Training code for Adversarial patch training using Faster RCNN based on mmdetection
Redo UPC in PyTorch

"""
import os
os.environ["CUDA_VISIBLE_DEVICES"] = '0,1,2,3'

import PIL
import cv2
# from load_data import *
import copy
from tqdm import tqdm
from mmdet import __version__
from mmdet.apis import init_detector, inference_detector, show_result_pyplot, get_Image_ready
import mmcv
from mmcv.ops import RoIAlign, RoIPool
from mmcv.parallel import collate, scatter
from mmcv.runner import load_checkpoint
from mmcv.parallel import MMDataParallel, MMDistributedDataParallel

import torch
import gc
import matplotlib.pyplot as plt
from torch import autograd
from torchvision import transforms

import subprocess
# from utils.utils import *
import numpy as np
import pickle

# import patch_config as patch_config
import sys
import time

# from brambox.io.parser.annotation import DarknetParser as anno_darknet_parse
# from utils import *
from skimage.segmentation import slic
from skimage.util import img_as_float
from skimage import io
from skimage.segmentation import mark_boundaries
import os
from PIL import Image
import torch.optim as optim
import torch.nn.functional as F
import skimage.io as io
from train_patch_frcn_measure_no_stop_gray_step_3_0920 import measure_region_with_attack
import math
from dataset.coco_train_1000_PERSON import CocoTrainPerson
from dataset.sp_geter_dataset import SuperPixelGet

from torch.utils.data import DataLoader,Dataset
from utils.tv_loss import TVLoss
from utils.iou import compute_iou_tensor

seed = 2020
torch.manual_seed(seed)
torch.cuda.manual_seed(seed)

csv_name = 'x_result2.csv'
# torch.cuda.set_device(0)

def gkern(kernlen=21, nsig=3):
    """Returns a 2D Gaussian kernel array."""
    import scipy.stats as st

    x = np.linspace(-nsig, nsig, kernlen)
    kern1d = st.norm.pdf(x)
    kernel_raw = np.outer(kern1d, kern1d)
    kernel = kernel_raw / kernel_raw.sum()
    return kernel

class PatchTrainer(object):
    def __init__(self, mode):
        # self.config = patch_config.patch_configs[mode]()

        # self.darknet_model = Darknet(self.config.cfgfile)
        # self.darknet_model.load_weights(self.config.weightfile)
        # self.darknet_model = self.darknet_model.eval().cuda() # TODO: Why eval?

        self.ocpy_list = []
        for i in range(4):
            if i ==0:
                ocpy_tensor = torch.rand(500,500,500,3).cuda(i)
                self.ocpy_list.append(ocpy_tensor)
                continue

            ocpy_tensor = torch.rand(500,500,500,6).cuda(i)
            self.ocpy_list.append(ocpy_tensor)

        self.config_file = './configs/faster_rcnn/faster_rcnn_r50_fpn_1x_coco.py'
        self.checkpoint_file = '../common_data/faster_rcnn_r50_fpn_1x_coco_20200130-047c8118.pth'
        self.Faster_RCNN = init_detector(self.config_file, self.checkpoint_file, device='cpu').cuda()

        # self.Faster_RCNN = MMDistributedDataParallel(
        #     self.Faster_RCNN.cuda(),
        #     device_ids=[torch.cuda.current_device()],
        #     broadcast_buffers=False)

        self.Faster_RCNN = torch.nn.DataParallel(self.Faster_RCNN)

        # self.patch_applier = PatchApplier().cuda()
        # self.patch_transformer = PatchTransformer().cuda()
        # self.prob_extractor = MaxProbExtractor(0, 80, self.config).cuda()
        # self.nps_calculator = NPSCalculator(self.config.printfile, self.config.patch_size).cuda()
        # self.total_variation = TotalVariation().cuda()
        self.mean = torch.Tensor([123.675, 116.28, 103.53]).unsqueeze(0).unsqueeze(-1).unsqueeze(-1).cuda()
        self.std = torch.Tensor([58.395, 57.12, 57.375]).unsqueeze(0).unsqueeze(-1).unsqueeze(-1).cuda()

        # self.writer = self.init_tensorboard(mode)

    def Transform_Patch(self, patch):
        # clamp_patch = torch.clamp(patch, 0.01, 254.99)
        clamp_patch = patch
        unsqueezed_patch = clamp_patch.unsqueeze(0)
        resized_patch = F.interpolate(unsqueezed_patch, (800, 800), mode='bilinear').cuda()
        normalized_patch = (resized_patch - self.mean) / self.std
        return normalized_patch

    def Transform_Patch_batch(self, patch):
        clamp_patch = torch.clamp(patch, 0.01, 254.99)
        resized_patch = F.interpolate(clamp_patch, (800, 800), mode='bilinear').cuda()
        normalized_patch = (resized_patch - self.mean) / self.std
        return normalized_patch

    def train(self):
        """
        Optimize a patch to generate an adversarial example.
        :return: Nothing
        """

        img_size = 800
        batch_size = 1
        n_epochs = 5000
        max_lab = 14

        ATTACK_TASK = 'target'

        # TARGET_CLASS = 'dog'
        TARGET_CLASS = 16
        # ATTACK_TASK = 'untarget'

        time_str = time.strftime("%Y%m%d-%H%M%S")

        conv_size = 2
        kernel1 = gkern(2*conv_size+1, 3).astype(np.float32)
        stack_kernel1 = np.stack([kernel1, kernel1, kernel1]).swapaxes(2, 0)
        stack_kernel1 = np.expand_dims(stack_kernel1, 3)
        stack_kernel1 = torch.from_numpy(stack_kernel1).permute(2,3,0,1).float()

        # Dataset prepare

        data_obj = CocoTrainPerson(dataType='train2017',num_use=500)
        dataloader_obj = DataLoader(data_obj, batch_size=1,shuffle=False) #使用DataLoader加载数据

        # img info prepare
        img_frcn = get_Image_ready(self.Faster_RCNN.module, '1016.png')
        img_frcn['img_metas'][0][0]['filename'] = None
        img_frcn['img_metas'][0][0]['ori_filename'] = None
        img_frcn['img_metas'][0][0]['ori_shape'] = None
        img_frcn['img_metas'][0][0]['pad_shape'] = None
        img_frcn['img_metas'][0][0]['scale_factor'] = None

        # attack_area_rate = 0.2
        ATTACK_AREA_RATE = 0.1
        evo_step_num = 20
        batch_size_sp = 12
        population_num = 72 # 36
        optim_step_num = 300
        k = 0
        for i_batch, batch_data in enumerate(dataloader_obj):
            img, mask, bbox, class_label = batch_data[0][0], batch_data[1][0], batch_data[2][0], batch_data[3][0]
            # img  : 3,500,500
            # mask : 500,500
            # bbox : x1,y1,w,h
            # class_label : tensor[]

            img_name = batch_data[4][0]

            print('---------------')
            print(img_name)
            print('---------------')

            save_dir = os.path.join('../common_data/NES_search_test_1107/', img_name.split('.')[0])
            if os.path.exists(save_dir):
                continue
            time0 = time.time()

            # use segment SLIC
            base_SLIC_seed_num = 3000
            img_np = img.numpy().transpose(1,2,0)
            mask_np = mask.numpy()
            numSegments = int(base_SLIC_seed_num/(500*500)*torch.sum(mask))
            segments_np = slic(image=img_np, n_segments=numSegments, sigma=0, slic_zero=True, mask=mask_np)
            segments_tensor = torch.from_numpy(segments_np).float().cuda()
            segments_label = torch.unique(segments_tensor)
            segments_label = segments_label[1:]

            # define theta_m
            # pay attention to the center and the boundary

            # (0) prepare stack of sp
            # (1) find the center sp
            # (2) find the boundary sp

            # (0) prepare stack of sp
            zero_layer = torch.zeros_like(segments_tensor)
            one_layer = torch.ones_like(segments_tensor)
            segments_stack = torch.stack([torch.where(segments_tensor==segments_label[j], segments_tensor, zero_layer) for j in range(segments_label.shape[0])], dim=0)
            segments_stack = segments_stack.cpu()

            mask_area = torch.sum(mask)
            # (1) find the center sp
            bbox_x1 = bbox[0]
            bbox_y1 = bbox[1]
            bbox_w = bbox[2]
            bbox_h = bbox[3]

            bbox_x_c = bbox_x1 + bbox_w/2
            bbox_y_c = bbox_y1 + bbox_h/2
            bbox_x_c_int = int(bbox_x_c)
            bbox_y_c_int = int(bbox_y_c)

            # if segments_tensor[bbox_y_c_int, bbox_x_c_int] == 0:
            #     # no sp in center
            #     center_sp = torch.Tensor().cuda()
            # else:
            #     center_sp = segments_tensor[bbox_y_c_int, bbox_x_c_int].unsqueeze(0)
            #     center_sp_layer = torch.where(segments_tensor==center_sp, one_layer, zero_layer)
            #     center_sp_layer_np = center_sp_layer.cpu().numpy()
            #     kernel = np.ones((3,3),np.uint8)
            #     center_sp_layer_dilate = cv2.dilate(center_sp_layer_np, kernel, iterations = 2)
            #     center_sp_layer_dilate = torch.from_numpy(center_sp_layer_dilate)
            #     center_sp_layer_dilate_stack = center_sp_layer_dilate.unsqueeze(0).repeat(segments_stack.shape[0],1,1)

            #     cross_stack = center_sp_layer_dilate_stack * segments_stack
            #     cross_ = torch.sum(cross_stack, dim=0)
            #     neighborhoods = torch.unique(cross_)[1:]
            #     center_sps = neighborhoods
            #     cross_stack = cross_stack.cpu()
            #     center_sp_layer_dilate_stack = center_sp_layer_dilate_stack.cpu()

            #     # we also need center sp's neighborhoods

            # # (2) find the boundary sp

            # # boundary_erode
            # kernel = np.ones((3,3),np.uint8)
            # mask_erosion = cv2.erode(mask_np, kernel, iterations = 2)
            # boundary_erode = mask_np - mask_erosion
            # boundary_erode = torch.from_numpy(boundary_erode)
            # # boundary_erode_pil = transforms.ToPILImage()(boundary_erode.cpu())
            # # boundary_erode_pil.show()

            # boundary_erode_stack = boundary_erode.unsqueeze(0).repeat(segments_stack.shape[0],1,1)
            # boundary_mul_segments_stack = boundary_erode_stack * segments_stack
            # boundary_mul_segments = torch.sum(boundary_mul_segments_stack, dim=0)
            # # boundary_mul_segments_pil = transforms.ToPILImage()(boundary_mul_segments.cpu())
            # # boundary_mul_segments_pil.show()
            # boundary_mul_segments_unique = torch.unique(boundary_mul_segments)
            # boundary_mul_segments_unique = boundary_mul_segments_unique[1:]
            # boundary_sp = boundary_mul_segments_unique

            # shan dian
            # init grid
            densy = 7
            unit_w = 13*densy
            unit_h = 13*densy
            sandian = torch.zeros(unit_w,unit_h)
            '''
            log:
            10,5,10,5 : 0.04   work! at 700
            10,5,10,6 : 0.0333 work! at 2040
            '''
            sandian = sandian.reshape(13,densy,13,densy)
            sandian[:,int((densy-1)/2),:,int((densy-1)/2)] = 1
            sandian = sandian.reshape(unit_w, unit_h)
            sandian = sandian.unsqueeze(0).unsqueeze(0)
            sandian = F.interpolate(sandian, (500, 500), mode='nearest').squeeze()
            sandian_stack = sandian.unsqueeze(0).repeat(segments_stack.shape[0],1,1)
            sandian_mul_segments_stack = sandian_stack * segments_stack
            sandian_mul_segments = torch.sum(sandian_mul_segments_stack, dim=0)
            sandian_mul_segments_pil = transforms.ToPILImage()(sandian_mul_segments.cpu())
            sandian_mul_segments_pil.show()
            sandian_mul_segments_unique = torch.unique(sandian_mul_segments)
            sandian_mul_segments_unique = sandian_mul_segments_unique[1:]
            sandian_sp = sandian_mul_segments_unique

            # pay attention
            # spot_sp = torch.cat((center_sps, boundary_sp), dim=0)
            spot_sp = sandian_sp
            spot_sp = torch.unique(spot_sp).long()

            sandian_stack = sandian_stack.cpu()
            # boundary_erode_stack = boundary_erode_stack.cpu()
            # boundary_mul_segments_stack = boundary_mul_segments_stack.cpu()
            segments_stack = segments_stack.cpu()

            torch.cuda.empty_cache()

            # show_tensor = img.clone().cuda()
            # for i in range(spot_sp.shape[0]):
            #     sp = spot_sp[i]
            #     show_tensor = torch.where(segments_tensor==sp, zero_layer.fill_(1).unsqueeze(0).repeat(3,1,1),show_tensor)

            # show_tensor_pil = transforms.ToPILImage()(show_tensor.cpu())
            # show_tensor_pil.show()

            # generate theta_m
            # for sp id from 1 to 128
            uniform_ratio = torch.Tensor([ATTACK_AREA_RATE])[0]
            higher_ratio =  uniform_ratio*1.5
            uniform_ratio_theta = 1/2*torch.log(uniform_ratio/(1-uniform_ratio))
            higher_ratio_theta = 1/2*torch.log(higher_ratio/(1-higher_ratio))
            theta_m = torch.zeros_like(segments_label).cpu().fill_(uniform_ratio_theta)
            theta_m[spot_sp-1] = higher_ratio_theta

            for evo_step in range(evo_step_num):
                # prepare sp dataset
                g_theta_m = 1/2*(torch.tanh(theta_m)+1)
                data_sp = SuperPixelGet(segments_label = segments_label,
                                    segments_tensor=segments_tensor,
                                    g_theta_m=g_theta_m,
                                    data_num=population_num)
                dataloader_sp = DataLoader(data_sp, batch_size=batch_size_sp,shuffle=False) #使用DataLoader加载数据

                F_value_restore = torch.Tensor()
                select_m_restore = torch.Tensor()

                for j_sp_batch , sp_batch_data in enumerate(dataloader_sp):

                    attack_mask_batch, select_m = sp_batch_data[0].unsqueeze(1), sp_batch_data[1]
                    select_m_restore = torch.cat((select_m_restore, select_m.cpu()))

                    batch_size_now = attack_mask_batch.shape[0]
                    ## start at gray
                    adv_patch_w_batch = torch.zeros(batch_size_now,3,500,500).cuda()

                    adv_patch_w_batch.requires_grad_(True)

                    ## optimizer and scheduler
                    # optimizer = optim.Adam([
                    #     {'params': adv_patch, 'lr': 0.01*255}
                    # ], amsgrad=True)
                    optimizer = optim.Adam([
                        {'params': adv_patch_w_batch, 'lr': 0.1}
                    ], amsgrad=True)

                    L_value_step = torch.ones(batch_size_now ,optim_step_num)*1000
                    for step in tqdm(range(optim_step_num)):

                        # prepare batch data to feed the frcn
                        img_batch = img.cuda().unsqueeze(0).repeat(batch_size_now,1,1,1)
                        adv_patch_batch = torch.sigmoid(adv_patch_w_batch/2)
                        patched_img_batch = torch.where(attack_mask_batch > 0, adv_patch_batch, img_batch)
                        patched_img_batch_255 = patched_img_batch * 255.
                        patched_img_batch_rsz = F.interpolate(patched_img_batch_255, (800, 800), mode='bilinear').cuda()
                        patched_img_batch_nom_rsz = (patched_img_batch_rsz - self.mean) / self.std

                        # output
                        img_new = copy.deepcopy(img_frcn)
                        img_new['img'][0] = patched_img_batch_nom_rsz
                        # [0] = patched_img_batch_nom_rsz
                        img_new['img_metas'][0] = [img_new['img_metas'][0][0] for i in range(batch_size_now)]
                        frcn_output = self.Faster_RCNN(return_loss=False, rescale=False,  **img_new)
                        # output formate is [x1,y1,x2,y2]

                        # compute loss
                        proposals_4507 = frcn_output[1]
                        proposals_score_4507 = frcn_output[2]
                        det_bboxes, det_labels, proposals = frcn_output[0]

                        det_bboxes = torch.cat(det_bboxes,dim=0)/800*500
                        proposals = torch.cat(proposals,dim=0)/800*500
                        det_labels = torch.cat(det_labels,dim=0)

                        attack_prob = det_labels[:,class_label]
                        training_confidence_threshold = 0.28
                        ov_thrs_index = torch.where(attack_prob > training_confidence_threshold)[0] # for certain class
                        pbbox_attack_cls = det_bboxes[:, class_label*4:(class_label+1)*4]

                        # cls loss
                        attack_class_score = det_labels[:,class_label]
                        top_sort_class_score = torch.sort(attack_class_score, descending=True)[0][:10]
                        cls_loss = torch.sum(top_sort_class_score)

                        # iou loss
                        bbox_x1 = bbox[0]/500*416
                        bbox_y1 = bbox[1]/500*416
                        bbox_w = bbox[2]/500*416
                        bbox_h = bbox[3]/500*416
                        ground_truth_bbox = [bbox_x1, bbox_y1, bbox_x1+bbox_w, bbox_y1 + bbox_h]
                        ground_truth_bbox = torch.Tensor(ground_truth_bbox).unsqueeze(0).cuda()
                        iou_all = compute_iou_tensor(det_bboxes[:, class_label*4:(class_label+1)*4], ground_truth_bbox)
                        iou_positive = iou_all[iou_all>0.15]
                        iou_loss = torch.sum(iou_all)

                        # class loss selected by IoU
                        attack_class_score = det_labels[:,class_label]
                        attack_class_score_iou = attack_class_score[iou_all>0.45]
                        attack_class_score_iou_sort = torch.sort(attack_class_score_iou, descending=True)[0][:30]
                        cls_iou_loss = torch.sum(attack_class_score_iou_sort)

                        final_roi = pbbox_attack_cls[ov_thrs_index]  # for certain class
                        final_roi = final_roi[:,:4]
                        final_ctx = (final_roi[:,0] + final_roi[:,2])/2
                        final_cty = (final_roi[:,1] + final_roi[:,3])/2
                        final_ctx = final_ctx.unsqueeze(-1)
                        final_cty = final_cty.unsqueeze(-1)
                        final_roi_target = torch.cat([final_ctx,final_cty,final_ctx,final_cty], dim=-1)
                        reg_loss = 10* l1_norm(final_roi-final_roi_target) / final_roi.shape[0] /500

                        # RPN loss
                        # r1 : from score
                        # r2 : from x,y,w,h

                        # rpn score target is 0
                        rpn_score = proposals[:, 4]
                        loss_r1 = l2_norm(rpn_score - 0)

                        # rpn box target is smaller the boxes
                        rpn_ctx = (proposals[:,0] + proposals[:,2])/2
                        rpn_cty = (proposals[:,1] + proposals[:,3])/2
                        rpn_box = proposals[:,:4]
                        rpn_ctx = rpn_ctx.unsqueeze(-1)
                        rpn_cty = rpn_cty.unsqueeze(-1)
                        rpn_box_target = torch.cat([rpn_ctx,rpn_cty,rpn_ctx,rpn_cty], dim=-1)
                        # loss_r2 = l1_norm(rpn_score.unsqueeze(-1).repeat(1,4)*(rpn_box - rpn_box_target)) / 500
                        loss_r2 = l1_norm((rpn_box - rpn_box_target)) / 500

                        lambda_balance1 = 0.02
                        # rpn_loss = loss_r1 + lambda_balance1 * loss_r2
                        rpn_loss = lambda_balance1 * loss_r2
                        # rpn_loss = loss_r1

                        total_loss = cls_loss + reg_loss + rpn_loss
                        total_loss = cls_loss + cls_iou_loss + reg_loss + rpn_loss

                        # if epoch > 500:
                        #     total_loss = rpn_loss + cls_loss_new + reg_loss

                        total_loss.backward()
                        optimizer.step()
                        optimizer.zero_grad()

                        # compute F(m,t*;y)

                        # L in disappearing
                        # i think we better use iou>0.45 in prediction, all bbox's person class confidence sum

                        bbox_x1 = bbox[0]
                        bbox_y1 = bbox[1]
                        bbox_w = bbox[2]
                        bbox_h = bbox[3]
                        ground_truth_bbox = [bbox_x1, bbox_y1, bbox_x1+bbox_w, bbox_y1 + bbox_h]
                        ground_truth_bbox = torch.Tensor(ground_truth_bbox).unsqueeze(0).cuda()
                        iou_all = compute_iou_tensor(det_bboxes[:, class_label*4:(class_label+1)*4], ground_truth_bbox)

                        iou_select_index = iou_all > 0.45

                        for ib in range(batch_size_now):

                            det_labels_one = det_labels[ib:(ib+1)*1000, class_label]

                            iou_select_index_one = iou_select_index[ib:(ib+1)*1000]

                            cls_conf_iou_select = det_labels_one[iou_select_index_one]

                            cls_conf_iou_select = cls_conf_iou_select[cls_conf_iou_select>0.25]

                            cls_conf_iou_select_top = torch.sort(cls_conf_iou_select, descending=True)[0][:10]

                            L_value_step[ib, step] = torch.sum(cls_conf_iou_select_top).detach().clone()
                            pass
                        if (torch.min(L_value_step, dim=1)[0] == torch.zeros(batch_size_now)).all():
                            break

                        pass
                    L_value = - torch.min(L_value_step, dim=1)[0]
                    F_value = L_value - torch.sum(torch.sum(sp_batch_data[0], dim=-1), dim=-1).cpu()/mask_area*10 ###!!!!
                    F_value_restore = torch.cat((F_value_restore, F_value))

                    pass
                # print(F_value_restore)
                # now we have F value

                delta_J_theta = 1/population_num * F_value_restore.unsqueeze(1) * 2 * (select_m_restore.cpu().float() - g_theta_m)

                delta_J_theta = torch.sum(delta_J_theta, dim=0)
                theta_m = theta_m + delta_J_theta
                g_theta_m = 1/2*(torch.tanh(theta_m)+1)

                select_sp_index = torch.sort(theta_m)[1][:int(theta_m.shape[0]*ATTACK_AREA_RATE)]

                attack_region_show = zero_layer.clone().squeeze().cuda()
                flag = torch.zeros_like(segments_label)
                flag[select_sp_index] = 1
                for i in range(segments_label.shape[0]):

                    color = g_theta_m[i]
                    color_layer = torch.zeros_like(segments_tensor).fill_(color)
                    sp =  segments_label[i]

                    attack_region_show = torch.where(segments_tensor==sp, color_layer, attack_region_show)
                attack_region_show = attack_region_show / torch.max(attack_region_show)

                attack_region_show = (attack_region_show + img.cuda())/2

                attack_region_show_pil = transforms.ToPILImage()(attack_region_show.cpu())

                save_dir = os.path.join('../common_data/NES_search_test_1107/', img_name.split('.')[0])
                if not os.path.exists(save_dir):
                    os.makedirs(save_dir)
                save_path = os.path.join(save_dir, str(evo_step)+'_pop_'+str(population_num)+'_'+str(base_SLIC_seed_num)+'.png')
                attack_region_show_pil.save(save_path)

                time1 = time.time()

                time_cost = time1 - time0

                if time_cost > 1800:
                    break

                print(g_theta_m)
        print(assasasasasasa)
        pass

    def generate_patch(self, type):
        """
        Generate a random patch as a starting point for optimization.

        :param type: Can be 'gray' or 'random'. Whether or not generate a gray or a random patch.
        :return:
        """
        if type == 'gray':
            adv_patch_cpu = torch.full((3, 500, 500), 0.5)
        elif type == 'random':
            adv_patch_cpu = torch.rand((3, 500, 500))
        if type == 'trained_patch':
            patchfile = 'patches/object_score.png'
            patch_img = Image.open(patchfile).convert('RGB')
            patch_size = self.config.patch_size
            tf = transforms.Resize((patch_size, patch_size))
            patch_img = tf(patch_img)
            tf = transforms.ToTensor()
            adv_patch_cpu = tf(patch_img)

        return adv_patch_cpu

    def read_image(self, path):
        """
        Read an input image to be used as a patch

        :param path: Path to the image to be read.
        :return: Returns the transformed patch as a pytorch Tensor.
        """
        patch_img = Image.open(path).convert('RGB')
        tf = transforms.Resize((self.config.patch_size, self.config.patch_size))
        patch_img = tf(patch_img)
        tf = transforms.ToTensor()

        adv_patch_cpu = tf(patch_img)
        return adv_patch_cpu

def connected_domin_detect(input_img):
    from skimage import measure
    # detection
    if input_img.shape[0] == 3:
        input_img_new = (input_img[0] + input_img[1] + input_img[2])
    else:
        input_img_new = input_img
    ones = torch.Tensor(input_img_new.size()).fill_(1)
    zeros = torch.Tensor(input_img_new.size()).fill_(0)
    input_map_new = torch.where((input_img_new != 0), ones, zeros)
    # img = transforms.ToPILImage()(input_map_new.detach().cpu())
    # img.show()
    input_map_new = input_map_new.cpu()
    labels = measure.label(input_map_new[:, :], background=0, connectivity=2)
    label_max_number = np.max(labels)
    return float(label_max_number)

def get_obj_min_score(boxes):
    if type(boxes[0][0]) is list:
        min_score_list = []
        for i in range(len(boxes)):
            score_list = []
            for j in range(len(boxes[i])):
                score_list.append(boxes[i][j][4])
            min_score_list.append(min(score_list))
        return np.array(min_score_list)
    else:
        score_list = []
        for j in range(len(boxes)):
            score_list.append(boxes[j][4])
        return np.array(min(score_list))

def l2_norm(tensor):
    return torch.sqrt(torch.sum(torch.pow(tensor,2)))

def l1_norm(tensor):
    return torch.sum(torch.abs(tensor))

def main():

    trainer = PatchTrainer('paper_obj')
    trainer.train()

if __name__ == '__main__':
    main()


from __future__ import print_function
from __future__ import division
import torch
import torch.nn as nn
import torch.optim as optim
import numpy as np
import torchvision
from torchvision import datasets, models, transforms
import matplotlib.pyplot as plt
import time
import os
import copy
## import pdb
## import tqdm
## import visdom

def save_checkpoint(state, filename):
    filename_checkpoint = 'checkpoints/'+ filename + '_checkpoint.pth.tar'
    if os.path.isfile(filename_checkpoint):
        filename_checkpoint = 'checkpoints/'+ filename + '_new_checkpoint.pth.tar'
        torch.save(state, filename_checkpoint)
    else:
        torch.save(state, filename_checkpoint)
    print("=> Checkpoint saved at: '{}'".format(filename_checkpoint))

def train_model(name, resume, model, dataloaders, criterion, optimizer, device, num_epochs, is_inception=False):
    """
       Train_model function loads the dataloa der to the choosen model, run train and validation steps
       model: model choosen
       dataloaders: a dictionary of data loaders one for the training epochs and on efor the validation ones
       criterion  loss function
       optimizer: derivates for backpropagation
       device: GPU or CPU
    """
    # optionally resume from a checkpoint
    if resume:
        if os.path.isfile(resume):
            print("=> loading checkpoint '{}'".format(resume))
            checkpoint = torch.load(resume)
            ## start_epoch = checkpoint['epoch']
            best_acc = checkpoint['best_acc']
            model.load_state_dict(checkpoint['state_dict'])
            optimizer.load_state_dict(checkpoint['optimizer'])
            print("=> loaded checkpoint '{}' (epoch {})"
                  .format(resume, checkpoint['epoch']))
        else:
            print("=> no checkpoint found at '{}'".format(resume))

    since = time.time()

    # decay learning rate when validation loss plateaus
    scheduler = torch.optim.lr_scheduler.ReduceLROnPlateau(optimizer, 'min', patience=10, verbose=True)

    history_acc = []
    history_loss = []
    history_val_acc = []
    history_val_loss = []

    best_model_wts = copy.deepcopy(model.state_dict())
    best_acc = 0.0

    for epoch in range(num_epochs):
        print('Epoch {}/{}'.format(epoch, num_epochs - 1))
        print('-' * 10)

        # Each epoch has a training and validation phase
        for phase in ['train', 'val']:
            if phase == 'train':
                model.train()  # Set model to training mode
            else:
                model.eval()   # Set model to evaluate mode

            running_loss = 0.0
            running_corrects = 0

            # Iterate over data.
            for inputs, labels in dataloaders[phase]:
            ## batch_idx, (inputs, labels) in tqdm.tqdm(enumerate(dataloaders[phase]),total=len(dataloaders[phase]), desc="Epoch: {}".format(epoch)):

                inputs = inputs[:,0:3,:,:]

                inputs = inputs.to(device)
                labels = labels.to(device)

                # zero the parameter gradients
                optimizer.zero_grad()

                # forward
                # track history if only in train
                with torch.set_grad_enabled(phase == 'train'):
                    # Get model outputs and calculate loss
                    # Special case for inception because in training it has an auxiliary output. In train
                    #   mode we calculate the loss by summing the final output and the auxiliary output
                    #   but in testing we only consider the final output.
                    if is_inception and phase == 'train':
                        # From https://discuss.pytorch.org/t/how-to-optimize-inception-model-with-auxiliary-classifiers/7958
                        outputs, aux_outputs = model(inputs)
                        loss1 = criterion(outputs, labels)
                        loss2 = criterion(aux_outputs, labels)
                        loss = loss1 + 0.4*loss2
                    else:
                        outputs = model(inputs)
                        print(outputs.shape)
                        loss = criterion(outputs, labels)

                    _, preds = torch.max(outputs, 1)
                    #outputs = np.array(outputs)
                    #preds = outputs >= 0.5
                    #preds.astype(np.int)
                    # backward + optimize only if in training phase
                    if phase == 'train':
                        loss.backward()
                        optimizer.step()

                # statistics
                running_loss += loss.item() * inputs.size(0)
                running_corrects += torch.sum(preds == labels.data)

            epoch_loss = running_loss / len(dataloaders[phase].dataset)
            epoch_acc = running_corrects.double() / len(dataloaders[phase].dataset)

            print('{} Loss: {:.5f} Acc: {:.4f}'.format(phase, epoch_loss, epoch_acc))

            if phase == 'train':
                history_acc.append(epoch_acc)
                history_loss.append(epoch_loss)
            if phase == 'val':
                history_val_acc.append(epoch_acc)
                history_val_loss.append(epoch_loss)

            # deep copy the model
            if phase == 'val' and epoch_acc > best_acc:
                best_acc = epoch_acc
                save_checkpoint({'epoch': epoch + 1,
                                'state_dict': model.state_dict(),
                                'optimizer' : optimizer.state_dict()},
                                name)
        # epoch end. Does learning rate need decay?
        scheduler.step(epoch_acc)

    ## print()
    ## name_file = 'vgg' + '.txt'
    ## datetime.datetime.now().strftime("%y%m%d%H%M") +
    time_elapsed = time.time() - since
    np.savetxt('./results/' + name + '.txt', (history_acc, history_loss, history_val_acc, history_val_loss), delimiter=',')
    print('Training complete in {:.0f}m {:.0f}s'.format(time_elapsed // 60, time_elapsed % 60))
    print('Best val Acc: {:4f}'.format(best_acc))

    # load best model weights
    ## model.load_state_dict(best_model_wts)
    return model


import tokens, ptypes

class Command(tokens.Token):
    def __init__(self, token, function, argtypes):
        super().__init__(token)
        self.function = function
        self.argtypes = argtypes

    def __str__(self):
        return self._text

    def __repr__(self):
        return "Command({},{},{})".format(self._text,
                                          self.function,
                                          self.argtypes)

class Operator(tokens.Token):
    def __init__(self, token, function, arity, precedence, associativity,
                 default=None, flags=0):
        super().__init__(token)
        self.function = function
        self.arity = arity
        self.precedence = precedence
        self.associativity = associativity
        self.flags = flags
        self.assign = False  # Turns + into +: for instance
        self.map = False  # Turns ! into !* for instance
        self.fold = False  # Turns + into $+ for instance

        # The default argument represents the value when folding an empty list
        # with this operator. For economy of space, defaults are specified as
        # Python objects in the operator table; convert them to Pip types here.
        if type(default) in (ptypes.Nil, ptypes.Scalar, ptypes.List):
            self.default = default
        elif default is None:
            self.default = ptypes.nil
        elif type(default) in (int, float, str):
            self.default = ptypes.Scalar(default)
        elif type(default) is list:
            self.default = ptypes.List(default)
        else:
            print("Unsupported operator default value type:", type(default))
            self.default = ptypes.nil

    def __str__(self):
        return (("$" if self.fold else "")
                + self._text
                + ("*" if self.map else "")
                + (":" if self.assign else ""))

    def __repr__(self):
        return "Operator({},{},{},{},{})".format(str(self),
                                                 self.function,
                                                 self.arity,
                                                 self.precedence,
                                                 self.associativity)

    def copy(self):
        cpy = Operator(self._text,
                       self.function,
                       self.arity,
                       self.precedence,
                       self.associativity,
                       self.default,
                       self.flags)
        cpy.assign = self.assign
        cpy.fold = self.fold
        cpy.map = self.map
        return cpy

cmdTable = [
    ("F", "FOR", ["NAME", "EXPR", "CODE"]),
    ("I", "IF", ["EXPR", "CODE", "ELSE"]),
    ("L", "LOOP", ["EXPR", "CODE"]),
    ("S", "SWAP", ["EXPR", "EXPR"]),
    ("T", "TILL", ["EXPR", "CODE"]),
    ("U", "UNIFY", ["NAMES", "WITH"]),
    ("W", "WHILE", ["EXPR", "CODE"]),
    ]

commands = {cmdSpecs[0]:Command(*cmdSpecs) for cmdSpecs in cmdTable}

# The precedence table contains a list of precedence levels from lowest to
# highest. The first entry in each level is the arity of all operators at that
# level. Each tuple that follows contains the operator symbol, the function,
# and optionally the associativity (default "L"). "C" associativity means
# "Chaining," used for comparison operators (e.g. 1<x<5).

# Flags for operator properties:
VALS = 0x01          # Evaluate all arguments (to lvals or rvals)
RVALS = 0x02         # Evaluate all arguments to rvals
RANGE_EACH = 0x04    # Convert Range arguments to equivalent List of ints
LIST_EACH = 0x08     # Perform operation item by item on Lists
IN_LAMBDA = 0x10     # Can be used to build lambda expressions from _

precedenceTable = [
    [1,
     ("O", "OUTPUT", "L", None, RVALS),
     ("P", "PRINT", "L", None, RVALS),
     ("Y", "YANK", "L", None, RVALS),
     ],
    [2,
     (":", "ASSIGN", "R", None, VALS),
     ],
    [3,
     ("?", "IFTE", "R"),
     ],
    [2,
     ("|", "OR", "L", 0),
     ],
    [2,
     ("&", "AND", "L", 1),
     ],
    [1,
     ("!", "NOT", "L", None, RVALS),
     ],
    [2,
     ("==", "OBJEQUAL", "L", 1, RVALS),  # NB: *not* a chaining operator!
     ],
    [2,
     ("M", "MAP", "R", [], RVALS),
     ("MM", "MAPMAP", "R", [], RVALS),
     ("MJ", "MAPJOIN", "R", "", RVALS),
     ("MS", "MAPSUM", "R", 0, RVALS),
     ("FI", "FILTER", "R", [], RVALS),
     ("SK", "SORTKEYED", "R", [], RVALS),
     ("V", "EVAL", "R", None, RVALS),
     ],
    [1,
     ("V", "EVAL", "L", None, RVALS),
     ],
    [2,
     ("<", "NUMLESS", "C", 1, RVALS),
     (">", "NUMGREATER", "C", 1, RVALS),
     ("=", "NUMEQUAL", "C", 1, RVALS),
     ("<=", "NUMLESSEQ", "C", 1, RVALS),
     (">=", "NUMGREATEREQ", "C", 1, RVALS),
     ("!=", "NUMNOTEQUAL", "C", 1, RVALS),
     ("LT", "STRLESS", "C", 1, RVALS),
     ("GT", "STRGREATER", "C", 1, RVALS),
     ("Q", "STREQUAL", "C", 1, RVALS),
     ("EQ", "STREQUAL", "C", 1, RVALS),  # Synonym for backward compatibility
     ("LE", "STRLESSEQ", "C", 1, RVALS),
     ("GE", "STRGREATEREQ", "C", 1, RVALS),
     ("NE", "STRNOTEQUAL", "C", 1, RVALS),
     ("#=", "LENEQUAL", "C", 1, RVALS),
     ("#<", "LENLESS", "C", 1, RVALS),
     ("#>", "LENGREATER", "C", 1, RVALS),
     ],
     # Note: comparison operators CAN also be used in lambdas, due to the
     # CHAIN pseudo-operator having the IN_LAMBDA flag (see below).
    [2,
     ("N", "IN", "L", None, RVALS | IN_LAMBDA),
     ("IN", "IN", "L", None, RVALS | IN_LAMBDA), # Synonym for backward compat
     ("NI", "NOTIN", "L", None, RVALS | IN_LAMBDA),
     ],
    [1,
     ("RP", "REPR", "L", None, RVALS),
     ("ST", "STR", "L", None, RVALS),
     ("X", "REGEX", "L", None, RVALS | IN_LAMBDA),
     ],
    [2,
     ("CB", "COMBINATIONS", "L", [], RVALS | IN_LAMBDA),
     ],
    [1,
     ("MX", "MAX", "L", None, RVALS | IN_LAMBDA),
     ("MN", "MIN", "L", None, RVALS | IN_LAMBDA),
     ("RC", "RANDCHOICE", "L", None, RVALS | IN_LAMBDA),
     ("SN", "SORTNUM", "L", None, RVALS | IN_LAMBDA),
     ("SS", "SORTSTRING", "L", None, RVALS | IN_LAMBDA),
     ("UQ", "UNIQUE", "L", None, RVALS | IN_LAMBDA),
     ("EN", "ENUMERATE", "L", None, RVALS | IN_LAMBDA),
     ("PM", "PERMUTATIONS", "L", None, RVALS | IN_LAMBDA),
     ],
    [2,
     ("AE", "APPENDELEM", "L", [], RVALS | IN_LAMBDA),
     ("AL", "APPENDLIST", "L", [], RVALS | IN_LAMBDA),
     ("PE", "PREPENDELEM", "L", [], RVALS | IN_LAMBDA),
     ("PU", "PUSH", "L", [], VALS | IN_LAMBDA),
     ("PB", "PUSHBACK", "L", [], VALS | IN_LAMBDA),
     ],
    [1,
     ("PO", "POP", "L", None, VALS | IN_LAMBDA),
     ("DQ", "DEQUEUE", "L", None, VALS | IN_LAMBDA),
     ],
    [2,
     ("^", "SPLIT", "L", [], RVALS | IN_LAMBDA | LIST_EACH | RANGE_EACH),
     ("^@", "SPLITAT", "L", [], RVALS | IN_LAMBDA),
     ("@?", "FIND", "L", None, RVALS | IN_LAMBDA),
     ("@*", "FINDALL", "L", [], RVALS | IN_LAMBDA),
     ("<>", "GROUP", "L", [], RVALS | IN_LAMBDA),
     ("J", "JOIN", "L", "", RVALS | IN_LAMBDA),
     ("RL", "REPEATLIST", "L", [], RVALS | IN_LAMBDA),
     ("Z", "ZIP", "L", [], RVALS | IN_LAMBDA),
     ("ZD", "ZIPDEFAULT", "L", [], RVALS | IN_LAMBDA),
     ("CP", "CARTESIANPRODUCT", "L", [], RVALS | IN_LAMBDA),
     ("CG", "COORDINATEGRID", "L", None, RVALS | IN_LAMBDA),
     ("ZG", "ZEROGRID", "L", None, RVALS | IN_LAMBDA),
     ],
    [1,
     ("^", "SPLIT", "L", None, RVALS | IN_LAMBDA | LIST_EACH | RANGE_EACH),
     ("J", "JOIN", "L", None, RVALS | IN_LAMBDA),
     ("RV", "REVERSE", "L", None, RVALS | IN_LAMBDA),
     ("Z", "ZIP", "L", None, RVALS | IN_LAMBDA),
     ("ZD", "ZIPDEFAULT", "L", None, RVALS | IN_LAMBDA),
     ("CP", "CARTESIANPRODUCT", "L", None, RVALS | IN_LAMBDA),
     ("CG", "COORDINATEGRID", "L", None, RVALS | IN_LAMBDA),
     ("ZG", "ZEROGRID", "L", None, RVALS | IN_LAMBDA),
     ],
    [3,
     ("R", "REPLACE", "L", None, RVALS),
     ],
    [2,
     (".", "CAT", "L", "", RVALS | IN_LAMBDA | RANGE_EACH | LIST_EACH),
     ],
    [2,
     ("RM", "REMOVE", "L", "", RVALS | IN_LAMBDA),
     ],
    [2,
     ("X", "STRMUL", "L", "", RVALS | IN_LAMBDA | RANGE_EACH | LIST_EACH),
     ],
    [2,
     ("||", "STRIP", "L", "", RVALS | IN_LAMBDA | LIST_EACH),
     ("|>", "LSTRIP", "L", "", RVALS | IN_LAMBDA | LIST_EACH),
     ("<|", "RSTRIP", "L", "", RVALS | IN_LAMBDA | LIST_EACH),
     ("TM", "TRIM", "L", "", RVALS | IN_LAMBDA | LIST_EACH),
     ],
    [1,
     ("||", "STRIP", "L", "", RVALS | IN_LAMBDA | LIST_EACH),
     ("|>", "LSTRIP", "L", "", RVALS | IN_LAMBDA | LIST_EACH),
     ("<|", "RSTRIP", "L", "", RVALS | IN_LAMBDA | LIST_EACH),
     ("TM", "TRIM", "L", "", RVALS | IN_LAMBDA | LIST_EACH),
     ("LC", "LOWERCASE", "L", None, RVALS | IN_LAMBDA | LIST_EACH),
     ("UC", "UPPERCASE", "L", None, RVALS | IN_LAMBDA | LIST_EACH),
     ],
    [2,
     (",", "RANGE", "L", None, RVALS | IN_LAMBDA | RANGE_EACH | LIST_EACH),
     ("RR", "RANDRANGE", "L", 0, RVALS | IN_LAMBDA | RANGE_EACH | LIST_EACH),
     ("TB", "TOBASE", "L", 0, RVALS | IN_LAMBDA | RANGE_EACH | LIST_EACH),
     ],
    [1,
     (",", "RANGETO", "L", None, RVALS | IN_LAMBDA | RANGE_EACH | LIST_EACH),
     ("RR", "RANDRANGETO", "L", None, RVALS | IN_LAMBDA | RANGE_EACH | LIST_EACH),
     ("TB", "TOBASE", "L", None, RVALS | IN_LAMBDA | RANGE_EACH | LIST_EACH),
     # Unary mnemonic: ToBinary
     ],
    [2,
     ("BA", "BITWISEAND", "L", -1, RVALS | IN_LAMBDA | RANGE_EACH | LIST_EACH),
     ("BO", "BITWISEOR", "L", 0, RVALS | IN_LAMBDA | RANGE_EACH | LIST_EACH),
     ("BX", "BITWISEXOR", "L", 0, RVALS | IN_LAMBDA | RANGE_EACH | LIST_EACH),
     ],
    [1,
     ("BN", "BITWISENOT", "L",None, RVALS | IN_LAMBDA | RANGE_EACH | LIST_EACH),
     ],
    [2,
     ("<=>", "NUMCMP", "L", 0, RVALS | IN_LAMBDA),
     ],
    [2,
     ("+", "ADD", "L", 0, RVALS | IN_LAMBDA | LIST_EACH),
     ("-", "SUB", "L", 0, RVALS | IN_LAMBDA | LIST_EACH),
     ],
    [2,
     ("*", "MUL", "L", 1, RVALS | IN_LAMBDA | RANGE_EACH | LIST_EACH),
     ("/", "DIV", "L", 1, RVALS | IN_LAMBDA | RANGE_EACH | LIST_EACH),
     ("%", "MOD", "L", 0, RVALS | IN_LAMBDA | RANGE_EACH | LIST_EACH),
     ("//", "INTDIV", "L", 1, RVALS | IN_LAMBDA | RANGE_EACH | LIST_EACH),
     ],
    [1,
     ("+", "POS", "L", None, RVALS | IN_LAMBDA | LIST_EACH),
     ("-", "NEG", "L", None, RVALS | IN_LAMBDA | RANGE_EACH | LIST_EACH),
     ("/", "INVERT", "L", None, RVALS | IN_LAMBDA | RANGE_EACH | LIST_EACH),
     ],
    [2,
     ("**", "POW", "R", 1, RVALS | IN_LAMBDA | RANGE_EACH | LIST_EACH),
     ("RT", "ROOT", "R", 1, RVALS | IN_LAMBDA | RANGE_EACH | LIST_EACH),
     ("AT", "ARCTAN", "L", None, RVALS | IN_LAMBDA | RANGE_EACH | LIST_EACH),
     ],
    [1,
     ("RT", "SQRT", "L", None, RVALS | IN_LAMBDA | RANGE_EACH | LIST_EACH),
     ("SI", "SINE", "L", None, RVALS | IN_LAMBDA | RANGE_EACH | LIST_EACH),
     ("CO", "COSINE", "L", None, RVALS | IN_LAMBDA | RANGE_EACH | LIST_EACH),
     ("TA", "TANGENT", "L", None, RVALS | IN_LAMBDA | RANGE_EACH | LIST_EACH),
     ("SE", "SECANT", "L", None, RVALS | IN_LAMBDA | RANGE_EACH | LIST_EACH),
     ("CS", "COSEC", "L", None, RVALS | IN_LAMBDA | RANGE_EACH | LIST_EACH),
     ("CT", "COTAN", "L", None, RVALS | IN_LAMBDA | RANGE_EACH | LIST_EACH),
     ("AT", "ARCTAN", "L", None, RVALS | IN_LAMBDA | RANGE_EACH | LIST_EACH),
     ("RD", "RADIANS", "L", None, RVALS | IN_LAMBDA | RANGE_EACH | LIST_EACH),
     ("DG", "DEGREES", "L", None, RVALS | IN_LAMBDA | RANGE_EACH | LIST_EACH),
     ],
    [2,
     ("FB", "FROMBASE", "L", 0, RVALS | IN_LAMBDA | RANGE_EACH | LIST_EACH),
     ],
    [2,
     ("@", "AT", "L", None, VALS | IN_LAMBDA),
     ("@<", "LEFTOF", "L", None, VALS | IN_LAMBDA),
     ("@>", "RIGHTOF", "L", None, VALS | IN_LAMBDA),
     ],
    [1,
     ("@<", "LEFTOF", "L", None, VALS | IN_LAMBDA),
     ("@>", "RIGHTOF", "L", None, VALS | IN_LAMBDA),
     ("++", "INC", "L", None, VALS | IN_LAMBDA),
     ("--", "DEC", "L", None, VALS | IN_LAMBDA),
     ("#", "LEN", "L", None, RVALS | IN_LAMBDA),
     ("A", "ASC", "L", None, RVALS | IN_LAMBDA | RANGE_EACH | LIST_EACH),
     ("C", "CHR", "L", None, RVALS | IN_LAMBDA | RANGE_EACH | LIST_EACH),
     ("AB", "ABS", "L", None, RVALS | IN_LAMBDA | RANGE_EACH | LIST_EACH),
     ("SG", "SIGN", "L", None, RVALS | IN_LAMBDA | RANGE_EACH | LIST_EACH),
     ("FB", "FROMBASE", "L", None, RVALS | IN_LAMBDA | RANGE_EACH | LIST_EACH),
     # Unary mnemonic: FromBinary
     ],
    ]

# Now take the information from the precedence table and rearrange it into
# structures more useful to the scanner and parser:

opsByArity = {1:{}, 2:{}, 3:{}}
operators = set()

# Some convenience operators for parsing that don't have straightforward
# equivalents in the syntax:
highestPrecedence = len(precedenceTable)
paren = Operator("PAREN", "PARENTHESIZE", 1, highestPrecedence, "L", None, VALS)
enlist = Operator("LIST", "LIST", 1, highestPrecedence, "L", None, RVALS)
block = Operator("BLOCK", "BLOCK", 1, highestPrecedence, "L", None)
send = Operator("SEND", "SEND", 1, highestPrecedence, "L", None, VALS)
chain = None  # We define this "operator" later to make sure it gets the same
              # precedence as the comparison operators it comprises

for precedence, (arity, *entries) in enumerate(precedenceTable):
    for entry in entries:
        text, function, associativity = entry[:3]
        if len(entry) > 3:
            default = entry[3]
        else:
            default = None
        if len(entry) > 4:
            flags = entry[4]
        else:
            flags = 0
        if arity > 2 and flags & (LIST_EACH | RANGE_EACH):
            # TODO: proper implementation error message
            msg = "Current implementation cannot handle LIST_EACH or RANGE_EACH"
            msg += "\nfor operators of arity greater than 2 (like %s)" % text
            print(msg)
            flags = flags & ~LIST_EACH & ~RANGE_EACH
        if flags & IN_LAMBDA and not flags & (VALS | RVALS):
            # TODO: proper implementation error message
            msg = "IN_LAMBDA may not be set for operators that do not also"
            msg += "\nset VALS or RVALS (%s)" % test
            print(msg)
            flags = flags & ~IN_LAMBDA
        op = Operator(text,
                      function,
                      arity,
                      precedence,
                      associativity,
                      default,
                      flags)
        if associativity == "C" and not chain:
            # Define the CHAIN pseudo-operator, which allows parsing of
            # chained comparisons like 1<x<5
            chain = Operator("CHAIN",
                             "CHAIN",
                             2,
                             precedence,
                             "C",
                             None,
                             RVALS | IN_LAMBDA)
        opsByArity[arity][text] = op
        operators.add(text)


from robopy import configs
from flask import Flask, json, request, Response

app = Flask(__name__)

@app.route('/init', methods=['POST'])
def entitiesNames():
    """
    Receives request 'POST/init',
    """

    resp_dict = dict()
    resp_dict['init'] = 'ok'

    resp = Response(json.dumps(resp_dict), status=200)

    return resp

@app.route('/kick', methods=['POST'])
def sentiment():
    """
    Receives request 'POST/sentiment', executes NLP pipeline with mode 'sa',
    responses back with result
    """

    request_json = request.json
    power = request_json['power']
    angle = request_json['angle']

    print(power, angle)

    resp_dict = dict()
    resp_dict['kick'] = 'ok'

    resp = Response(json.dumps(resp_dict), status=200)

    return resp

if __name__ == '__main__':
    app.run(host=configs.DEFAULT_HOST,
            port=configs.DEFAULT_PORT,
            debug=True,
            threaded=True)


from django.shortcuts import render
from django.http import HttpResponseRedirect
from django.urls import reverse_lazy, reverse

from django.views.generic import (
    CreateView,
    UpdateView,
    TemplateView,
    DetailView,
    FormView,
)

#app torneo
from applications.torneo.models import Tournament

#.
from .models import Team

class TeamAddTemplateView(TemplateView):
    """vista que agregar nuevo equipo"""
    template_name = 'team/add.html'

    def get_context_data(self, **kwargs):
        context = super(TeamAddTemplateView, self).get_context_data(**kwargs)
        context['flat_menu'] = 'E'
        return context

class TeamByDelegateView(DetailView):
    """vista que lista equipos por delegado y torneo"""
    model = Tournament
    template_name = 'team/team_for_inscribed.html'

    def get_context_data(self, **kwargs):
        context = super(TeamByDelegateView, self).get_context_data(**kwargs)
        context['flat_menu'] = 'PR'
        return context

class TeamByUserView(TemplateView):
    """vista que lista equipos de un usuario"""

    def get_template_names(self):
        #determinamos la ruta del template
        if self.template_name is None:
            filtro = self.kwargs.get('filtro', 0)
            if filtro == 'mios':
                self.template_name = 'team/by_usuario/mis_equipos.html'
            elif filtro == 'deportista':
                self.template_name = 'team/by_usuario/donde_juego.html'
            elif filtro == 'siguiendo':
                self.template_name = 'team/by_usuario/siguiendo.html'
            else:
                self.template_name = 'home/404.html'
            return self.template_name
        else:
            return [self.template_name]

    def get_context_data(self, **kwargs):
        context = super(TeamByUserView, self).get_context_data(**kwargs)
        context['flat'] = 'PR'
        context['flat_menu'] = 'E'
        return context

class TeamByTournamentView(DetailView):
    """vista que lista equipos de un torneo"""
    model = Tournament

    def get_template_names(self):
        #determinamos la ruta del template
        if self.template_name is None:
            filtro = self.kwargs.get('filtro', 0)
            if filtro == 'inscritos':
                self.template_name = 'team/by_torneo/by_torneo_inscrito.html'
            elif filtro == 'solicitud':
                self.template_name = 'team/by_torneo/by_torneo_solicitud.html'
            elif filtro == 'rechazados':
                self.template_name = 'team/by_torneo/by_torneo_rechazado.html'
            else:
                self.template_name = 'home/404.html'
            return self.template_name
        else:
            return [self.template_name]

    def get_context_data(self, **kwargs):
        context = super(TeamByTournamentView, self).get_context_data(**kwargs)
        context['flat'] = 'T'
        context['flat_menu'] = 'TE'
        return context


# PoseNet + face_recognition

# main.py
import cv2
import os
import math
import face_recognition
from PIL import Image, ImageDraw

BODY_PARTS_BODY_25 = {0: "Nose", 1: "Neck", 2: "RShoulder", 3: "RElbow", 4: "RWrist",
                      5: "LShoulder", 6: "LElbow", 7: "LWrist", 8: "MidHip", 9: "RHip",
                      10: "RKnee", 11: "RAnkle", 12: "LHip", 13: "LKnee", 14: "LAnkle",
                      15: "REye", 16: "LEye", 17: "REar", 18: "LEar", 19: "LBigToe",
                      20: "LSmallToe", 21: "LHeel", 22: "RBigToe", 23: "RSmallToe", 24: "RHeel", 25: "Background"}

POSE_PAIRS_BODY_25 = [[0, 1], [0, 15], [0, 16], [1, 2], [1, 5], [1, 8], [8, 9], [8, 12], [9, 10], [12, 13], [2, 3],
                      [3, 4], [5, 6], [6, 7], [10, 11], [13, 14], [15, 17], [16, 18], [14, 21], [19, 21], [20, 21],
                      [11, 24], [22, 24], [23, 24]]

# 신경 네트워크의 구조를 지정하는 prototxt 파일 (다양한 계층이 배열되는 방법 등)
protoFile_body_25 = "C:/Users/user/Documents/GitHub/blackbox/models/pose_deploy.prototxt"

# 훈련된 모델의 weight 를 저장하는 caffemodel 파일
weightsFile_body_25 = "C:/Users/user/Documents/GitHub/blackbox/models/pose_iter_584000.caffemodel"

# 네트워크 불러오기
net = cv2.dnn.readNetFromCaffe(protoFile_body_25, weightsFile_body_25)

def output_keypoints(image, threshold, BODY_PARTS):
    global points

    # 입력 이미지의 사이즈 정의
    image_height = 368
    image_width = 368

    # 네트워크에 넣기 위한 전처리
    input_blob = cv2.dnn.blobFromImage(image, 1.0 / 255, (image_width, image_height), (0, 0, 0), swapRB=False,
                                       crop=False)

    # 전처리된 blob 네트워크에 입력
    net.setInput(input_blob)

    # 결과 받아오기
    out = net.forward()
    # The output is a 4D matrix :
    # The first dimension being the image ID ( in case you pass more than one image to the network ).
    # The second dimension indicates the index of a keypoint.
    # The model produces Confidence Maps and Part Affinity maps which are all concatenated.
    # For COCO model it consists of 57 parts – 18 keypoint confidence Maps + 1 background + 19*2 Part Affinity Maps. Similarly, for MPI, it produces 44 points.
    # We will be using only the first few points which correspond to Keypoints.
    # The third dimension is the height of the output map.
    out_height = out.shape[2]
    # The fourth dimension is the width of the output map.
    out_width = out.shape[3]

    # 원본 이미지의 높이, 너비를 받아오기
    frame_height, frame_width = image.shape[:2]

    # 포인트 리스트 초기화
    points = []

    for i in range(len(BODY_PARTS)):

        # 신체 부위의 confidence map
        prob_map = out[0, i, :, :]

        # 최소값, 최대값, 최소값 위치, 최대값 위치
        min_val, prob, min_loc, point = cv2.minMaxLoc(prob_map)

        # 원본 이미지에 맞게 포인트 위치 조정
        x = (frame_width * point[0]) / out_width
        x = int(x)
        y = (frame_height * point[1]) / out_height
        y = int(y)

        if prob > threshold:  # [pointed]
            cv2.circle(image, (x, y), 5, (0, 255, 255), thickness=-1, lineType=cv2.FILLED)
            points.append((x, y))

        else:  # [not pointed]
            cv2.circle(image, (x, y), 5, (0, 255, 255), thickness=-1, lineType=cv2.FILLED)
            points.append(None)

    return image

def output_keypoints_with_lines(POSE_PAIRS, frame):
    # 프레임 복사
    frame_line = frame.copy()

    # Neck 과 MidHeap 의 좌표값이 존재한다면
    # if (points[1] is not None) and (points[8] is not None):
    #     calculate_degree(point_1=points[1], point_2=points[8], frame=frame_line)

    for pair in POSE_PAIRS:
        part_a = pair[0]  # 0 (Head)
        part_b = pair[1]  # 1 (Neck)
        if points[part_a] and points[part_b]:
            print(f"[linked] {part_a} {points[part_a]} <=> {part_b} {points[part_b]}")
            # Neck 과 MidHip 이라면 분홍색 선
            if part_a == 1 and part_b == 8:
                cv2.line(frame, points[part_a], points[part_b], (255, 0, 255), 3)
            else:  # 노란색 선
                cv2.line(frame, points[part_a], points[part_b], (0, 255, 0), 3)
        else:
            print(f"[not linked] {part_a} {points[part_a]} <=> {part_b} {points[part_b]}")

    return frame_line

# def calculate_degree(point_1, point_2, frame):
#     # 역탄젠트 구하기
#     dx = point_2[0] - point_1[0]
#     dy = point_2[1] - point_1[1]
#     rad = math.atan2(abs(dy), abs(dx))
#
#     # radian 을 degree 로 변환
#     deg = rad * 180 / math.pi
#
#     # degree 가 45'보다 작으면 허리가 숙여졌다고 판단
#     if deg < 45:
#         string = "Bend Down"
#         cv2.putText(frame, string, (0, 25), cv2.FONT_HERSHEY_DUPLEX, 1, (255, 0, 255))
#         print(f"[degree] {deg} ({string})")
#     else:
#         string = "Stand"
#         cv2.putText(frame, string, (0, 25), cv2.FONT_HERSHEY_DUPLEX, 1, (255, 0, 255))
#         print(f"[degree] {deg} ({string})")

# 영상 처리
def video_processing(video_path, background):

    cap = cv2.VideoCapture(video_path)
    fourcc = cv2.VideoWriter_fourcc('m', 'p', '4', 'v')

    frame_count = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))
    now_frame = 1

    if not os.path.exists('outputs'):
        os.mkdir('outputs')

    out = None

    colors = [(0, 255, 0), (0, 0, 255)] # 초록, 빨강
    #labels = ['with_mask', 'without_mask']

    while cap.isOpened():
        ret, image = cap.read()
        if not ret:
            break

        face_locations = face_recognition.face_locations(image)

        face_image = Image.fromarray(image)
        draw = ImageDraw.Draw(face_image)

        result_image = image.copy()

        for face_location in face_locations:
            top = face_location[0]
            right = face_location[1]
            bottom = face_location[2]
            left = face_location[3]

            cv2.rectangle(
                result_image,
                pt1=(left, top),
                pt2=(right, bottom),
                thickness=2,
                color=colors[0],
                lineType=cv2.LINE_AA
            )

        points = []

        frame_man = output_keypoints(image=result_image, threshold=0.1, BODY_PARTS=BODY_PARTS_BODY_25)
        image_frame = output_keypoints_with_lines(POSE_PAIRS=POSE_PAIRS_BODY_25, frame=frame_man)

        #print(type(result_image))

        if out is None:
            out = cv2.VideoWriter(
                'C:/Users/user/Documents/GitHub/blackbox/outputs/output_test.wmv',
                fourcc,
                cap.get(cv2.CAP_PROP_FPS),
                # (image_frame.shape[1], image_frame.shape[0])
                (image.shape[1], image.shape[0])
            )
        else:
            out.write(result_image)

        # (10/400): 11%
        print('(' + str(now_frame) + '/' + str(frame_count) + '): ' + str(now_frame * 100 // frame_count) + '%')
        now_frame += 1

        if not background:
            cv2.imshow('result', result_image)
            if cv2.waitKey(1) == ord('q'):
                break

    out.release()
    cap.release()

if __name__ == '__main__':
    #video_processing('C:/Users/user/PycharmProjects/OpenCV/doc/watershed_TestVideo/success_1.mp4', False)
    #video_processing('C:/Users/user/Documents/GitHub/blackbox/data/05.mp4', False)
    video_processing('C:/Users/user/Documents/GitHub/blackbox/data/baby1.mp4', False)


import pommerman
from pommerman import agents
from enum import Enum
import numpy as np
import tensorflow as tf

class Item(Enum):
	"""The Items in the game.

	When picked up:
	  - ExtraBomb increments the agent's ammo by 1.
	  - IncrRange increments the agent's blast strength by 1.
	  - Kick grants the agent the ability to kick items.

	AgentDummy is used by team games to denote the third enemy and by ffa to
	denote the teammate.
	"""
	Passage = 0
	Rigid = 1
	Wood = 2
	Bomb = 3
	Flames = 4
	Fog = 5
	ExtraBomb = 6
	IncrRange = 7
	Kick = 8
	AgentDummy = 9
	Agent0 = 10
	Agent1 = 11
	Agent2 = 12
	Agent3 = 13

def state_to_matrix(obs):
    #In this implementation I just concatenate everything in one big matrix

    #for e in obs['enemies']:
    #    print(e)
    #    print(Item(e))
    #TODO enemies
    my_position = np.asmatrix(obs['position'])
    bomb_life = np.asmatrix(obs['bomb_life'])
    board = np.asmatrix(obs['board'])
    bombs = np.asmatrix(obs['bomb_blast_strength'])
    #enemies = np.asmatrix([Item_en(e) for e in obs['enemies']])
    can_kick = np.asmatrix(int(1 if obs['can_kick'] else 0))
    ammo = np.asmatrix(int(obs['ammo']))
    blast_strength = np.asmatrix(int(obs['blast_strength']))

    m = np.max([my_position.shape[1], bomb_life.shape[1], board.shape[1], bombs.shape[1],  can_kick.shape[1], ammo.shape[1], blast_strength.shape[1]])

    my_position1 = np.concatenate((my_position, np.zeros(( my_position.shape[0], m - my_position.shape[1]))), axis=1)
    bomb_life1 = np.concatenate((bomb_life, np.zeros((bomb_life.shape[0], m - bomb_life.shape[1]))), axis=1)
    board1 = np.concatenate((board, np.zeros((board.shape[0], m - board.shape[1]))), axis=1)
    bombs1 = np.concatenate((bombs, np.zeros((bombs.shape[0], m - bombs.shape[1]))), axis=1)
    #enemies1 = np.concatenate((enemies, np.zeros((enemies.shape[0], m - enemies.shape[1]))), axis=1)
    can_kick1 = np.concatenate((can_kick, np.zeros((can_kick.shape[0], m - can_kick.shape[1]))), axis=1)
    ammo1 = np.concatenate((ammo, np.zeros((ammo.shape[0], m - ammo.shape[1]))), axis=1)
    blast_strength1 = np.concatenate((blast_strength, np.zeros((blast_strength.shape[0], m - blast_strength.shape[1]))), axis=1)

    result = np.concatenate((my_position1, bomb_life1, board1, board1, bombs1, can_kick1, ammo1, blast_strength1), axis=0)
    return np.asmatrix(result)

# def state_to_matrixes(obs):
#     #In this implementation I put all vector in matrix and concatenate this matrixes
#     def get_map(board, item):
#         map = np.zeros(shape)
#         map[board == item] = 1
#         return map

#     def convert_bombs(bomb_map):
#         ret = []
#         locations = np.where(bomb_map > 0)
#         for r, c in zip(locations[0], locations[1]):
#             ret.append({
#                 'position': (r, c),
#                 'blast_strength': int(bomb_map[(r, c)])
#             })
#         return ret

#     board = np.array(obs['board'])

#     rigid_map = get_map(board, 1)  # Rigid = 1
#     wood_map = get_map(board, 2)  # Wood = 2
#     bomb_map = get_map(board, 3)  # Bomb = 3
#     flames_map = get_map(board, 4)  # Flames = 4
#     fog_map = get_map(board, 5)  # Fog = 5
#     extra_bomb_map = get_map(board, 6)  # ExtraBomb = 6
#     incr_range_map = get_map(board, 7)  # IncrRange = 7
#     kick_map = get_map(board, 8)  # Kick = 8
#     skull_map = get_map(board, 9)  # Skull = 9

#     my_position = np.asmatrix(tuple(obs['position']))
#     bomb_life = np.array(obs['bomb_life'])
#     bombs = convert_bombs(np.array(obs['bomb_blast_strength']))
#     enemies = np.asmatrix([constants.Item(e) for e in obs['enemies']])
#     can_kick = np.asmatrix(int(obs['Can Kick']))
#     ammo = np.asmatrix(int(obs['ammo']))
#     blast_strength = np.asmatrix(int(obs['blast_strength']))

#     m = np.max([my_position.shape[1], bomb_life.shape[1], board.shape[1], bombs.shape[1], enemies.shape[1], can_kick.shape[1], ammo.shape[1], blast_strength.shape[1]])

#     my_position1 = np.concatenate((my_position, np.zeros(( my_position.shape[0], m - my_position.shape[1]))), axis=1)
#     bomb_life1 = np.concatenate((bomb_life, np.zeros((bomb_life.shape[0], m - bomb_life.shape[1]))), axis=1)
#     board1 = np.concatenate((board, np.zeros((board.shape[0], m - board.shape[1]))), axis=1)
#     bombs1 = np.concatenate((bombs, np.zeros((bombs.shape[0], m - bombs.shape[1]))), axis=1)
#     enemies1 = np.concatenate((enemies, np.zeros((enemies.shape[0], m - enemies.shape[1]))), axis=1)
#     can_kick1 = np.concatenate((can_kick, np.zeros((can_kick.shape[0], m - can_kick.shape[1]))), axis=1)
#     ammo1 = np.concatenate((ammo, np.zeros((ammo.shape[0], m - ammo.shape[1]))), axis=1)
#     blast_strength1 = np.concatenate((blast_strength, np.zeros((blast_strength.shape[0], m - blast_strength.shape[1]))), axis=1)

#     result = np.concatenate((my_position1, bomb_life1, board1, board1, bombs1, enemies1, can_kick1, ammo1, blast_strength1, rigid_map, wood_map,\
#                              bomb_map, flames_map, fog_map, fog_map, extra_bomb_map, incr_range_map, kick_map, skull_map), axis=0)
#     return result

def state_to_matrix_with_action(obs, action):
    # In this implementation I just concatenate everything in one big matrix

    # for e in obs['enemies']:
    #    print(e)
    #    print(Item(e))
    # TODO enemies
    my_position = np.asmatrix(obs['position'])
    bomb_life = np.asmatrix(obs['bomb_life'])
    board = np.asmatrix(obs['board'])
    bombs = np.asmatrix(obs['bomb_blast_strength'])
    # enemies = np.asmatrix([Item_en(e) for e in obs['enemies']])
    can_kick = np.asmatrix(int(1 if obs['can_kick'] else 0))
    ammo = np.asmatrix(int(obs['ammo']))
    blast_strength = np.asmatrix(int(obs['blast_strength']))

    m = np.max([my_position.shape[1], bomb_life.shape[1], board.shape[1], bombs.shape[1], can_kick.shape[1], ammo.shape[1],
         blast_strength.shape[1]])

    my_position1 = np.concatenate((my_position, np.zeros((my_position.shape[0], m - my_position.shape[1]))), axis=1)
    bomb_life1 = np.concatenate((bomb_life, np.zeros((bomb_life.shape[0], m - bomb_life.shape[1]))), axis=1)
    board1 = np.concatenate((board, np.zeros((board.shape[0], m - board.shape[1]))), axis=1)
    bombs1 = np.concatenate((bombs, np.zeros((bombs.shape[0], m - bombs.shape[1]))), axis=1)
    # enemies1 = np.concatenate((enemies, np.zeros((enemies.shape[0], m - enemies.shape[1]))), axis=1)
    can_kick1 = np.concatenate((can_kick, np.zeros((can_kick.shape[0], m - can_kick.shape[1]))), axis=1)
    ammo1 = np.concatenate((ammo, np.zeros((ammo.shape[0], m - ammo.shape[1]))), axis=1)
    blast_strength1 = np.concatenate((blast_strength, np.zeros((blast_strength.shape[0], m - blast_strength.shape[1]))), axis=1)
    action = np.asmatrix(action)
    action1 = np.concatenate((action, np.zeros((action.shape[0], m - action.shape[1]))), axis=1)
#     print(my_position1.shape, bomb_life1.shape, board1.shape, bombs1.shape, can_kick1.shape, ammo1.shape,
#           blast_strength1.shape, action.shape, action1.shape)
    result = np.concatenate((my_position1, bomb_life1, board1, bombs1, can_kick1, ammo1, blast_strength1, action1),
                            axis=0)
    return np.asmatrix(result)

def normalize(inputs,
              epsilon = 1e-8,
              scope="ln",
              reuse=None):
    '''Applies layer normalization.

    Args:
      inputs: A tensor with 2 or more dimensions, where the first dimension has
        `batch_size`.
      epsilon: A floating number. A very small number for preventing ZeroDivision Error.
      scope: Optional scope for `variable_scope`.
      reuse: Boolean, whether to reuse the weights of a previous layer
        by the same name.

    Returns:
      A tensor with the same shape and data dtype as `inputs`.
    '''
    with tf.variable_scope(scope, reuse=reuse):
        inputs_shape = inputs.get_shape()
        params_shape = inputs_shape[-1:]

        mean, variance = tf.nn.moments(inputs, [-1], keep_dims=True)
#         print('mean', mean, 'variance', variance)
        beta= tf.Variable(tf.zeros(params_shape))
#         print('beta', beta)
        gamma = tf.Variable(tf.ones(params_shape))
#         print('gamma', gamma)
        normalized = (inputs - mean) / ( (variance + epsilon) ** (.5) )
#         print('normalized before', normalized)
        ##Because normalized has float64 type. Need to change to float32 to operate.
        normalized = tf.cast(normalized, tf.float32)
#         print('normalized after', normalized)
        outputs = gamma * normalized + beta
#         print('outputs', outputs)

    return outputs

def embedding(inputs,
              vocab_size,
              num_units,
              zero_pad=True,
              scale=True,
              scope="embedding",
              reuse=None):
    '''Embeds a given tensor.

    Args:
      inputs: A `Tensor` with type `int32` or `int64` containing the ids
         to be looked up in `lookup table`.
      vocab_size: An int. Vocabulary size.
      num_units: An int. Number of embedding hidden units.
      zero_pad: A boolean. If True, all the values of the fist row (id 0)
        should be constant zeros.
      scale: A boolean. If True. the outputs is multiplied by sqrt num_units.
      scope: Optional scope for `variable_scope`.
      reuse: Boolean, whether to reuse the weights of a previous layer
        by the same name.

    Returns:
      A `Tensor` with one more rank than inputs's. The last dimensionality
        should be `num_units`.

    For example,

    ```
    import tensorflow as tf

    inputs = tf.to_int32(tf.reshape(tf.range(2*3), (2, 3)))
    outputs = embedding(inputs, 6, 2, zero_pad=True)
    with tf.Session() as sess:
        sess.run(tf.global_variables_initializer())
        print sess.run(outputs)
    >>
    [[[ 0.          0.        ]
      [ 0.09754146  0.67385566]
      [ 0.37864095 -0.35689294]]

     [[-1.01329422 -1.09939694]
      [ 0.7521342   0.38203377]
      [-0.04973143 -0.06210355]]]
    ```

    ```
    import tensorflow as tf

    inputs = tf.to_int32(tf.reshape(tf.range(2*3), (2, 3)))
    outputs = embedding(inputs, 6, 2, zero_pad=False)
    with tf.Session() as sess:
        sess.run(tf.global_variables_initializer())
        print sess.run(outputs)
    >>
    [[[-0.19172323 -0.39159766]
      [-0.43212751 -0.66207761]
      [ 1.03452027 -0.26704335]]

     [[-0.11634696 -0.35983452]
      [ 0.50208133  0.53509563]
      [ 1.22204471 -0.96587461]]]
    ```
    '''
    with tf.variable_scope(scope, reuse=reuse):
        lookup_table = tf.get_variable('lookup_table',
                                       dtype=tf.float32,
                                       shape=[vocab_size, num_units],
                                       initializer=tf.contrib.layers.xavier_initializer())
        if zero_pad:
            lookup_table = tf.concat((tf.zeros(shape=[1, num_units]),
                                      lookup_table[1:, :]), 0)
            print('lookup_table', lookup_table)
            print('inputs', inputs)
        outputs = tf.nn.embedding_lookup(lookup_table, inputs)
        print('outputs', outputs)

        if scale:
            outputs = outputs * (num_units ** 0.5)

    return outputs

def positional_encoding(inputs,
                        num_units,
                        zero_pad=True,
                        scale=True,
                        scope="positional_encoding",
                        reuse=None):
    '''Sinusoidal Positional_Encoding.

    Args:
      inputs: A 2d Tensor with shape of (N, T).
      num_units: Output dimensionality
      zero_pad: Boolean. If True, all the values of the first row (id = 0) should be constant zero
      scale: Boolean. If True, the output will be multiplied by sqrt num_units(check details from paper)
      scope: Optional scope for `variable_scope`.
      reuse: Boolean, whether to reuse the weights of a previous layer
        by the same name.

    Returns:
        A 'Tensor' with one more rank than inputs's, with the dimensionality should be 'num_units'
    '''

    N, T = inputs.get_shape().as_list()
    with tf.variable_scope(scope, reuse=reuse):
        position_ind = tf.tile(tf.expand_dims(tf.range(T), 0), [N, 1])

        # First part of the PE function: sin and cos argument
        position_enc = np.array([
            [pos / np.power(10000, 2.*i/num_units) for i in range(num_units)]
            for pos in range(T)])

        # Second part, apply the cosine to even columns and sin to odds.
        position_enc[:, 0::2] = np.sin(position_enc[:, 0::2])  # dim 2i
        position_enc[:, 1::2] = np.cos(position_enc[:, 1::2])  # dim 2i+1

        # Convert to a tensor
        lookup_table = tf.convert_to_tensor(position_enc)

        if zero_pad:
            lookup_table = tf.concat((tf.zeros(shape=[1, num_units]),
                                      lookup_table[1:, :]), 0)
        outputs = tf.nn.embedding_lookup(lookup_table, position_ind)

        if scale:
            outputs = outputs * num_units**0.5

        return outputs

def multihead_attention(queries,
                        keys,
                        num_units=None,
                        num_heads=8,
                        dropout_rate=0,
                        is_training=True,
                        causality=False,
                        scope="multihead_attention",
                        reuse=None):
    '''Applies multihead attention.

    Args:
      queries: A 3d tensor with shape of [N, T_q, C_q].
      keys: A 3d tensor with shape of [N, T_k, C_k].
      num_units: A scalar. Attention size.
      dropout_rate: A floating point number.
      is_training: Boolean. Controller of mechanism for dropout.
      causality: Boolean. If true, units that reference the future are masked.
      num_heads: An int. Number of heads.
      scope: Optional scope for `variable_scope`.
      reuse: Boolean, whether to reuse the weights of a previous layer
        by the same name.

    Returns
      A 3d tensor with shape of (N, T_q, C)
    '''
    with tf.variable_scope(scope, reuse=reuse):
        # Set the fall back option for num_units
        if num_units is None:
            num_units = queries.get_shape().as_list[-1]

        # Linear projections
        Q = tf.layers.dense(queries, num_units, activation=tf.nn.relu) # (N, T_q, C)
        K = tf.layers.dense(keys, num_units, activation=tf.nn.relu) # (N, T_k, C)
        V = tf.layers.dense(keys, num_units, activation=tf.nn.relu) # (N, T_k, C)

        # Split and concat
        Q_ = tf.concat(tf.split(Q, num_heads, axis=2), axis=0) # (h*N, T_q, C/h)
        K_ = tf.concat(tf.split(K, num_heads, axis=2), axis=0) # (h*N, T_k, C/h)
        V_ = tf.concat(tf.split(V, num_heads, axis=2), axis=0) # (h*N, T_k, C/h)

        # Multiplication
        outputs = tf.matmul(Q_, tf.transpose(K_, [0, 2, 1])) # (h*N, T_q, T_k)

        # Scale
        outputs = outputs / (K_.get_shape().as_list()[-1] ** 0.5)

        # Key Masking
        key_masks = tf.sign(tf.abs(tf.reduce_sum(keys, axis=-1))) # (N, T_k)
        key_masks = tf.tile(key_masks, [num_heads, 1]) # (h*N, T_k)
        key_masks = tf.tile(tf.expand_dims(key_masks, 1), [1, tf.shape(queries)[1], 1]) # (h*N, T_q, T_k)

        paddings = tf.ones_like(outputs)*(-2**32+1)
        outputs = tf.where(tf.equal(key_masks, 0), paddings, outputs) # (h*N, T_q, T_k)

        # Causality = Future blinding
        if causality:
            diag_vals = tf.ones_like(outputs[0, :, :]) # (T_q, T_k)
            tril = tf.linalg.LinearOperatorLowerTriangular(diag_vals).to_dense() # (T_q, T_k)
            masks = tf.tile(tf.expand_dims(tril, 0), [tf.shape(outputs)[0], 1, 1]) # (h*N, T_q, T_k)

            paddings = tf.ones_like(masks)*(-2**32+1)
            outputs = tf.where(tf.equal(masks, 0), paddings, outputs) # (h*N, T_q, T_k)

        # Activation
        outputs = tf.nn.softmax(outputs) # (h*N, T_q, T_k)

        # Query Masking
        query_masks = tf.sign(tf.abs(tf.reduce_sum(queries, axis=-1))) # (N, T_q)
        query_masks = tf.tile(query_masks, [num_heads, 1]) # (h*N, T_q)
        query_masks = tf.tile(tf.expand_dims(query_masks, -1), [1, 1, tf.shape(keys)[1]]) # (h*N, T_q, T_k)
        outputs *= query_masks # broadcasting. (N, T_q, C)

        # Dropouts
        outputs = tf.layers.dropout(outputs, rate=dropout_rate, training=tf.convert_to_tensor(is_training))

        # Weighted sum
        outputs = tf.matmul(outputs, V_) # ( h*N, T_q, C/h)

        # Restore shape
        outputs = tf.concat(tf.split(outputs, num_heads, axis=0), axis=2 ) # (N, T_q, C)

        # Residual connection
        outputs += queries
#         print('outputs', outputs)
        # Normalize
        outputs = normalize(outputs) # (N, T_q, C)

    return outputs

def feedforward(inputs,
                num_units=[2048, 512],
                scope="multihead_attention",
                reuse=None):
    '''Point-wise feed forward net.

    Args:
      inputs: A 3d tensor with shape of [N, T, C].
      num_units: A list of two integers.
      scope: Optional scope for `variable_scope`.
      reuse: Boolean, whether to reuse the weights of a previous layer
        by the same name.

    Returns:
      A 3d tensor with the same shape and dtype as inputs
    '''
    with tf.variable_scope(scope, reuse=reuse):
        # Inner layer
        params = {"inputs": inputs, "filters": num_units[0], "kernel_size": 1,
                  "activation": tf.nn.relu, "use_bias": True}
        outputs = tf.layers.conv1d(**params)

        # Readout layer
        params = {"inputs": outputs, "filters": num_units[1], "kernel_size": 1,
                  "activation": None, "use_bias": True}
        outputs = tf.layers.conv1d(**params)

        # Residual connection
        outputs += inputs

        # Normalize
        outputs = normalize(outputs)

    return outputs

def label_smoothing(inputs, epsilon=0.1):
    '''Applies label smoothing. See https://arxiv.org/abs/1512.00567.

    Args:
      inputs: A 3d tensor with shape of [N, T, V], where V is the number of vocabulary.
      epsilon: Smoothing rate.

    For example,

    ```
    import tensorflow as tf
    inputs = tf.convert_to_tensor([[[0, 0, 1],
       [0, 1, 0],
       [1, 0, 0]],

      [[1, 0, 0],
       [1, 0, 0],
       [0, 1, 0]]], tf.float32)

    outputs = label_smoothing(inputs)

    with tf.Session() as sess:
        print(sess.run([outputs]))

    >>
    [array([[[ 0.03333334,  0.03333334,  0.93333334],
        [ 0.03333334,  0.93333334,  0.03333334],
        [ 0.93333334,  0.03333334,  0.03333334]],

       [[ 0.93333334,  0.03333334,  0.03333334],
        [ 0.93333334,  0.03333334,  0.03333334],
        [ 0.03333334,  0.93333334,  0.03333334]]], dtype=float32)]
    ```
    '''
    K = inputs.get_shape().as_list()[-1] # number of channels
    return ((1-epsilon) * inputs) + (epsilon / K)

def model_NN1(features, labels, mode):
    #Implement 3 seperate preprocessing and combine them together
    # implementation of NN with 3 inputs: https://stackoverflow.com/questions/40318457/how-to-build-a-multiple-input-graph-with-tensor-flow

    # First state preprocessing
    input_layer_state1 = tf.reshape(features["state1"], [-1, 49, 11, 1])
    conv1_state1 = tf.layers.conv2d(
        inputs=input_layer_state1,
        filters=32,
        kernel_size=[5, 5],
        padding="same",
        activation=tf.nn.relu)
    # Output Tensor Shape: [batch_size, 49, 11, 32]
    pool1_state1 = tf.layers.max_pooling2d(inputs=conv1_state1, pool_size=[2, 2], strides=2)
    # Output Tensor Shape: [batch_size, 24, 6, 32]
    conv2_state1 = tf.layers.conv2d(
        inputs=pool1_state1,
        filters=64,
        kernel_size=[5, 5],
        padding="same",
        activation=tf.nn.relu)
    # Output Tensor Shape: [batch_size, 24, 5, 64]
    pool2_state1 = tf.layers.max_pooling2d(inputs=conv2_state1, pool_size=[2, 2], strides=2)
    pool2_flat_state1 = tf.reshape(pool2_state1, [-1, 12 * 2 * 64])

    # Second state preprocessing
    # Outputs equal to state1
    input_layer_state2 = tf.reshape(features["state2"], [-1, 49, 11, 1])
    conv1_state2 = tf.layers.conv2d(
        inputs=input_layer_state2,
        filters=32,
        kernel_size=[5, 5],
        padding="same",
        activation=tf.nn.relu)
    pool1_state2 = tf.layers.max_pooling2d(inputs=conv1_state2, pool_size=[2, 2], strides=2)
    conv2_state2 = tf.layers.conv2d(
        inputs=pool1_state2,
        filters=64,
        kernel_size=[5, 5],
        padding="same",
        activation=tf.nn.relu)

    pool2_state2 = tf.layers.max_pooling2d(inputs=conv2_state2, pool_size=[2, 2], strides=2)
    pool2_flat_state2 = tf.reshape(pool2_state2, [-1, 12 * 2 * 64])

    whole_model = tf.concat([pool2_flat_state2, pool2_flat_state1], 1) # (1,3072)

    dense = tf.layers.dense(inputs=whole_model, units=1024, activation=tf.nn.relu) #(1, 1024)

    #################################################################
    #################################################################
    #################################################################

    decoder_inputs = tf.concat((tf.ones_like(labels[:, :1])*2, labels[:, :-1]), -1) #(1,120)
# #     print('decoder_inputs', decoder_inputs)

    # Encoder
    with tf.variable_scope("encoder"):

        ## Positional Encoding
        enc = positional_encoding(dense, num_units=512, zero_pad=False, scale=False, scope="enc_pe")

        ## Dropout
        enc = tf.layers.dropout(enc, rate=0.1, training=tf.convert_to_tensor(True))

        ## Blocks
        for i in range(6):
            with tf.variable_scope("num_blocks_{}".format(i)):
                ### Multihead Attention
                enc = multihead_attention(queries=enc, keys=enc, num_units=512, num_heads=8,
                                          dropout_rate=0.1, is_training=True, causality=False)

                ### Feed Forward
                enc = feedforward(enc, num_units=[4*512, 512])

    # Decoder
    with tf.variable_scope("decoder"):

        ## Positional Encoding
        dec = positional_encoding(decoder_inputs, num_units=512, zero_pad=False, scale=False, scope="dec_pe")

        ## Dropout
        dec = tf.layers.dropout(dec, rate=0.1, training=tf.convert_to_tensor(True))

        ## Blocks
        for i in range(6):
            with tf.variable_scope("num_blocks_{}".format(i)):
                ## Multihead Attention ( self-attention)
                dec = multihead_attention(queries=dec, keys=dec, num_units=512, num_heads=8, dropout_rate=0.1,
                                          is_training=True,causality=True, scope="self_attention")

                ## Multihead Attention ( vanilla attention)
                dec = multihead_attention(queries=dec, keys=dec, num_units=512, num_heads=8, dropout_rate=0.1,
                                          is_training=True,causality=False, scope="vanilla_attention")

                ## Feed Forward
                dec = feedforward(dec, num_units=[4*512, 512])

        # Final linear projection
        dec = tf.layers.dense(dec, 1024)
        dec = tf.reshape(dec, [1, -1])

    #################################################################
    #################################################################
    #################################################################

    # Output Tensor Shape: [batch_size, 1024]
    dropout = tf.layers.dropout(inputs=dec, rate=0.4, training=mode == tf.estimator.ModeKeys.TRAIN)
    print('dropout', dropout) #(1, 120, 1024)

    logits = tf.layers.dense(inputs=dropout, units=120) # shape of matrix -- 30 * 4

    predictions = {
        "classes": tf.argmax(input=logits, axis=1),
        "probabilities": tf.nn.softmax(logits, name="softmax_tensor")
    }

    if mode == tf.estimator.ModeKeys.PREDICT:
        return tf.estimator.EstimatorSpec(mode=mode, predictions=predictions)

    print('labels', labels, 'logits', logits) #labels: (1,120), logits: (1,120,120)

    loss = tf.losses.mean_squared_error(labels=labels, predictions=logits)

    if mode == tf.estimator.ModeKeys.TRAIN:
        optimizer = tf.train.GradientDescentOptimizer(learning_rate=0.001)
        train_op = optimizer.minimize(
            loss=loss,
            global_step=tf.train.get_global_step())
        return tf.estimator.EstimatorSpec(mode=mode, loss=loss, train_op=train_op)

    eval_metric_ops = {
        "accuracy": tf.metrics.accuracy(
            labels=labels, predictions=predictions["classes"])}

    return tf.estimator.EstimatorSpec(mode=mode, loss=loss, eval_metric_ops=eval_metric_ops)

def model_NN2(features, labels, mode):
    # Implement 2 seperate preprocessing of state and graph and combine them together

    # implementation of NN with 3 inputs: https://stackoverflow.com/questions/40318457/how-to-build-a-multiple-input-graph-with-tensor-flow

    # First state preprocessing
    input_layer_state1 = tf.reshape(features["state1"], [-1, 49, 11, 1])
    conv1_state1 = tf.layers.conv2d(
        inputs=input_layer_state1,
        filters=32,
        kernel_size=[5, 5],
        padding="same",
        activation=tf.nn.relu)
    pool1_state1 = tf.layers.max_pooling2d(inputs=conv1_state1, pool_size=[2, 2], strides=2)
    conv2_state1 = tf.layers.conv2d(
        inputs=pool1_state1,
        filters=64,
        kernel_size=[5, 5],
        padding="same",
        activation=tf.nn.relu)
    pool2_state1 = tf.layers.max_pooling2d(inputs=conv2_state1, pool_size=[2, 2], strides=2)
    pool2_flat_state1 = tf.reshape(pool2_state1, [-1, 12 * 2 * 64])

    # Graph preprocessing
    input_layer = tf.reshape(features["graph"], [-1, 4, 30, 1])
    conv1 = tf.layers.conv2d(
        inputs=input_layer,
        filters=32,
        kernel_size=[5, 5],
        padding="same",
        activation=tf.nn.relu)
    pool1 = tf.layers.max_pooling2d(inputs=conv1, pool_size=[2, 2], strides=2)
    conv2 = tf.layers.conv2d(
        inputs=pool1,
        filters=64,
        kernel_size=[5, 5],
        padding="same",
        activation=tf.nn.relu)
    conv2_flat = tf.reshape(conv2, [-1, 2 * 15 * 64])

    whole_model = tf.concat([pool2_flat_state1, conv2_flat], 1)

    dense = tf.layers.dense(inputs=whole_model, units=1024, activation=tf.nn.relu)
    dropout = tf.layers.dropout(
        inputs=dense, rate=0.4, training=mode == tf.estimator.ModeKeys.TRAIN)
    logits = tf.layers.dense(inputs=dropout, units=10)

    predictions = {
        "classes": tf.argmax(input=logits, axis=1),
        "probabilities": tf.nn.softmax(logits, name="softmax_tensor")
    }

    if mode == tf.estimator.ModeKeys.PREDICT:
        return tf.estimator.EstimatorSpec(mode=mode, predictions=predictions)

    onehot_labels = tf.one_hot(indices=tf.cast(labels, tf.int32), depth=10)
#     print('H')
    loss = tf.losses.softmax_cross_entropy(onehot_labels=onehot_labels, logits=logits)
#     print(loss.shape, onehot_labels.shape)

    # Configure the Training Op (for TRAIN mode)
    if mode == tf.estimator.ModeKeys.TRAIN:
        optimizer = tf.train.GradientDescentOptimizer(learning_rate=0.001)
        train_op = optimizer.minimize(
            loss=loss,
            global_step=tf.train.get_global_step())
        return tf.estimator.EstimatorSpec(mode=mode, loss=loss, train_op=train_op)

    if mode == tf.estimator.ModeKeys.PREDICT:
        return tf.estimator.EstimatorSpec(mode=mode, loss=loss, predictions=predictions["classes"])

    return tf.estimator.EstimatorSpec(mode=mode, loss=loss, predictions=predictions)

def main():

	# Print all possible environments in the Pommerman registry
	print(pommerman.registry)
	# sess = tf.Session()
	#sess = tf_debug.TensorBoardDebugWrapperSession(sess, 'localhost:6064')

	# Create a set of agents (exactly four)
	agent_list = [
	    agents.SimpleAgent(),
	    agents.SimpleAgent(),
	    agents.RandomAgent(),
	    agents.RandomAgent(),
	    # agents.DockerAgent("pommerman/simple-agent", port=12345),
	]
	env = pommerman.make('PommeFFACompetition-v0', agent_list)

	# Create the Estimator
	estimator_nn1 = tf.estimator.Estimator(model_fn=model_NN1, model_dir="/tmp/sa_nn1")
	# Set up logging for predictions
	tensors_to_logNN1 = {"probabilities": "softmax_tensor"}
	logging_hook_nn1 = tf.train.LoggingTensorHook(tensors=tensors_to_logNN1, every_n_iter=50)

	# Create the Estimator
	estimator_nn2 = tf.estimator.Estimator(model_fn=model_NN2, model_dir="/tmp/sa_nn2")
	# Set up logging for predictions
	tensors_to_logNN2 = {"probabilities": "softmax_tensor"}
	logging_hook_nn2 = tf.train.LoggingTensorHook(tensors=tensors_to_logNN2, every_n_iter=50)

	r_sum = np.zeros(1)
	for i in range(1):
	    # Make the "Free-For-All" environment using the agent list
	    env.reset()
	    # Run the episodes just like OpenAI Gym
	    for i_episode in range(1):
	        state = env.reset()

	        done = False
	        curr_state = None
	        prev_state = None
	        graph = np.random.rand(4, 30).astype("float32") + 0.0001
	#         print(graph)
	        pr_action = None
	        pr_pr_action = None

	        while not done:
	#             env.render()
	            actions = env.act(state)
	            state, reward, done, info = env.step(actions)
	            r_sum[i] += reward[0]

	            # as basic implementation I consider only one agent
	            prev_state = curr_state
	            curr_state = state

	            if pr_pr_action is not None:
	                # Train the model
	                for agent_num in range(4):
	                    train_input_NN2 = tf.estimator.inputs.numpy_input_fn(
	                        x={"state1": np.resize(
	                                state_to_matrix_with_action(curr_state[agent_num], action=pr_action[agent_num])\
	                                .astype("float32"), (1, 49*11)),
	                           "graph": np.resize(graph, (1, 4*30))
	                          },
	                        y=np.asarray([actions[agent_num]]),
	                        batch_size=1,
	                        num_epochs=None,
	                        shuffle=True)

	                    train_input_NN1 = tf.estimator.inputs.numpy_input_fn(
	                        x={"state1": np.resize(
	                                state_to_matrix_with_action(prev_state[agent_num], action=pr_pr_action[agent_num])\
	                                .astype("float32"), (1, 49 * 11)),
	                           "state2": np.resize(
	                                state_to_matrix_with_action(curr_state[agent_num], action=pr_action[agent_num])\
	                               .astype("float32"), (1, 49 * 11))},
	                        y=np.asmatrix(graph.flatten()),
	                        batch_size=1,
	                        num_epochs=None,
	                        shuffle=True)

	                    estimator_nn1.train(
	                        input_fn=train_input_NN1,
	                        steps=200,
	                        hooks=[logging_hook_nn1])

	                    estimator_nn2.train(
	                        input_fn=train_input_NN2,
	                        steps=200,
	                        hooks=[logging_hook_nn2])
	                    predictions = estimator_nn2.predict(input_fn=train_input_NN2)
	                    #next_action = np.array(list(p['classes'] for p in predictions))
	            pr_pr_action = pr_action
	            pr_action = actions
	        if i_episode > 300:
	            break
	    print('Game {} finished'.format(i))

	np.savetxt('result_2simple_2random.csv', r_sum, fmt='%1.4e')
	env.close()

if __name__ == '__main__':
	main()