import numpy as np
#import matplotlib.pyplot as plt
#import tensorflow as tf
#from tensorflow.keras.layers import *
import pickle
import math
#from tensorflow.keras import backend as K
import sys
import cv2
import imgaug.augmenters as iaa
from PIL import Image
def part_label_map():
  bird_labels = {'head':1, 'leye':2, 'reye':3, 'beak':4, 'torso':5, 'neck':6, 'lwing':7, 'rwing':8, 'lleg':9, 'lfoot':10, 'rleg':11, 'rfoot':12, 'tail':13}
  cat_labels = {'head':1, 'leye':2, 'reye':3, 'lear':4, 'rear':5, 'nose':6, 'torso':7, 'neck':8, 'lfleg':9, 'lfpa':10, 'rfleg':11, 'rfpa':12, 'lbleg':13, 'lbpa':14, 'rbleg':15, 'rbpa':16, 'tail':17}
  cow_labels = {'head':1, 'leye':2, 'reye':3, 'lear':4, 'rear':5, 'muzzle':6, 'lhorn':7, 'rhorn':8, 'torso':9, 'neck':10, 'lfuleg':11, 'lflleg':12, 'rfuleg':13, 'rflleg':14, 'lbuleg':15, 'lblleg':16, 'rbuleg':17, 'rblleg':18, 'tail':19}
  dog_labels = {'head':1, 'leye':2, 'reye':3, 'lear':4, 'rear':5, 'nose':6, 'torso':7, 'neck':8, 'lfleg':9, 'lfpa':10, 'rfleg':11, 'rfpa':12, 'lbleg':13, 'lbpa':14, 'rbleg':15, 'rbpa':16, 'tail':17, 'muzzle':18}
  horse_labels = {'head':1, 'leye':2, 'reye':3, 'lear':4, 'rear':5, 'muzzle':6, 'lfho':7, 'rfho':8, 'torso':9, 'neck':10, 'lfuleg':11, 'lflleg':12, 'rfuleg':13, 'rflleg':14, 'lbuleg':15, 'lblleg':16, 'rbuleg':17, 'rblleg':18, 'tail':19, 'lbho':20, 'rbho':21}
  9
  bottle_labels = {'cap':1, 'body':2}
  person_labels = {'head':1, 'leye':2,  'reye':3, 'lear':4, 'rear':5, 'lebrow':6, 'rebrow':7,  'nose':8,  'mouth':9,  'hair':10, 'torso':11, 'neck': 12, 'llarm': 13, 'luarm': 14, 'lhand': 15, 'rlarm':16, 'ruarm':17, 'rhand': 18, 'llleg': 19, 'luleg':20, 'lfoot':21, 'rlleg':22, 'ruleg':23, 'rfoot':24}
  aeroplane_labels = {'body': 1, 'stern': 2, 'lwing': 3, 'rwing':4, 'tail':5}
  for ii in range(1, 10):
      aeroplane_labels['engine_{}'.format(ii)] = 5+ii
  for ii in range(1, 10):
      aeroplane_labels['wheel_{}'.format(ii)] = 14+ii
  motorbike_labels = {'fwheel': 1, 'bwheel': 2, 'handlebar': 3, 'saddle': 4}
  for ii in range(0,10):
      motorbike_labels['headlight_{}'.format(ii+1)] = 5+ii
  motorbike_labels['body']=15
  bicycle_labels = {'fwheel': 1, 'bwheel': 2, 'saddle': 3, 'handlebar': 4, 'chainwheel': 5}
  for ii in range(0,10):
      bicycle_labels['headlight_{}'.format(ii+1)] = 6+ii
  bicycle_labels['body']=16
  sheep_labels = cow_labels

max_num_node = 24
canvas_size = 550

flip_bird = [1, 3, 2, 4, 5, 6, 10, 7, 11, 12, 9, 10, 13]
flip_cow = [1, 3, 2, 5, 4, 6, 10, 7, 9, 10, 13, 14, 11, 12, 17, 110, 15, 16, 19]
flip_cat = [1, 3, 2, 5, 4, 6, 7, 8, 11, 12, 9, 10, 15, 16, 13, 14, 17]
flip_dog = [1, 3, 2, 5, 4, 6, 7, 8, 11, 12, 9, 10, 15, 16, 13, 14, 17, 18]
flip_horse = [1, 3, 2, 5, 4, 6, 8, 7, 9, 10, 13, 14, 11, 12, 17, 18, 15, 16, 19, 21, 20]
flip_person = [1, 3, 2, 5, 4, 7, 6, 8, 9, 10, 11, 12, 16, 17, 18, 13, 14, 15, 22, 23, 24, 19, 20, 21]
flip_sheep = flip_cow

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

bird_tree = {}
bird_tree[0] = [1, 2, 3, 4, 5]
bird_tree[1] = [0, 2, 3]
bird_tree[2] = [0, 1, 3]
bird_tree[3] = [0, 1, 2]
bird_tree[4] = [0, 5, 6, 7, 8, 10, 12]
bird_tree[5] = [0, 4]
bird_tree[6] = [7, 4]
bird_tree[7] = [6, 4]
bird_tree[8] = [4, 9]
bird_tree[9] = [8]
bird_tree[10] = [11, 4]
bird_tree[11] = [10]
bird_tree[12] = [4]

dog_tree = {}
dog_tree[0] = [1, 2, 3, 4, 5, 6, 7, 17]
dog_tree[1] = [2, 0, 3]
dog_tree[2] = [0, 1, 4]
dog_tree[3] = [0, 1]
dog_tree[4] = [0, 2]
dog_tree[5] = [0, 1, 2]
dog_tree[6] = [0, 8, 10, 12, 14, 16, 7]
dog_tree[7] = [0, 6]
dog_tree[8] = [9, 6]
dog_tree[9] = [8]
dog_tree[10] = [6, 11]
dog_tree[11] = [10]
dog_tree[12] = [13, 6]
dog_tree[13] = [12]
dog_tree[14] = [6, 15]
dog_tree[15] = [14]
dog_tree[16] = [6]
dog_tree[17] = [0]

cat_tree = {}
cat_tree[0] = [1, 2, 3, 4, 5, 6, 7]
cat_tree[1] = [2, 0, 3]
cat_tree[2] = [0, 1, 4]
cat_tree[3] = [0, 1]
cat_tree[4] = [0, 2]
cat_tree[5] = [0, 1, 2]
cat_tree[6] = [0, 8, 10, 12, 14, 16, 7]
cat_tree[7] = [0, 6]
cat_tree[8] = [9, 6]
cat_tree[9] = [8]
cat_tree[10] = [6, 11]
cat_tree[11] = [10]
cat_tree[12] = [13, 6]
cat_tree[13] = [12]
cat_tree[14] = [6, 15]
cat_tree[15] = [14]
cat_tree[16] = [6]

horse_tree = {}
horse_tree[0] = [1, 2, 3, 4, 5, 8, 9]
horse_tree[1] = [2, 0, 3]
horse_tree[2] = [0, 1, 4]
horse_tree[3] = [0, 1]
horse_tree[4] = [0, 2]
horse_tree[5] = [0, 1, 2]
horse_tree[6] = [11]  # lfho
horse_tree[7] = [13]  # rfho
horse_tree[8] = [0, 10, 12, 14, 16, 18]
horse_tree[9] = [0, 8]
horse_tree[10] = [8, 11, 12]
horse_tree[11] = [10, 6]
horse_tree[12] = [10, 8, 13]
horse_tree[13] = [7]
horse_tree[14] = [8, 15, 16]
horse_tree[15] = [14, 19]
horse_tree[16] = [14, 17]
horse_tree[17] = [16, 20]
horse_tree[18] = [8]
horse_tree[19] = [15]
horse_tree[20] = [17]

cow_tree = {}
cow_tree[0] = [1, 2, 3, 4, 5, 6, 7, 8, 9]
cow_tree[1] = [2, 0, 3, 5]
cow_tree[2] = [0, 1, 4, 5]
cow_tree[3] = [0, 1, 6]
cow_tree[4] = [0, 2, 7]
cow_tree[5] = [0, 1, 2]
cow_tree[6] = [0, 3]  # lfho
cow_tree[7] = [0, 4, 13]  # rfho
cow_tree[8] = [0, 9, 10, 12, 14, 16, 18]
cow_tree[9] = [0, 8]
cow_tree[10] = [8, 11, 12]
cow_tree[11] = [10, 6]
cow_tree[12] = [10, 8, 13]
cow_tree[13] = [7, 12]
cow_tree[14] = [8, 15, 16]
cow_tree[15] = [14]
cow_tree[16] = [8, 14, 17]
cow_tree[17] = [16]
cow_tree[18] = [8]

motorbike_tree = {}
motorbike_tree[0] = [14, 1, 2]
motorbike_tree[1] = [14, 0]
motorbike_tree[2] = [14, 0]
motorbike_tree[3] = [14]
motorbike_tree[4] = [14]
motorbike_tree[5] = [14]
motorbike_tree[6] = [14]
motorbike_tree[7] = [14]
motorbike_tree[8] = [14]
motorbike_tree[9] = [14]
motorbike_tree[10] = [14]
motorbike_tree[11] = [14]
motorbike_tree[12] = [14]
motorbike_tree[13] = [14]
motorbike_tree[14] = [0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13]

bicycle_tree = {}
bicycle_tree[0] = [15, 1, 3]
bicycle_tree[1] = [15, 0, 2, 4]
bicycle_tree[2] = [15, 1]
bicycle_tree[3] = [15, 0]
bicycle_tree[4] = [15, 1]
bicycle_tree[5] = [15]
bicycle_tree[6] = [15]
bicycle_tree[7] = [15]
bicycle_tree[8] = [15]
bicycle_tree[9] = [15]
bicycle_tree[10] = [15]
bicycle_tree[11] = [15]
bicycle_tree[12] = [15]
bicycle_tree[13] = [15]
bicycle_tree[14] = [15]
bicycle_tree[15] = [0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14]

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

tree = {'aeroplane': aeroplane_tree, 'motorbike': motorbike_tree, 'bicycle': bicycle_tree, 'person': person_tree,
        'cow': cow_tree, 'dog': dog_tree, 'cat': cat_tree, 'sheep': cow_tree, 'bird': bird_tree, 'horse': horse_tree}

object_names = ['cow', 'sheep', 'bird', 'person', 'cat', 'dog', 'horse', 'aeroplane', 'motorbike', 'bicycle']

class_dic = {'cow': 0, 'sheep': 1, 'bird': 2, 'person': 3, 'cat': 4, 'dog': 5, 'horse': 6, 'aeroplane': 7,
             'motorbike': 8, 'bicycle': 9, 'car': 10}

def get_pos(bx):
    temp_pos = []
    for i in bx:
        if i.tolist() != [0, 0, 0, 0]:
            temp_pos.append([1])
        elif i.tolist() == [0, 0, 0, 0]:
            temp_pos.append([0])

    return np.asarray(temp_pos)

colors = [(229, 184, 135), (0, 0, 255), (0, 255, 0), (255, 0, 0), (0, 255, 255), (255, 255, 0), (255, 0, 255),
          (130, 0, 75), (0, 128, 128), (128, 128, 0), (128, 128, 128), (0, 0, 0), (30, 105, 210),
          (30, 105 // 2, 210 // 2), (180, 105, 255), (180 // 2, 105 // 2, 255), (100, 100, 30), (0, 100 // 2, 20),
          (128, 0, 128), (30, 105, 210), (255 // 2, 105, 255), (180 // 2, 105, 255 // 2), (50, 100, 0),
          (229 // 2, 184, 135 // 2), (229, 184, 135), (0, 0, 255), (0, 255, 0), (255, 0, 0), (0, 255, 255),
          (255, 255, 0), (255, 0, 255), (130, 0, 75), (0, 128, 128), (128, 128, 0), (128, 128, 128), (0, 0, 0),
          (30, 105, 210), (30, 105 // 2, 210 // 2), (180, 105, 255), (180 // 2, 105 // 2, 255), (100, 100, 30),
          (0, 100 // 2, 20), (128, 0, 128), (30, 105, 210), (255 // 2, 105, 255), (180 // 2, 105, 255 // 2),
          (50, 100, 0), (229 // 2, 184, 135 // 2)]

def arrangement(a, b, object_name):
    if object_name == 'cow' or object_name == 'sheep':
        p = [10, 11, 18, 13, 12, 14, 16, 15, 17, 9, 0, 7, 3, 4, 5, 6, 1, 2, 8]
    elif object_name == 'bird':
        p = [10, 11, 12, 9, 0, 7, 3, 4, 5, 6, 1, 2, 8]
    elif object_name == 'person':
        p = [10, 11, 19, 18, 20, 22, 21, 23, 13, 12, 14, 16, 15, 17, 9, 0, 7, 3, 4, 5, 6, 1, 2, 8]
    elif object_name == 'cat':
        p = [10, 11, 13, 12, 14, 16, 15, 9, 0, 7, 3, 4, 5, 6, 1, 2, 8]
    elif object_name == 'dog':
        p = [10, 11, 13, 12, 14, 16, 15, 17, 9, 0, 7, 3, 4, 5, 6, 1, 2, 8]
    elif object_name == 'horse':
        p = [10, 11, 19, 18, 20, 13, 12, 14, 16, 15, 17, 9, 0, 7, 3, 4, 5, 6, 1, 2, 8]
    elif object_name == 'aeroplane':
        p = [10, 11, 19, 18, 20, 22, 21, 13, 12, 14, 16, 15, 17, 9, 0, 7, 3, 4, 5, 6, 1, 2, 8]
    elif object_name == 'car':
        p = [10, 11, 19, 18, 20, 22, 21, 23, 24, 25, 26, 27, 28, 13, 12, 14, 16, 15, 17, 9, 0, 7, 3, 4, 5, 6, 1, 2, 8]
    elif object_name == 'motorbike':
        p = [10, 11, 13, 12, 14, 9, 0, 7, 3, 4, 5, 6, 1, 2, 8]
    elif object_name == 'bicycle':
        p = [10, 11, 13, 12, 14, 15, 9, 0, 7, 3, 4, 5, 6, 1, 2, 8]
    else:
        print("error")
    return a[p], b[p]

def rearrange(lbl, bx, mask, object_name):
    if object_name == 'cow' or object_name == 'sheep':
        p = np.asarray([1, 3, 2, 5, 4, 6, 8, 7, 9, 10, 13, 14, 11, 12, 17, 18, 15, 16, 19]) - 1
    elif object_name == 'bird':
        p = np.asarray([1, 3, 2, 4, 5, 6, 8, 7, 11, 12, 9, 10, 13]) - 1
    elif object_name == 'person':
        p = np.asarray([1, 3, 2, 5, 4, 7, 6, 8, 9, 10, 11, 12, 16, 17, 18, 13, 14, 15, 22, 23, 24, 19, 20, 21]) - 1
    elif object_name == 'cat':
        p = np.asarray([1, 3, 2, 5, 4, 6, 7, 8, 11, 12, 9, 10, 15, 16, 13, 14, 17]) - 1
    elif object_name == 'dog':
        p = np.asarray([1, 3, 2, 5, 4, 6, 7, 8, 11, 12, 9, 10, 15, 16, 13, 14, 17, 18]) - 1
    elif object_name == 'horse':
        p = np.asarray([1, 3, 2, 5, 4, 6, 8, 7, 9, 10, 13, 14, 11, 12, 17, 18, 15, 16, 19, 21, 20]) - 1
    elif object_name == 'aeroplane':
        p = np.asarray([1, 3, 2, 5, 4, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16, 17, 18, 19, 20, 21, 22, 23]) - 1
    elif object_name == 'car':
        p = np.asarray(
            [1, 3, 2, 4, 5, 7, 6, 8, 9, 10, 11, 12, 13, 14, 15, 16, 17, 18, 19, 20, 21, 22, 23, 24, 25, 26, 27, 28,
             29]) - 1
    elif object_name == 'motorbike':
        p = np.asarray([1, 3, 2, 4, 5, 7, 6, 8, 9, 10, 11, 12, 13, 14, 15]) - 1
    elif object_name == 'bicycle':
        p = np.asarray([1, 3, 2, 4, 5, 7, 6, 8, 9, 10, 11, 12, 13, 14, 15, 16]) - 1
    else:
        print("error")
    return lbl[p], bx[p], mask[p]

def pad_along_axis(array, target_length, axis=0):
    pad_size = target_length - array.shape[axis]
    axis_nb = len(array.shape)

    if pad_size < 0:
        return array

    npad = [(0, 0) for x in range(axis_nb)]
    npad[axis] = (0, pad_size)

    b = np.pad(array, pad_width=npad, mode='constant', constant_values=0)

    return b
def clip_rgb(img):
    result = np.where(img < 0.2)
    listOfCoordinates = list(zip(result[0], result[1], result[2]))
    for cord in listOfCoordinates:
        img[cord] = 0
    return img

def bounder(img):
    result = np.where(img < 0.5)
    listOfCoordinates = list(zip(result[0], result[1]))
    for cord in listOfCoordinates:
        img[cord] = 0
    result1 = np.where(img >= 0.5)
    listOfCoordinates1 = list(zip(result1[0], result1[1]))
    for cord in listOfCoordinates1:
        img[cord] = 1
    return img

def add_images(canvas, img, ii):
    result = np.where(img != 0)
    listOfCoordinates = list(zip(result[0], result[1]))
    for cord in listOfCoordinates:
        canvas[cord] = ii
    return canvas

def add_images_rgb(canvas, img):
    #print("Canvas : " + str(canvas.shape))
    #print("Image :" + str(img.shape))
    #print(img)
    result = np.where(img != 0)
    listOfCoordinates = list(zip(result[0], result[1], result[2]))
    #print(len(listOfCoordinates))
    for cord in listOfCoordinates:
        canvas[cord] = img[cord]
    return canvas

label_to_color = {0: (1, 1, 1),
                  1: (0.941, 0.973, 1),
                  2: (0.98, 0.922, 0.843),
                  3: (0, 1, 1),
                  4: (0.498, 1, 0.831),
                  5: (0.941, 1, 1),
                  6: (0.961, 0.961, 0.863),
                  7: (1, 0.894, 0.769),
                  8: (0.251, 0.878, 0.816),
                  9: (1, 0.388, 0.278),
                  10: (0, 0, 1),
                  11: (0.541, 0.169, 0.886),
                  12: (0.647, 0.165, 0.165),
                  13: (0.871, 0.722, 0.529),
                  14: (0.373, 0.62, 0.627),
                  15: (0.498, 1, 0),
                  16: (0.824, 0.412, 0.118),
                  17: (1, 0.498, 0.314),
                  18: (0.392, 0.584, 0.929),
                  19: (0.275, 0.51, 0.706),
                  20: (0.863, 0.0784, 0.235),
                  21: (0, 1, 1),
                  22: (0, 0, 0.545),
                  23: (0.824, 0.706, 0.549),
                  24: (0.251, 0.878, 0.816)}

def label_2_image(img):
    rgb_img = np.zeros((img.shape[0], img.shape[1], 3))
    for key in label_to_color.keys():
        rgb_img[img == key] = label_to_color[key]
    return rgb_img

def make_rgb_mask(mask, box):
    box = box * canvas_size
    max_parts = box.shape[0]
    #print(box.shape)
    #print(mask.shape)
    xmax = int(np.max(box[:, 2]))
    ymax = int(np.max(box[:, 3]))
    w = xmax
    h = ymax
    print(xmax)
    print(ymax)
    canvas = np.zeros((int(ymax), int(xmax), 3), np.float32)
    # b_in, mx = arrangement(b_in, mx,object_name)
    for i in range(max_parts):
        new_mask = mask[i]
        #cv2.imwrite('mask_output' + str(i) + '.jpg', mask[i]*255)
        x_min, y_min, x_max, y_max = box[i][0], box[i][1], box[i][2], box[i][3]
        x_min = 0 if x_min < 0 else x_min
        y_min = 0 if y_min < 0 else y_min
        if x_max - x_min > 0 and y_max - y_min > 0:
            x, y, c = canvas[int(y_min):int(y_max), int(x_min):int(x_max),:].shape
            print(x)
            print(y)
            print(c)
            if x==0:
              print("{} {} {} {}".format(y_min,y_max, x_min, x_max))
           # print(canvas[int(y_min):int(y_max), int(x_min):int(x_max)].shape)
            if x > 0 and y > 0:
            #print(mask[i].shape)
              img = cv2.resize(clip_rgb(new_mask.copy()), dsize=(y, x))
              #print(img.shape)
              canvas[int(y_min):int(y_max), int(x_min):int(x_max)] = add_images_rgb(
                canvas[int(y_min):int(y_max), int(x_min):int(x_max)],
                img)

    #plt.imshow(canvas)
    #plt.savefig('test.jpg')
    #plt.show()

    # plt.imshow(label_2_image(canvas))
    # plt.show()
    # return label_2_image(canvas)
    return canvas

def make_mask(box, mask):
    b_in = np.copy(box)
    mx = np.copy(mask)
    max_parts = len(box)
    xmax = max(box[:, 2])
    ymax = max(box[:, 3])
    canvas = np.zeros((int(ymax), int(xmax)), np.float32)
    b_in, mx = arrangement(b_in, mx, object_name)
    for i in range(max_parts):
        x_min, y_min, x_max, y_max = b_in[i]
        if x_max - x_min > 0 and y_max - y_min > 0:
            x, y = canvas[int(y_min):int(y_max), int(x_min):int(x_max)].shape
            canvas[int(y_min):int(y_max), int(x_min):int(x_max)] = add_images(
                canvas[int(y_min):int(y_max), int(x_min):int(x_max)],
                cv2.resize(bounder(np.squeeze(mx[i])) * (i + 1), (y, x)), i + 1)
    plt.imshow(label_2_image(canvas))
    plt.show()
    return label_2_image(canvas)

def plot_image_bbx(bbx, image):
    canvas = np.copy(image)
    i = 0
    for coord in bbx:
        x_minp, y_minp, x_maxp, y_maxp = coord
        if [x_minp, y_minp, x_maxp, y_maxp] != [0, 0, 0, 0]:
            cv2.rectangle(canvas, ((x_minp), (y_minp)), ((x_maxp), (y_maxp)), colors[i], 4)
        i = i + 1
    plt.imshow(canvas)
    plt.show()
    return canvas

def flip_mask(mask):
    mx = np.copy(mask)
    for i in range(len(mx)):
        mx[i] = mx[i][:, ::-1]
    return mx

def flip_bbx(label, bbx, img):
    bx = np.copy(bbx)
    x_min = min(bbx[:, 0])
    y_min = min(bbx[:, 1])
    x_max = max(bbx[:, 2])
    y_max = max(bbx[:, 3])
    img_center = np.asarray([((x_max + x_min) / 2), ((y_max + y_min) / 2)])
    img_center = np.hstack((img_center, img_center))
    bx[:, [0, 2]] += 2 * (img_center[[0, 2]] - bx[:, [0, 2]])
    box_w = abs(bx[:, 0] - bx[:, 2])
    bx[:, 0] -= box_w
    bx[:, 2] += box_w
    for i in range(len(label)):
        if sum(label[i]) == 0:
            bx[i][0] = 0
            bx[i][1] = 0
            bx[i][2] = 0
            bx[i][3] = 0
    return bx

def flip_data_instance(label, box, mask, image, object_name):
    bx = np.copy(flip_bbx(label, box, image))
    mx = np.copy(flip_mask(mask))
    ix = np.copy(image[:, ::-1])
    lx = np.copy(label)
    lx, bx, mx = rearrange(lx, bx, mx, object_name)
    return lx, bx, mx, ix

def cordinates(img):
    y_min = 0
    y_max = 0
    x_min = 0
    x_max = 0

    for i in img:
        if np.count_nonzero(i) is not 0:
            break
        y_min += 1

    for i in img.T:
        if np.count_nonzero(i) is not 0:
            break
        x_min += 1

    for i in img[::-1]:
        if np.count_nonzero(i) is not 0:
            break
        y_max += 1
    y_max = img.shape[0] - y_max - 1

    for i in img.T[::-1]:
        if np.count_nonzero(i) is not 0:
            break
        x_max += 1
    x_max = img.shape[1] - x_max - 1

    return x_min, y_min, x_max, y_max

def rotate_im(image, angle):
    # grab the dimensions of the image and then determine the
    # centre
    (h, w) = image.shape[:2]
    (cX, cY) = (w // 2, h // 2)

    # grab the rotation matrix (applying the negative of the
    # angle to rotate clockwise), then grab the sine and cosine
    # (i.e., the rotation components of the matrix)
    M = cv2.getRotationMatrix2D((cX, cY), angle, 1.0)
    cos = np.abs(M[0, 0])
    sin = np.abs(M[0, 1])

    # compute the new bounding dimensions of the image
    nW = int((h * sin) + (w * cos))
    nH = int((h * cos) + (w * sin))

    # adjust the rotation matrix to take into account translation
    M[0, 2] += (nW / 2) - cX
    M[1, 2] += (nH / 2) - cY

    # perform the actual rotation and return the image
    image = cv2.warpAffine(image, M, (nW, nH))

    #    image = cv2.resize(image, (w,h))
    return image

def get_corners(bboxes):
    width = (bboxes[:, 2] - bboxes[:, 0]).reshape(-1, 1)
    height = (bboxes[:, 3] - bboxes[:, 1]).reshape(-1, 1)

    x1 = bboxes[:, 0].reshape(-1, 1)
    y1 = bboxes[:, 1].reshape(-1, 1)

    x2 = x1 + width
    y2 = y1

    x3 = x1
    y3 = y1 + height

    x4 = bboxes[:, 2].reshape(-1, 1)
    y4 = bboxes[:, 3].reshape(-1, 1)

    corners = np.hstack((x1, y1, x2, y2, x3, y3, x4, y4))

    return corners

def clip_box(bbox, clip_box, alpha):
    ar_ = (bbox_area(bbox))
    x_min = np.maximum(bbox[:, 0], clip_box[0]).reshape(-1, 1)
    y_min = np.maximum(bbox[:, 1], clip_box[1]).reshape(-1, 1)
    x_max = np.minimum(bbox[:, 2], clip_box[2]).reshape(-1, 1)
    y_max = np.minimum(bbox[:, 3], clip_box[3]).reshape(-1, 1)

    bbox = np.hstack((x_min, y_min, x_max, y_max, bbox[:, 4:]))

    delta_area = ((ar_ - bbox_area(bbox)) / ar_)

    mask = (delta_area < (1 - alpha)).astype(int)

    bbox = bbox[mask == 1, :]

    return bbox

def rotate_box(corners, angle, cx, cy, h, w):
    corners = corners.reshape(-1, 2)
    corners = np.hstack((corners, np.ones((corners.shape[0], 1), dtype=type(corners[0][0]))))

    M = cv2.getRotationMatrix2D((cx, cy), angle, 1.0)

    cos = np.abs(M[0, 0])
    sin = np.abs(M[0, 1])

    nW = int((h * sin) + (w * cos))
    nH = int((h * cos) + (w * sin))
    # adjust the rotation matrix to take into account translation
    M[0, 2] += (nW / 2) - cx
    M[1, 2] += (nH / 2) - cy
    # Prepare the vector to be transformed
    calculated = np.dot(M, corners.T).T

    calculated = calculated.reshape(-1, 8)

    return calculated

def get_enclosing_box(corners):
    x_ = corners[:, [0, 2, 4, 6]]
    y_ = corners[:, [1, 3, 5, 7]]

    xmin = np.min(x_, 1).reshape(-1, 1)
    ymin = np.min(y_, 1).reshape(-1, 1)
    xmax = np.max(x_, 1).reshape(-1, 1)
    ymax = np.max(y_, 1).reshape(-1, 1)

    final = np.hstack((xmin, ymin, xmax, ymax, corners[:, 8:]))

    return final

def bbox_area(bbox):
    return (bbox[:, 2] - bbox[:, 0]) * (bbox[:, 3] - bbox[:, 1])

def rtt(angle, label, img, bboxes):
    w, h = img.shape[1], img.shape[0]
    cx, cy = w // 2, h // 2

    img = rotate_im(img, angle)

    corners = get_corners(bboxes)

    corners = np.hstack((corners, bboxes[:, 4:]))

    corners[:, :8] = rotate_box(corners[:, :8], angle, cx, cy, h, w)

    new_bbox = get_enclosing_box(corners)

    scale_factor_x = img.shape[1] / w

    scale_factor_y = img.shape[0] / h

    img = cv2.resize(img, (w, h), interpolation=cv2.INTER_NEAREST)

    new_bbox[:, :4] = np.true_divide(new_bbox[:, :4], [scale_factor_x, scale_factor_y, scale_factor_x, scale_factor_y])

    for i in range(len(label)):
        if (label[i]) == 0:
            new_bbox[i][0] = 0
            new_bbox[i][1] = 0
            new_bbox[i][2] = 0
            new_bbox[i][3] = 0

    return img, new_bbox

def render_mask(box, mask, angle):
    mx = np.copy(mask)
    b_in = np.copy(box)
    max_parts = len(box)
    xmax = max(box[:, 2])
    ymax = max(box[:, 3])
    temp_mx_list = []
    temp_bx_list = []
    for i in range(max_parts):
        canvas = np.zeros((int(ymax), int(xmax)), np.float32)
        x_min, y_min, x_max, y_max = b_in[i]
        if x_max - x_min > 0 and y_max - y_min > 0:
            x, y = canvas[int(y_min):int(y_max), int(x_min):int(x_max)].shape
            canvas[int(y_min):int(y_max), int(x_min):int(x_max)] = add_images(
                canvas[int(y_min):int(y_max), int(x_min):int(x_max)],
                cv2.resize(bounder(np.squeeze(mx[i])) * (i + 1), (y, x)), i + 1)
            canvas = rotate_im(canvas, angle)
            x_min, y_min, x_max, y_max = cordinates(canvas)
            # canvas = canvas[int(y_min):int(y_max), int(x_min):int(x_max)]
            # resized_cropped = np.expand_dims(cv2.resize(canvas, (64, 64)), axis = 3)
        temp_bx_list.append([x_min, y_min, x_max, y_max])
        # temp_mx_list.append(resized_cropped)
        # plt.imshow(canvas)
        # plt.show()
    return np.asarray(temp_bx_list, dtype="float32")

def scale(bbx, scaling_factor):
    height = max(bbx[:, 3])
    width = max(bbx[:, 2])

    pos = get_pos(bbx)

    fold_a = np.copy(bbx)
    fold_b = np.copy(bbx)
    fold_c = np.copy(bbx)
    fold_d = np.copy(bbx)

    scale_height = scaling_factor
    scale_width = scaling_factor

    fold_a[:, 0] = (fold_a[:, 0] - scale_width)
    fold_b[:, 1] = (fold_b[:, 1] - scale_height)
    fold_c[:, 2] = (fold_c[:, 2] + scale_width)
    fold_d[:, 3] = (fold_d[:, 3] + scale_height)

    return fold_a * pos, fold_b * pos, fold_c * pos, fold_d * pos

def centre_object(bbx, canvas_size):
    pos = get_pos(bbx)
    bx = np.copy(bbx)

    h, w = canvas_size

    h_o = max(bbx[:, 3])
    w_o = max(bbx[:, 2])

    h_shift = int(h / 2 - h_o / 2)
    w_shift = int(w / 2 - w_o / 2)

    bx[:, 0] = (bx[:, 0] + w_shift)
    bx[:, 1] = (bx[:, 1] + h_shift)
    bx[:, 2] = (bx[:, 2] + w_shift)
    bx[:, 3] = (bx[:, 3] + h_shift)

    return bx * pos

def append_labels(box):
    all_box = []
    for bbx in box:
        pos = get_pos(bbx)
        bbx = (((bbx / canvas_size))) * pos

        temp = []
        for bx in bbx:
            if bx.tolist() != [0, 0, 0, 0]:
                temp.append([1] + bx.tolist())
            else:
                temp.append([0] + bx.tolist())
        all_box.append(temp)
    return np.asarray(all_box)

def plot_bbx(bbx):

    bbx = bbx * canvas_size
    xmax = int(np.max(bbx[:, 2]))
    ymax = int(np.max(bbx[:, 3]))
    w = xmax
    h = ymax
    canvas = np.ones((ymax, xmax, 3), np.uint8) * 255
    bbx = bbx.tolist()

    i = 0
    for coord in bbx:
        x_minp, y_minp, x_maxp, y_maxp = coord
        if [x_minp, y_minp, x_maxp, y_maxp] != [0, 0, 0, 0]:
            cv2.rectangle(canvas, (int(x_minp), int(y_minp)), (int(x_maxp), int(y_maxp)), colors[i], 4)
        i = i + 1
    return canvas

def plot_bbx_new(bbx):
    canvas = np.ones((canvas_size, canvas_size, 3), np.uint8) * 255
    i = 0
    p = [ 18,19,20,21,22,23,6, 7, 16, 8, 9, 10, 11, 12, 13, 14, 15, 0, 17,1, 2, 3, 4, 5]
    bbx = bbx[p]
    for coord in bbx:
        x_minp, y_minp, x_maxp, y_maxp = coord
        if [x_minp, y_minp, x_maxp, y_maxp] != [0, 0, 0, 0]:
            cv2.rectangle(canvas, (int(x_minp), int(y_minp)), (int(x_maxp), int(y_maxp)), colors[i], 4)
        i = i + 1
    return canvas

def transform_bbx(bbx1):
    eps = 0.00001
    bbx = np.copy(bbx1)
    bxx = np.copy(bbx)

    bbx[:, 0] = np.exp(bbx[:, 0])
    bbx[:, 1] = np.exp(bbx[:, 1])
    bbx[:, 2] = np.exp(bbx[:, 2])
    bbx[:, 3] = np.exp(bbx[:, 3])

    bxx[:, 0] = bbx[:, 0]
    bxx[:, 1] = bbx[:, 1]
    bxx[:, 2] = bbx[:, 0] + (bbx[:, 3])
    bxx[:, 3] = bbx[:, 1] + (bbx[:, 2])

    return bxx

def make_full_mask_list(mask_list, size_of_list):
    ml = []
    for mask in mask_list:
        masks = [np.zeros((64, 64, 1))] * size_of_list
        masks[0:len(mask)] = mask
        ml.append(np.asarray(masks))

    return np.asarray(ml)

    class_v = {}

def get_rgbmask(box, mask, images, object_name):
    """
    box: Nx24x5 bounding boxes
    mask: Nx24x64x64x1 part separated masks
    images: List of N rgb images
    """
    print(box.shape[0])
    print(len(images))

    box = box.astype('float32')

    cropped = []  # list of part separated rgb for all objects: Nx24x64x64x3

    for idx, img in enumerate(images):

        images[idx] = img.astype('uint8')  # change to uint8 due to memory issues

        mmx = mask[idx, :, :, :, 0]  # full-mask of object at position idx

        bbx = box[idx, :, 1:] * canvas_size  # xmin,ymin,xmax,ymax format
        avail = box[idx, :, 0] == 1  # part availability

        bbx = (np.round(bbx)).astype('int')

        h, w = img.shape[:2]
        parts = []  # list of part separated rgb for one object: 24x64x64x3

        for j in range(24):
            x1, y1, x2, y2 = bbx[j]

            #####Handle zero-height/width due to rounding errors -- inflate bb by 1 or 2 pixel####
            if (x1 == x2):
                x1 -= 1
                x2 += 1
            if (y1 == y2):
                y1 -= 1
                y2 += 1

            x1 = max(0, x1)
            y1 = max(0, y1)
            x2 = min(w, x2)
            y2 = min(h, y2)

            #print("{} {} {} {}".format(x1, y1, x2, y2))
            ##############################################################################

            try:
                if avail[j]:
                    part = img[y1:y2, x1:x2]  # crop a part from rgb image
                    part = cv2.resize(part, (64, 64))  # resize to part-mask dimensions
                    part[mmx[j] == 0] = np.array([0, 0, 0])  # keep only required part of rgb using mask
                else:
                    part = np.zeros((64, 64, 3))  # black part-rgb if some part is not there in rgb
            except:
                part = np.zeros(
                    (64, 64, 3))  # Handle exceptional cases of negaitive values in bb: make a black part rgb
                mask[idx, j, :, :, 0] = np.zeros((64, 64))  # Black part mask
                box[idx, j] = np.zeros(5)  # delete gibberish bounding box-- zero out

                print(
                    'object {} at index {} skipped due to error in bounding box dimensions :({},{}),({},{})....'.format(
                        object_name, idx, x1, x2, y1, y2))

            parts.append(part)

        """   #uncomment to change bb values to a fraction of rgb image instead of fraction of canvas size
        box[idx,:,1:]*=canvas_size
        box[idx,:,1:]/=np.array([h,w,h,w])
        """  # uncomment to change bb values to a fraction of rgb image instead of fraction of canvas size

        parts = np.array(parts).astype('uint8')  # for memory issues change to uint8
        cropped.append(parts)

    cropped = np.array(cropped)

    return box, mask.astype('bool'), cropped, images

def visualise(num_samples):
    idx = np.random.randint(0, class_v.shape[0], num_samples)
    for i in idx:
        img = rgbs[i]
        for j in range(24):
            if X_train[i][j][0] == 1:
                h, w = img.shape[:2]
                box = np.round(X_train[i, j, 1:] * canvas_size).astype('uint16')
                st = tuple(box[:2])
                en = tuple(box[2:])
                img1 = img.copy()
                cv2.rectangle(img1, pt1=st, pt2=en, color=(255, 0, 0), thickness=2)
                plt.imshow(img1)
                plt.show()
                plt.imshow(masks[i, j, :, :, 0])
                plt.show()
                plt.imshow(rgb_masks[i, j])
                plt.show()

def make_data(object_class_names, train_mode):
    class_v = {}
    X_train = {}
    nX_train = {}
    adj_train = {}
    masks = {}
    rgb_masks = {}
    rgbs = {}
    train_split = 0.75
    valid_split = 0.10
    test_split = 0.15
    for object_name in object_class_names:
        path = 'data/data_custom/'
        print(type(object_class_names))
        if object_name == 'motorbike' or object_name == 'bicycle':
            path = 'data/motor_cycle/'
        with open(path + object_name + '_images', 'rb') as f:
            o_images = pickle.load(f)
        with open(path + object_name + '_part_separated_labels', 'rb') as f:
            o_labels = pickle.load(f)
        with open(path + object_name + '_part_separated_masks', 'rb') as f:
            o_masks = pickle.load(f)
        with open(path + object_name + '_part_separated_bbx', 'rb') as f:
            o_bbx = pickle.load(f)

        samples = o_labels.shape[0]
        print("samples = " + str(samples))
        train_set_limit = int(train_split * samples)
        valid_set_limit = int(valid_split * samples)
        test_set_limit = int(test_split * samples)

        if train_mode=='train':
          label = o_labels[0:train_set_limit]
          box = o_bbx[0:train_set_limit]
          mask = o_masks[0:train_set_limit]
          image = o_images[0:train_set_limit]
        if train_mode=='valid':
          label = o_labels[train_set_limit:train_set_limit + valid_set_limit]
          box = o_bbx[train_set_limit:train_set_limit + valid_set_limit]
          mask = o_masks[train_set_limit:train_set_limit + valid_set_limit]
          image = o_images[train_set_limit:train_set_limit + valid_set_limit]
        if train_mode=='test':
          label = o_labels[train_set_limit + valid_set_limit:]
          box = o_bbx[train_set_limit + valid_set_limit:]
          mask = o_masks[train_set_limit + valid_set_limit:]
          image = o_images[train_set_limit + valid_set_limit:]

        print("Masks= " + str(len(label)))
        print("Images= " + str(len(image)))

        max_parts = len(label[0])

        flipped_label = []
        flipped_box = []
        flipped_mask = []
        flipped_image = []

        angle = 3
        if train_mode=='train':
         for l, b, m, i in zip(label, box, mask, image):
            ll, bb, mm, ii = flip_data_instance(l, b, m, i, object_name)
            flipped_label.append(ll)
            flipped_box.append(bb)
            flipped_mask.append(mm)
            flipped_image.append(ii)

         flipped_label = np.asarray(flipped_label)
         flipped_box = np.asarray(flipped_box)
         flipped_mask = np.asarray(flipped_mask)
         flipped_image = np.asarray(flipped_image)

         label = np.concatenate((label, flipped_label), axis=0)

         box = np.concatenate((box, flipped_box), axis=0)
         mask = np.concatenate((mask, flipped_mask), axis=0)
         image = np.concatenate((image, flipped_image), axis=0)

        box = append_labels(box)

        box = pad_along_axis(box, 24, axis=1)
        mask = make_full_mask_list(mask, 24)

        numparts = 24
        final_input = []
        for i, x in enumerate(box):
            l = numparts
            temp = np.zeros((l, l), dtype=np.float32)
            for j in range(l):
                temp[j][j] = 1
                if x[j][0] == 1:
                    for y in tree[object_name][j]:
                        if x[y][0] == 1:
                            temp[j][y] = 1
                            temp[y][j] = 1
            final_input.append(temp)
        adj = np.asarray(final_input)

        box, mask, rgb_mask, image = get_rgbmask(box, mask, image, object_name)
        class_v[object_name] = (np.asarray([np.eye(10)[class_dic[object_name]]] * len(box))).astype('bool')
        X_train[object_name] = box
        adj_train[object_name] = (adj).astype('bool')
        masks[object_name] = mask.astype('bool')
        rgb_masks[object_name] = rgb_mask.astype('uint8')
        rgbs[object_name] = image

    class_v = np.concatenate(list(class_v.values()), axis=0)
    X_train = np.concatenate(list(X_train.values()), axis=0)
    rgb_imgs = np.concatenate(list(rgbs.values()), axis=0)
    masks = np.concatenate(list(masks.values()), axis=0)
    adj_train = np.concatenate(list(adj_train.values()), axis=0)
    rgb_masks = np.concatenate(list(rgb_masks.values()), axis=0)

    # np.save('masks.npy', masks)
    # np.save('X_train.npy', X_train)
    # np.save('rgb_masks.npy', rgb_mask)
    # np.save('rgbs.npy',  rgbs)
    # print("Part Masks Size = :" + str(masks.shape))
    # print("BBox Size = :" + str(X_train.shape))
    # print("RGB Part Size = :" + str(rgb_masks.shape))
    # print("RGB Shape = :" + str(rgb_imgs.shape))
    # print("Classes Shape " + str(class_v.shape))
    # print(class_v[0:5])
    if train_mode=='train':
     rand_indices = np.random.permutation(masks.shape[0])
     X_train = X_train[rand_indices]
     masks = masks[rand_indices]
     rgb_masks = rgb_masks[rand_indices]
     rgb_imgs = rgb_imgs[rand_indices]
     class_v = class_v[rand_indices]

    return X_train, masks, rgb_masks, rgb_imgs, class_v.astype(dtype=np.float32)

    # visualise(5)

def preprocess_make_data(object_class_names, train_mode, mode='boxVAE'):
    if train_mode=='train':
     return preprocess_make_data_augment(object_class_names, train_mode)
    X_train, masks, rgb_masks, rgb_imgs, class_v = make_data(object_class_names, train_mode)
    num_samples = X_train.shape[0]
    # print(X_train.shape)
    # print(masks.shape)
    # print(rgb_masks.shape)
    # print(rgb_imgs.shape)
    # print(class_v.shape)
    if mode=='boxVAE' and train_mode=='test':
      cls = pickle.load(open('data/class_vec_sketch_16','rb'))
      lbl = pickle.load(open('data/label_vec_sketch_16','rb'))
      bbx = pickle.load(open('data/bbx_gen_sketch_16','rb'))
     # nonzero_x1 = np.nonzero(bbx[:,:,0])
     # nonzero_y1 = np.nonzero(bbx[:,:,1])
     # nonzero_x2 = np.nonzero(bbx[:,:,2])
     # nonzero_y2 = np.nonzero(bbx[:,:,3])

     # xmin = np.min(bbx[nonzero_x1])
    #  ymin =  np.min(bbx[nonzero_y1])
    #  for i in range(bbx.shape[0]):
    #    xmin = 10.0
    #    ymin = 10.0
    #    for j in range(bbx.shape[1]):
    #      if np.sum(bbx[i][j])!=0.0:
    #        if bbx[i][j][0] < xmin:
    #          xmin = bbx[i][j][0]
    #        if bbx[i][j][1] < ymin:
    #          ymin = bbx[i][j][1]

        #print(bbx[i])
    #    for j in range(bbx.shape[1]):
    #      if np.sum(bbx[i][j])!=0.0:

    #        bbx[i,j,0] = bbx[i,j,0]-xmin
    #        bbx[i,j,2] = bbx[i,j,2]-xmin
    #        bbx[i,j,1] = bbx[i,j,1]-ymin
    #        bbx[i,j,3] = bbx[i,j,3]-ymin

      #print("box xmin {} ymin {} ".format(xmin, ymin))
      print("boxvae class {}".format(cls.shape))
      print("boxvae label {}".format(lbl.shape))
      print("boxvae bbx {}".format(bbx.shape))

      class_v = cls.astype(dtype=np.float32)
      X_train = np.concatenate((lbl,bbx), axis=2)
      #rgb_imgs = np.tile(rgb_imgs,(5,1,1,1))
      #rgb_masks = np.tile(rgb_masks,(5,1,1,1,1))
      #masks = np.tile(masks,(5,1,1,1,1))
      num_samples = X_train.shape[0]
      print("sample boxvae min {} max {}".format(np.min(X_train[0:100,:,1:]), np.max(X_train[0:100,:,1:])))

    data_dict = {}
    for i in range(num_samples):
        data_dict[str(i)] = {}
    dog_labels = {'head': 1, 'leye': 2, 'reye': 3, 'lear': 4, 'rear': 5, 'nose': 6, 'torso': 7, 'neck': 8, 'lfleg': 9,
                  'lfpa': 10, 'rfleg': 11, 'rfpa': 12, 'lbleg': 13, 'lbpa': 14, 'rbleg': 15, 'rbpa': 16, 'tail': 17,
                  'muzzle': 18}

    def get_key(val):
        for key, value in dog_labels.items():
            if val == value:
                return key
    for i in range(num_samples):
        if mode=='boxVAE':
          idx = i//5
        else:
          idx = i
        img = rgb_imgs[idx]
        canvas = np.zeros((img.shape[0], img.shape[1]), dtype=np.uint8)
        data_dict[str(i)]['image'] = img
        result = np.where(img != 0)
        listOfCoordinates = list(zip(result[0], result[1]))
        for cord in listOfCoordinates:
          canvas[cord] = 255
        data_dict[str(i)]['image_mask'] = canvas

        indices = np.where(X_train[i, :, 0] == 1)
        label_names = [get_key(val + 1) for val in indices[0]]
        labels = [(val + 1) for val in indices[0]]
        data_dict[str(i)]['boxes'] =  X_train[i][indices][:, 1:]
        data_dict[str(i)]['label_names'] = label_names
        data_dict[str(i)]['labels'] = labels
        data_dict[str(i)]['class'] = class_v[i]
        data_dict[str(i)]['masks'] = None#masks[i][indices]

    #print("Image size " + (str(data_dict['70']['image'].size)))
    # print(data_dict['70']['image'].shape)
    # print(data_dict['65']['image'].shape)
    # print(data_dict['10']['image'].shape)
    # print(data_dict['50']['image'].shape)

    return data_dict

def preprocess_make_data_augment(object_class_names, train_mode):
    X_train, masks, rgb_masks, rgb_imgs, class_v = make_data(object_class_names, train_mode)

    num_samples = X_train.shape[0]
    # print(X_train.shape)
    # print(masks.shape)
    # print(rgb_masks.shape)
    # print(rgb_imgs.shape)
    # print(class_v.shape)
    data_dict = {}
    for i in range(num_samples * 4):
        data_dict[str(i)] = {}
    dog_labels = {'head': 1, 'leye': 2, 'reye': 3, 'lear': 4, 'rear': 5, 'nose': 6, 'torso': 7, 'neck': 8, 'lfleg': 9,
                  'lfpa': 10, 'rfleg': 11, 'rfpa': 12, 'lbleg': 13, 'lbpa': 14, 'rbleg': 15, 'rbpa': 16, 'tail': 17,
                  'muzzle': 18}

    def get_key(val):
        for key, value in dog_labels.items():
            if val == value:
                return key
    for i in range(num_samples):
        data_dict[str(i)]['image'] = rgb_imgs[i]
        indices = np.where(X_train[i, :, 0] == 1)
        label_names = [get_key(val + 1) for val in indices[0]]
        labels = [(val + 1) for val in indices[0]]
        data_dict[str(i)]['boxes'] =  X_train[i][indices][:, 1:]
        data_dict[str(i)]['label_names'] = label_names
        data_dict[str(i)]['labels'] = labels
        data_dict[str(i)]['class'] = class_v[i]
        data_dict[str(i)]['masks'] = masks[i][indices]
        data_dict[str(i + num_samples)]['image'] = iaa.Add((-20, 20), per_channel=0.5)(image=rgb_imgs[i])
        data_dict[str(i + 2 * num_samples)]['image'] = iaa.AddToHueAndSaturation((-30, 30))(image=rgb_imgs[i])
        data_dict[str(i + 3 * num_samples)]['image'] = iaa.LinearContrast((0.5, 2.0), per_channel=0.5)(image=rgb_imgs[i])
        data_dict[str(i + num_samples)]['label_names'] = data_dict[str(i)]['label_names']
        data_dict[str(i + num_samples)]['boxes'] = data_dict[str(i)]['boxes']
        data_dict[str(i + num_samples)]['labels'] = data_dict[str(i)]['labels']
        data_dict[str(i + num_samples)]['class'] = data_dict[str(i)]['class']
        data_dict[str(i + num_samples)]['masks'] = data_dict[str(i)]['masks']
        data_dict[str(i + 2 * num_samples)]['label_names'] = data_dict[str(i)]['label_names']
        data_dict[str(i + 2 * num_samples)]['boxes'] = data_dict[str(i)]['boxes']
        data_dict[str(i + 2 * num_samples)]['labels'] = data_dict[str(i)]['labels']
        data_dict[str(i + 2 * num_samples)]['class'] = data_dict[str(i)]['class']
        data_dict[str(i + 2 * num_samples)]['masks'] = data_dict[str(i)]['masks']
        data_dict[str(i + 3 * num_samples)]['label_names'] = data_dict[str(i)]['label_names']
        data_dict[str(i + 3 * num_samples)]['boxes'] = data_dict[str(i)]['boxes']
        data_dict[str(i + 3 * num_samples)]['labels'] = data_dict[str(i)]['labels']
        data_dict[str(i + 3 * num_samples)]['class'] = data_dict[str(i)]['class']
        data_dict[str(i + 3 * num_samples)]['masks'] = data_dict[str(i)]['masks']

    print(data_dict['70']['image'].shape)
    print(data_dict['65']['image'].shape)
    print(data_dict['10']['image'].shape)
    print(data_dict['50']['image'].shape)
    print(type(data_dict['70']['image']))
    img = Image.fromarray(data_dict[str(70)]['image'])
    img.save('test_input.png')
    img = Image.fromarray(data_dict[str(70 + num_samples)]['image'])
    img.save('test_brightness.png')
    img = Image.fromarray(data_dict[str(70 + 2 * num_samples)]['image'])
    img.save('test_saturationhue.png')
    img = Image.fromarray(data_dict[str(70 + 3 * num_samples)]['image'])
    img.save('test_contrast.png')

    return data_dict

def data_prep(X_train, masks, rgb_masks, classes, train_ratio=0.9, mode='train', model='seq2seq'):
    X_train_n = X_train
    masks_n = masks
    rgb_masks_n = rgb_masks

    for i in range(X_train.shape[0]):
        cnt = 0
        for j in range(X_train.shape[1]):
            if X_train[i][j][0] == 1:
                X_train_n[i][cnt] = X_train[i][j]
                masks_n[i][cnt] = masks[i][j]
                rgb_masks_n[i][cnt] = rgb_masks[i][j]
                cnt = cnt + 1

        while cnt < X_train.shape[1]:
            X_train_n[i][cnt] = np.zeros(X_train[i][0].shape, dtype=np.float32)
            masks_n[i][cnt] = np.zeros(masks[i][0].shape, dtype=np.float32)
            rgb_masks_n[i][cnt] = np.zeros(rgb_masks[i][0].shape, dtype=np.float32)

            cnt = cnt + 1

    # X_train = X_train_n
    # masks = masks_n
    # rgb_masks = rgb_masks_n

    indices = np.where(X_train[:, :, 0] == 1)
    # print(len(indices[0]))
    # print(len(indices[1]))
    if model == 'partae':
        part_ae_masks = masks[indices]
        # signed_distance_field = skfmm.distance(part_ae_masks)
        part_ae_rgb_masks = rgb_masks[indices]
        classes = np.expand_dims(classes, axis=1)
        classes = np.tile(classes, (1, 24, 1))
        classes = classes[indices]

        # print(part_ae_masks.shape)
        # print(coords.shape)

        train_samples = int(part_ae_masks.shape[0] * train_ratio)

        if mode == 'train':
            part_ae_masks = part_ae_masks[:train_samples]
            part_ae_rgb_masks = part_ae_rgb_masks[:train_samples]
            classes = classes[:train_samples]
        else:
            part_ae_masks = part_ae_masks[train_samples:]
            part_ae_rgb_masks = part_ae_rgb_masks[train_samples:]
            classes = classes[train_samples:]
        return part_ae_masks, part_ae_rgb_masks, classes

    else:
        part_number = np.zeros(shape=(masks.shape[0], 24, 24))

        for i in range(X_train.shape[0]):
            cnt = 0
            for j in range(X_train.shape[1]):
                if X_train[i][j][0] == 1:
                    part_number[i][j][j] = 1

        sign = np.zeros(shape=(masks.shape[0], masks.shape[1], 1))
        for i in range(X_train.shape[0]):
            stop = 0
            for j in range(X_train.shape[1]):
                if X_train[i][j][0] == 1:
                    stop = j

            sign[i][stop][0] = 1

        bce_mask = X_train[:, :, 0]
        affine_input = X_train[:, :, 1:]
        affine_target = X_train[:, :, 1:]

        train_samples = int(bce_mask.shape[0] * train_ratio)
        if mode == 'train':
            masks = rgb_masks[:train_samples]
            part_number = part_number[:train_samples]
            bce_mask = bce_mask[:train_samples]
            affine_input = affine_input[:train_samples]
            affine_target = affine_target[:train_samples]
            sign = sign[:train_samples]
            classes = classes[:train_samples]
        else:
            masks = rgb_masks[train_samples:]
            part_number = part_number[train_samples:]
            bce_mask = bce_mask[train_samples:]
            affine_input = affine_input[train_samples:]
            affine_target = affine_target[train_samples:]
            sign = sign[train_samples:]
            classes = classes[train_samples:]

        data = {
            'vox2d': masks / 255.0,  # (B, T, vox_dim, vox_dim, 3),
            'sign': sign,
            'affine_input': affine_input,
            'affine_target': affine_target,
            'cond': part_number,
            'bce_mask': bce_mask,
            'classes': classes
        }
        print("Data dict prepared")

        return data

def get_data(class_name, mode, model):
    a_file = open('final_data/' + class_name + '_' + str(mode) + '_' + str(model) + '.pkl', "rb")
    dict_load = pickle.load(a_file)
    # print(dict_load['part_ae_masks'].shape)
    a_file.close()
    return dict_load

def get_data_full(type, iter, mode, model):
    a_file = open('full_data/' + type + '_' + str(iter) + '_' + str(mode) + '_' + str(model) + '.pkl', "rb")
    dict_load = pickle.load(a_file)
    print(dict_load['part_ae_masks'].shape)
    a_file.close()
    return dict_load

if __name__ == "__main__":

    modes = ['train', 'test']
    class_names = object_names
    X_train, masks, rgb_masks, rgb_images, classes = make_data(class_names)
    for mode in modes:
        model = 'partae'
        part_ae_masks, part_ae_rgb_masks, part_classes = data_prep(X_train, masks, rgb_masks,
                                                                   classes,
                                                                   mode=mode, model=model)
        samples = part_ae_masks.shape[0]
        per_part_size = 1000
        iterations = int(np.ceil(samples / per_part_size))
        for iter in range(iterations):
            dict = {}
            if iter == iterations - 1:
                dict['part_ae_masks'] = part_ae_masks[iter * per_part_size:]
                dict['part_ae_rgb_masks'] = part_ae_rgb_masks[iter * per_part_size:]
                dict['part_classes'] = part_classes[iter * per_part_size:]
            else:
                dict['part_ae_masks'] = part_ae_masks[iter * per_part_size:iter * per_part_size + per_part_size]
                dict['part_ae_rgb_masks'] = part_ae_rgb_masks[iter * per_part_size:iter * per_part_size + per_part_size]
                dict['part_classes'] = part_classes[iter * per_part_size:iter * per_part_size + per_part_size]

            a_file = open('full_data/all_' + str(iter) + '_' + str(mode) + '_' + str(model) + '.pkl', "wb")
            pickle.dump(dict, a_file)
            a_file.close()

        model = 'seq2seq'
        dict_seq = data_prep(X_train, masks, rgb_masks, classes, 200, mode=mode, model=model)
        for k in dict_seq.keys():
            print(dict_seq[k].shape)

        samples = dict_seq['vox2d'].shape[0]
        per_part_size = 1000
        iterations = int(np.ceil(samples / per_part_size))
        for iter in range(iterations):
            dict = {}
            for v in dict_seq.keys():
                if iter == iterations - 1:
                    dict[v] = dict_seq[v][iter * per_part_size:]
                else:
                    dict[v] = dict_seq[v][iter * per_part_size:iter * per_part_size + per_part_size]

            a_file = open('full_data/all_' + str(iter) + '_' + str(mode) + '_' + str(model) + '.pkl', "wb")
            pickle.dump(dict, a_file)
            a_file.close()

        #
        # a_file = open('final_data/' + obj + '_' + str(mode) + '_' + str(model) + '.pkl', "wb")
        # pickle.dump(dict_seq, a_file)
        # a_file.close()
        # a_file = open('final_data/' + obj + '_' + str(mode) + '_' + str(model) + '.pkl', "rb")
        # dict_seq = pickle.load(a_file)
        # print(dict_seq['vox2d'].shape)
        # a_file.close()


import os
import sys
import random
import math
import numpy as np
import skimage.io
import matplotlib
import matplotlib.pyplot as plt

# Root directory of the project
ROOT_DIR = os.path.abspath("../")
ROOT_DIR

# Import Mask RCNN
sys.path.append(ROOT_DIR)  # To find local version of the library
from mrcnn import utils
from mrcnn.config import Config
import mrcnn.model as modellib
from mrcnn import visualize, visualize_custom

# Directory to save logs and trained model
MODEL_DIR = os.path.join(ROOT_DIR, "logs")

# Local path to trained weights file

TUMOR_MODEL_PATH = os.path.join(ROOT_DIR, "maskrcnn_tumor.h5")
#/content/drive/My Drive/data/data/mask-Rcnn/maskrcnn_tumor.h5
# Directory of images to run detection on
# %mkdir ../images
IMAGE_DIR = os.path.join(ROOT_DIR, "images")

allowed_ext= ['.jpg', '.png']
path= sys.argv[1]
if path[-4:] not in allowed_ext:
	print('file extension allowed:{}, given:[{}]'.format(allowed_ext, path[-4:]))
	exit()

class Model_Config(Config):
    """Configuration for training on the brain tumor dataset.
    """
    # Give the configuration a recognizable name
    NAME = 'Tumor_detector'
    # Train on 1 GPU and 3 images per GPU. We can put multiple images on each
    #GPU_COUNT = 1
    IMAGES_PER_GPU = 1
    # Number of classes (including background)
    NUM_CLASSES = 1 + 1  # background + tumor
    IMAGE_MIN_DIM = 256
    IMAGE_MAX_DIM= 256

    # Use smaller anchors because our image and objects are small
    RPN_ANCHOR_SCALES = (8, 16, 32, 64, 128)  # anchor side in pixels

    DETECTION_MIN_CONFIDENCE = 0.70
    LEARNING_RATE = 0.001

config = Model_Config()
#config.display()

# Create model object in inference mode.
model = modellib.MaskRCNN(mode="inference", model_dir=MODEL_DIR, config=config)
# Load weights trained on MS-COCO
model.load_weights(TUMOR_MODEL_PATH, by_name=True)

# Load a random image from the `images` folder
image = skimage.io.imread(path)

# Run detection
results = model.detect([image], verbose=1)
class_names= ['-', 'tumor']
# Visualize results
r = results[0]
visualize.display_instances(image, r['rois'], r['masks'], r['class_ids'],
                           class_names , r['scores'])


from os import listdir
from os.path import isfile, join
import operator
import argparse
import pickle
import pyBigWig
import numpy as np

def parse_args():
    parser=argparse.ArgumentParser()
    parser.add_argument("--path_to_pickle")
    parser.add_argument("--chromsizes",default="/mnt/data/annotations/by_release/hg38/hg38.chrom.sizes")
    parser.add_argument("--output_dir")
    return parser.parse_args()

def main():
    args=parse_args()
    #prepare chromsizes for bigwig header use
    chromsizes=[i.split('\t') for i in open(args.chromsizes,'r').read().strip().split('\n')]
    chromsizes=[tuple([i[0],int(i[1])]) for i in chromsizes]
    chromsizes=sorted(chromsizes,key=operator.itemgetter(0))
    data=pickle.load(open(args.path_to_pickle,'rb'))
    cur_file=args.path_to_pickle.split('/')[-1]
    outf_label_prof=pyBigWig.open('/'.join([args.output_dir,cur_file+'.label_prof.bw']),'w')
    outf_label_prof.addHeader(chromsizes)
    outf_pred_prof=pyBigWig.open('/'.join([args.output_dir,cur_file+'.pred_prof.bw']),'w')
    outf_pred_prof.addHeader(chromsizes)
    outf_profile_shap=pyBigWig.open('/'.join([args.output_dir,cur_file+'.shap_profile.bw']),'w')
    outf_profile_shap.addHeader(chromsizes)
    outf_count_shap=pyBigWig.open('/'.join([args.output_dir,cur_file+'.shap_count.bw']),'w')
    outf_count_shap.addHeader(chromsizes)

    label_prof=data['label_prof']
    pred_prof=data['pred_prof']
    pred_sum=data['pred_sum']
    label_sum=data['label_sum']
    seq=data['seq']
    profile_shap=data['profile_shap']
    count_shap=data['count_shap']

    print("iterating keys")
    total_keys=len(label_prof)
    key_index=0
    for key in sorted(label_prof):
        print(str(key_index))
        key_index+=1
        chrom=key[0]
        summit=key[1]

        label_sum_val=np.exp(np.squeeze(label_sum[key])) # counts
        pred_sum_val=np.exp(np.squeeze(pred_sum[key])) # counts

        label_prof_val=label_prof[key]*label_sum_val
        pred_prof_val=pred_prof[key]*pred_sum_val

        seq_val=seq[key]

        if (len(profile_shap[key].shape)>1) and (profile_shap[key].shape[1]==4):
                profile_shap_val=np.sum(profile_shap[key]*seq_val,axis=1)
                count_shap_val=np.sum(count_shap[key]*seq_val,axis=1)
        else:
            profile_shap_val=profile_shap[key]
            count_shap_val=count_shap[key]
        n_entries_output=label_prof_val.shape[0]
        start_pos_output=summit-n_entries_output//2
        n_entries_interpretation=profile_shap_val.shape[0]
        start_pos_interpretation=summit-n_entries_interpretation//2
        added=False
        offset=0
        while not added:
            try:
                outf_label_prof.addEntries(chrom,start_pos_output+offset,values=label_prof_val[offset::],span=1,step=1)
                outf_pred_prof.addEntries(chrom,start_pos_output+offset,values=pred_prof_val[offset::],span=1,step=1)
                added=True
            except:
                offset+=1
                #print("offset:"+str(offset))
                #print(str(label_prof_val[offset::].shape))
                if offset > label_prof_val.shape[0]:
                    raise Exception()
        added=False
        offset=0
        while not added:
            try:
                outf_profile_shap.addEntries(chrom,start_pos_interpretation+offset,values=profile_shap_val[offset::],span=1,step=1)
                outf_count_shap.addEntries(chrom,start_pos_interpretation+offset,values=count_shap_val[offset::],span=1,step=1)
                added=True
            except:
                offset+=1
                if offset > profile_shap_val.shape[0]:
                    raise Exception()
    outf_label_prof.close()
    outf_pred_prof.close()
    outf_profile_shap.close()
    outf_count_shap.close()

if __name__=="__main__":
    main()


'''
Author: Jin
Goal: DQN
Following https://www.zhihu.com/people/fuckyou-74/posts
'''

import numpy as np
import tensorflow as tf
# np.random.seed(2020)

class DoubleDQN(object):
    def __init__(self, config):
        self.state_maxlength = config['STATE_MAXLENGTH']
        self.action_space = config['ACTION_SPACE']
        self.action_feature = None
        if 'ACTION_FEATURE' in config:
            self.action_feature = config['ACTION_FEATURE']
            self.action_dim = self.action_feature.shape[1]
        else:
            self.action_dim = config['ACTION_DIM'] # todo: pretrain action feature?
        self.rnn_state_dim = config['RNN_STATE_DIM']
        self.memory_size = config['MEMORY_SIZE']
        self.memory = np.zeros((self.memory_size, self.state_maxlength*4+4), dtype='object')
        self.gamma = config['GAMMA']
        self.lr = config['LEARNING_RATE']
        self.lr_decay_step = int(config['lr_decay_step'])
        self.epsilon_min = config['EPSILON']
        self.epsilon = 0.8
        self.epsilon_decay_step = int(config['epsilon_decay_step'])
        self.batch_size = config['BATCH_SIZE']
        self.learn_step_counter = 0
        self.replace_TargetNet_iter = config['REPLACE_TARGETNET']
        self.optimizer_name = config['OPTIMIZER']
        self.state_encoder = config["state_encoder"]
        self.double_q = True

        self.save_model_file = 'DoubleDQN'
        if 'SAVE_MODEL_FILE' in config:
            self.save_model_file = config['SAVE_MODEL_FILE']

        # tf.set_random_seed(2020)
        # init sess
        self._build_net()
        mainNet_params = tf.get_collection(tf.GraphKeys.GLOBAL_VARIABLES, scope='MainNet')
        TargetNet_params = tf.get_collection(tf.GraphKeys.GLOBAL_VARIABLES, scope='TargetNet')
        with tf.variable_scope('replacement'):
            self.target_replace_op = [tf.assign(t, e) for t, e in zip(TargetNet_params, mainNet_params)] # update TargetNet params by using MainNet params

        gpu_options = tf.GPUOptions(allow_growth=True)
        self.sess = tf.Session(config=tf.ConfigProto(gpu_options=gpu_options))
        self.sess.run(tf.global_variables_initializer())
        self.saver = tf.train.Saver()
        # self.writer = tf.summary.FileWriter("log/logs/", self.sess.graph)
        # exit(0)
        self.with_userinit = False
        print("DQN with user_init:", self.with_userinit)

    def _build_net(self):
        # ---- input ----
        self.s = tf.placeholder(tf.int32, [None, self.state_maxlength], name='s') # todo: return s as a representation, not do rnn O(n) everytime
        self.s_= tf.placeholder(tf.int32, [None, self.state_maxlength], name='s_')
        self.f = tf.placeholder(tf.int32, [None, self.state_maxlength], name='f')
        self.f_= tf.placeholder(tf.int32, [None, self.state_maxlength], name='f_')
        # self.r = tf.placeholder(tf.float32, [None, ], name='r')
        self.a = tf.placeholder(tf.int32, [None, ], name='a')
        self.q_target = tf.placeholder(tf.float32, [None, ], name='Q_next')
        self.len_s = tf.placeholder(tf.int32, [None, ], name='len_s')
        self.len_s_ = tf.placeholder(tf.int32, [None, ], name='len_s_')

        # ----- embedding -----
        if self.action_feature is not None:
            self.action_embeddings = tf.constant(dtype=tf.float32, value=self.action_feature)
        else:
            self.action_embeddings = tf.Variable(tf.random_normal(shape=[self.action_space, self.action_dim], mean=0.0, stddev=0.01), name='action_embeddings', dtype=tf.float32)
        # deal with the masked item:
        self.action_embeddings = tf.concat([self.action_embeddings, tf.constant(np.zeros([1, self.action_dim]), dtype='float32')], axis=0)
        self.feedback_embeddings = tf.Variable(tf.random_normal(shape=[2, self.action_dim], mean=0.0, stddev=0.01), name='feedback_embeddings', dtype=tf.float32)

        # --------- get the feedback embeddings
        self.input_f = tf.nn.embedding_lookup(self.feedback_embeddings, tf.reshape(self.f, [-1]))
        self.input_f_ = tf.nn.embedding_lookup(self.feedback_embeddings, tf.reshape(self.f_, [-1]))
        # -------- be careful: here merge the state (items) with feedback_emb
        self.input_s = tf.reshape(tf.nn.embedding_lookup(self.action_embeddings, tf.reshape(self.s, [-1])) * self.input_f, [-1, self.state_maxlength, self.action_dim]) #(None, d) -> (None, s_lengh, d)
        self.input_s_ = tf.reshape(tf.nn.embedding_lookup(self.action_embeddings, tf.reshape(self.s_, [-1])) * self.input_f_, [-1, self.state_maxlength, self.action_dim]) #(None, d) -> (None, s_lengh, d)

        # # ********* here we use rnn, it can be MLP, CNN ....
        if self.state_encoder.lower() == 'gru':
            w_init, b_init = tf.random_normal_initializer(), tf.constant_initializer()
            # ----- build MainNet -----
            with tf.variable_scope('MainNet'):
                cell_main = tf.contrib.rnn.GRUCell(num_units=self.rnn_state_dim)
                _, h_s = tf.nn.dynamic_rnn(cell_main, dtype=tf.float32, sequence_length=self.len_s, inputs=self.input_s)
                self.q_eval = tf.layers.dense(h_s, self.action_space, kernel_initializer=w_init, bias_initializer=b_init, name='q')
            # ----- build TargetNet -----
            with tf.variable_scope('TargetNet'):
                cell_target = tf.contrib.rnn.GRUCell(num_units=self.rnn_state_dim)
                _, h_s_ = tf.nn.dynamic_rnn(cell_target, dtype=tf.float32, sequence_length=self.len_s_, inputs=self.input_s_)
                self.q_next = tf.layers.dense(h_s_, self.action_space, kernel_initializer=w_init, bias_initializer=b_init, name='t2')

        # ********* here we use 1-layer dense
        elif self.state_encoder.lower() == 'mlp':
            w_init, b_init = tf.random_normal_initializer(), tf.constant_initializer()
            ## before the code is h_s = tf.reduce_mean(self.input_s, axis=1)
            h_s = tf.reduce_sum(self.input_s, axis=1) / tf.reshape(tf.cast(self.len_s, dtype=tf.float32), [-1, 1]) # (None, d)
            h_s_ = tf.reduce_sum(self.input_s_, axis=1) / tf.reshape(tf.cast(self.len_s_, dtype=tf.float32), [-1, 1]) # (None, d)
            # ----- build MainNet -----
            with tf.variable_scope('MainNet'):
                self.q_eval = tf.layers.dense(h_s, self.action_space, kernel_initializer=w_init, bias_initializer=b_init, name='q_eval')
            # ----- build TargetNet -----
            with tf.variable_scope('TargetNet'):
                self.q_next = tf.layers.dense(h_s_, self.action_space, kernel_initializer=w_init, bias_initializer=b_init, name='q_next')
        else:
            print("error state_encoder")
            exit(1)

        # ----- DQN-loss -----
        with tf.variable_scope('q_eval'):
            a_indices = tf.stack([tf.range(tf.shape(self.a)[0],dtype=tf.int32), self.a], axis=1)
            self.q_eval_s_a = tf.gather_nd(self.q_eval, a_indices)
        # with tf.variable_scope('q_next'):
        #     self.q_target = self.r + self.gamma * tf.reduce_max(self.q_next, axis=1, name='Qmax_s_')
        #     self.q_target = tf.stop_gradient(self.q_target)

        with tf.variable_scope('loss'):
            # self.loss = tf.reduce_mean(tf.squared_difference(self.q_target, self.q_eval_s_a))
            # smooth-l1-loss: https://blog.csdn.net/u014365862/article/details/79924201
            diff = self.q_target - self.q_eval_s_a
            abs_diff = tf.abs(diff)
            smoothL1_sign = tf.stop_gradient(tf.to_float(tf.less(abs_diff, 1.)))
            loss = tf.pow(diff, 2) * 0.5 * smoothL1_sign + (abs_diff - 0.5) * (1 - smoothL1_sign)
            self.loss = tf.reduce_mean(loss)

        with tf.variable_scope('optimize'):
            self.global_step = tf.Variable(0, trainable=False)
            self.decayed_lr = tf.train.exponential_decay(self.lr, \
                                        self.global_step, 1, 0.9, staircase=True)
            ## self.decayed_lr = self.lr
            if self.optimizer_name.lower() == 'adam':
                self.optimizer = tf.train.AdamOptimizer(self.decayed_lr).minimize(self.loss)
            elif self.optimizer_name.lower() == 'rmsprop':
                self.optimizer = tf.train.RMSPropOptimizer(learning_rate=self.decayed_lr, decay=0.99999).minimize(self.loss)
            else:
                print("Wrong optimizer. Stopped")
                exit(1)
            self.add_global = self.global_step.assign_add(1)

    # def _pad_(self, state, max_length, value):
    def _pad_(self, state, value):
        return state + [value] * (self.state_maxlength - len(state))
    # store (s, a, r, s') in Memory
    def store_memory(self, s, a, r, s_): # s = [[item], [feedback]]
        len_s, len_s_ = len(s[0]), len(s_[0])
        if not hasattr(self, 'memory_counter'): # first (s, a, r, s')
            self.memory_counter = 0
        # replace the first input (first in first out)
        index = self.memory_counter % self.memory_size
        state = self._pad_(s[0], self.action_space)
        state_ = self._pad_(s_[0], self.action_space)
        f = self._pad_(s[1], 0) # feedback is 0 or 1 now, it is ok to set 0 or 1, because refering s is [0, 0, 0] so the sum should be 0 too.
        f_ = self._pad_(s_[1], 0)
        self.memory[index, :] = state + state_ + f + f_ + [a, r, len_s, len_s_]
        self.memory_counter += 1
    def store_memorys(self, states, actions, rewards, states_, len_s, len_s_): # s = [[item], [feedback]]
        if (len_s == 0) and (self.with_userinit == False):
            return
        u_batch_size = states.shape[0]
        assert (len_s + 1) == len_s_
        if not hasattr(self, 'memory_counter'): # first (s, a, r, s')
            self.memory_counter = 0
        # replace the first input (first in first out)
        index = self.memory_counter % self.memory_size
        s = states[:, 0]
        f = states[:, 1]
        s_ = states_[:, 0]
        f_ = states_[:, 1]
        r = np.expand_dims(rewards, axis=-1)
        a = np.expand_dims(actions, axis=-1)
        len_s = np.full((u_batch_size, 1), len_s, dtype=np.int32)
        len_s_ = np.full((u_batch_size, 1), len_s_, dtype=np.int32)
        if (index+u_batch_size) > self.memory_size:
            part2_len = index+u_batch_size - self.memory_size
            indices = list(range(index, self.memory_size)) + list(range(part2_len))
            assert len(indices) == u_batch_size
            self.memory[indices, :] = np.hstack((s, s_, f, f_, a, r, len_s, len_s_))
        else:
            self.memory[index:(index+u_batch_size), :] = np.hstack((s, s_, f, f_, a, r, len_s, len_s_))
        self.memory_counter += u_batch_size

    def choose_action(self, s, greedy='false'):
        if ((np.random.uniform() < self.epsilon) and (greedy == 'false')) or (len(s[0]) == 0):
            action = np.random.randint(0, self.action_space)
            while action in s[0]:
                action = np.random.randint(0, self.action_space)
        else:
            state = np.array(self._pad_(s[0], self.action_space), dtype=np.int32)
            f = np.array(self._pad_(s[1], 0), dtype=np.int32)
            action_value = self.sess.run(self.q_eval, feed_dict={self.f:f[None, :], self.s:state[None, :], self.len_s:np.array([len(s[0])], dtype=np.int32)})
            action_value = action_value[0, :]
            assert len(action_value) == self.action_space
            action = np.argmax(action_value)
            change_v = np.min(action_value) - 1.0
            while action in s[0]:
                action_value[action] = change_v
                action = np.argmax(action_value)
        return action
    def choose_actions(self, states, len_s, greedy='false'):
        if not hasattr(self, 'interact_count'):
            self.interact_count = 0
        # self.epsilon = max(self.epsilon_min, self.epsilon - self.epsilon_delta)
        self.interact_count += 1
        # if greedy is false, then do \epsilon-greedy. or just use policy to max
        s = states[:, 0]
        f = states[:, 1]
        batch_size = s.shape[0]
        ## indices for items can not be choosed
        indices_x = np.repeat(range(batch_size), len_s).astype('int')
        indices_y = s[:, :len_s].flatten()

        select_p = np.random.uniform(size=(batch_size, self.action_space))
        select_p[indices_x, indices_y] = -1.
        actions_random = np.argmax(select_p, axis=1)

        # # decrease the epsilon for exploration
        if self.interact_count % self.epsilon_decay_step == 0:
            self.epsilon = max(self.epsilon_min, self.epsilon-0.1)

        if (not self.with_userinit) and (len_s == 0):
            return actions_random

        ### use policy
        # --- normal way ----
        action_values = self.sess.run(self.q_eval, feed_dict={self.f:f, self.s:s, self.len_s:np.full((batch_size), len_s, dtype=np.int32)})

        # set the historical action very small value, deal with repeat items
        replace_value = np.min(action_values) - 1.
        action_values[indices_x, indices_y] = replace_value
        actions = np.argmax(action_values, axis=1)
        if greedy.lower() == 'true':
            return actions

        return np.where(np.random.uniform(size=(s.shape[0],))<self.epsilon, actions_random, actions)

    def rerank(self, s, itemList):
        if len(s[0]) == 0:
            print("to get scores, the state must has something")
            exit(1)
        state = np.array(self._pad_(s[0], self.action_space), dtype=np.int32)
        f = np.array(self._pad_(s[1], 0), dtype=np.int32)
        action_value = self.sess.run(self.q_eval, feed_dict={self.f:f[None, :], self.s:state[None, :], self.len_s:np.array([len(s[0])], dtype=np.int32)})
        action_value = action_value[0, :]
        assert len(action_value) == self.action_space
        scores = action_value[itemList]
        rerankOrder = np.flip(np.argsort(scores))
        return rerankOrder
    def reranks(self, states, len_s, itemLists=None):
        if (self.with_userinit == False) and (len_s == 0):
            print("to get scores, the state must has something")
            exit(1)
        s = states[:, 0]
        f = states[:, 1]
        action_values = self.sess.run(self.q_eval, feed_dict={self.f:f, self.s:s, self.len_s:np.full((s.shape[0]), len_s, dtype=np.int32)})
        assert action_values.shape[1] == self.action_space
        if itemLists == None:
            scores = action_values
            return np.flip(np.argsort(scores)), scores
        else:
            rerankOrder, scores = [], []
            for i in range(itemLists.shape[0]):
                score = action_values[i][itemLists[i]]
                scores.append(score)
                rerankOrder.append(np.flip(np.argsort(score)))
            return rerankOrder, scores

    def learn(self):
        if self.memory_counter < self.batch_size:
            return
        # get batch memory
        if self.memory_counter > self.memory_size:
            sample_index = np.random.choice(self.memory_size, size=self.batch_size)
        else:
            sample_index = np.random.choice(self.memory_counter, size=self.batch_size)
        # print(self.memory_counter, sample_index)
        batch_memory = self.memory[sample_index, :]
        # deal with states with unfixed length
        # states = tf.keras.preprocessing.sequence.pad_sequences(batch_memory[:, 0], \
        #             maxlen=self.state_maxlength, padding='post', value=self.action_space)
        # states_ = tf.keras.preprocessing.sequence.pad_sequences(batch_memory[:, 3], \
        #             maxlen=self.state_maxlength, padding='post', value=self.action_space)
        states = np.array(batch_memory[:, :self.state_maxlength], dtype='int')
        states_ = np.array(batch_memory[:, self.state_maxlength:(self.state_maxlength*2)], dtype='int')
        f = np.array(batch_memory[:, (self.state_maxlength*2):(self.state_maxlength*3)], dtype='int')
        f_ = np.array(batch_memory[:, (self.state_maxlength*3):(self.state_maxlength*4)], dtype='int')
        # learn
        actions = batch_memory[:, -4]
        rewards = batch_memory[:, -3]
        len_s, len_s_ = batch_memory[:, -2], batch_memory[:, -1]
        # print("s, s_:\n", states[:5], '\n', states_[:5], '\n')
        # print("actions, rewards:\n", actions[:5], '\n', rewards[:5], '\n')
        # print("len_s, len_s_", len_s[:5], len_s_[:5])
        q_next, q_eval_next = self.sess.run([self.q_next, self.q_eval], \
                feed_dict={self.s_:states_, self.f_:f_, self.s:states_, self.f:f_, \
                self.len_s:len_s_, self.len_s_:len_s_})

        batch_index = np.arange(self.batch_size, dtype=np.int32)
        if self.double_q:
            max_act4_next=np.argmax(q_eval_next, axis=1)
            selected_q_next = q_next[batch_index, max_act4_next]
        else:
            selected_q_next = np.max(q_next, axis=1)

        q_target = rewards + self.gamma * selected_q_next

        _, loss = self.sess.run([self.optimizer, self.loss], \
                feed_dict={self.s:states, self.f:f, self.a:actions, self.len_s:len_s,self.q_target:q_target})
        # todo: can \epislon can be decreased

        self.learn_step_counter += 1
        # check if change TargetNet params or not
        if self.learn_step_counter % self.replace_TargetNet_iter == 0:
            self.sess.run(self.target_replace_op)
        if self.learn_step_counter % self.lr_decay_step == 0:
            self.lr_decay()
        return loss

    def lr_decay(self):
        _, lr = self.sess.run([self.add_global, self.decayed_lr])
        return lr
    def get_lr(self):
        return self.sess.run(self.decayed_lr)

    def load_pretrain_model(self):
        self.saver.restore(self.sess, './checkpoint_dir/' + self.save_model_file)
    def save_model(self):
        self.saver.save(self.sess, './checkpoint_dir/' + self.save_model_file)


import tensorflow as tf
import numpy as np
import joblib
from rl.acktr.utils import Scheduler, find_trainable_variables
from rl.acktr.utils import fc, mse
from rl.acktr import kfac

disc_types = ['decentralized', 'centralized', 'single']

class Discriminator(object):
    def __init__(self, sess, ob_spaces, ac_spaces,
                 nstack, index, disc_type='decentralized', hidden_size=128,
                 lr_rate=0.01, total_steps=50000, scope="discriminator", kfac_clip=0.001, max_grad_norm=0.5):
        self.lr = Scheduler(v=lr_rate, nvalues=total_steps, schedule='linear')
        self.disc_type = disc_type
        if disc_type not in disc_types:
            assert False
        self.scope = scope
        self.index = index
        self.sess = sess
        ob_space = ob_spaces[index]
        ac_space = ac_spaces[index]
        self.ob_shape = ob_space.shape[0] * nstack
        try:
            nacts = [ac.n for ac in ac_spaces]
            nact = ac_space.n
        except:
            nacts = [ac.shape[0] for ac in ac_spaces]
            nact = ac_space.shape[0]

        self.ac_shape = nact * nstack
        self.all_ob_shape = sum([obs.shape[0] for obs in ob_spaces]) * nstack
        self.all_ac_shape = sum(nacts) * nstack
        self.hidden_size = hidden_size

        if disc_type == 'decentralized':
            input_shape = self.ob_shape + self.ac_shape
        elif disc_type == 'centralized':
            input_shape = self.all_ob_shape + self.all_ac_shape
        elif disc_type == 'single':
            input_shape = self.all_ob_shape + self.all_ac_shape
        else:
            assert False

        self.g = tf.placeholder(tf.float32, (None, input_shape))
        self.e = tf.placeholder(tf.float32, (None, input_shape))
        self.lr_rate = tf.placeholder(tf.float32, ())
        self.adv = tf.placeholder(tf.float32, ())

        num_outputs = len(ob_spaces) if disc_type == 'centralized' else 1

        logits = self.build_graph(tf.concat([self.g, self.e], axis=0), num_outputs, reuse=False)
        labels = tf.concat([tf.ones([tf.shape(self.g)[0], 1]), -tf.ones([tf.shape(self.e)[0], 1])], axis=0)

        g_logits = self.build_graph(self.g, num_outputs, reuse=True)
        e_logits = self.build_graph(self.e, num_outputs, reuse=True)

        self.g_loss = tf.reduce_mean(g_logits)
        self.e_loss = tf.reduce_mean(-e_logits)

        # self.g_loss = tf.reduce_mean(tf.nn.sigmoid_cross_entropy_with_logits(
        #     logits=g_logits, labels=tf.zeros_like(g_logits)))
        # self.e_loss = tf.reduce_mean(tf.nn.sigmoid_cross_entropy_with_logits(
        #     logits=e_logits, labels=tf.ones_like(e_logits)))

        self.total_loss = logits * labels # tf.reduce_mean(tf.nn.sigmoid_cross_entropy_with_logits(logits=logits, labels=labels))

        sample_net = logits + tf.random_normal(tf.shape(logits))
        fisher_loss = -tf.reduce_mean(tf.pow(logits - tf.stop_gradient(sample_net), 2))

        self.reward_op = tf.sigmoid(g_logits)
        # self.reward_op = tf.nn.sigmoid_cross_entropy_with_logits(logits=g_logits, labels=tf.zeros_like(g_logits))

        self.var_list = self.get_trainable_variables()
        params = find_trainable_variables(self.scope)
        grads = tf.gradients(self.total_loss, params)

        # self.d_optim = tf.train.AdamOptimizer(self.lr_rate, beta1=0.5, beta2=0.9).minimize(self.total_loss, var_list=self.var_list)
        with tf.variable_scope(self.scope + '/d_optim'):
            d_optim = kfac.KfacOptimizer(
                learning_rate=self.lr_rate, clip_kl=kfac_clip,
                momentum=0.9, kfac_update=1, epsilon=0.01,
                stats_decay=0.99, async=0, cold_iter=10,
                max_grad_norm=max_grad_norm)
            update_stats_op = d_optim.compute_and_apply_stats(fisher_loss, var_list=params)
            train_op, q_runner = d_optim.apply_gradients(list(zip(grads, params)))
            self.q_runner = q_runner
        self.d_optim = train_op
        self.saver = tf.train.Saver(self.get_variables())

        self.params_flat = self.get_trainable_variables()
        self.clip = [tf.assign(v, tf.clip_by_value(v, -0.05, 0.05)) for v in self.get_trainable_variables()]
        self.clip = tf.group(*self.clip)

    def build_graph(self, x, num_outputs=1, reuse=False):
        with tf.variable_scope(self.scope):
            if reuse:
                tf.get_variable_scope().reuse_variables()
            p_h1 = fc(x, 'fc1', nh=self.hidden_size)
            # p_h2 = fc(p_h1, 'fc2', nh=self.hidden_size)
            logits = fc(p_h1, 'out', nh=num_outputs, act=lambda x: x)
        return logits

    def get_variables(self):
        return tf.get_collection(tf.GraphKeys.GLOBAL_VARIABLES, self.scope)

    def get_trainable_variables(self):
        return tf.get_collection(tf.GraphKeys.TRAINABLE_VARIABLES, self.scope)

    def get_reward(self, obs, acs):
        if len(obs.shape) == 1:
            obs = np.expand_dims(obs, 0)
        if len(acs.shape) == 1:
            acs = np.expand_dims(acs, 0)
        feed_dict = {self.g: np.concatenate([obs, acs], axis=1)}
        return self.sess.run(self.reward_op, feed_dict)

    def train(self, g_obs, g_acs, e_obs, e_acs):
        feed_dict = {self.g: np.concatenate([g_obs, g_acs], axis=1),
                     self.e: np.concatenate([e_obs, e_acs], axis=1), self.lr_rate: self.lr.value()}
        loss, _ = self.sess.run([self.total_loss, self.d_optim], feed_dict)
        self.sess.run(self.clip)
        g_loss, e_loss = self.sess.run([self.g_loss, self.e_loss], feed_dict)
        return g_loss, e_loss, None, None

    def restore(self, path):
        print('restoring from:' + path)
        self.saver.restore(self.sess, path)

    def save(self, save_path):
        ps = self.sess.run(self.params_flat)
        joblib.dump(ps, save_path)

    def load(self, load_path):
        loaded_params = joblib.load(load_path)
        restores = []
        for p, loaded_p in zip(self.params_flat, loaded_params):
            restores.append(p.assign(loaded_p))
        self.sess.run(restores)


import logging
import os

import astropy.units as u
import h5py
import matplotlib
import matplotlib.pyplot as plt
import numpy as np
import seaborn as sns
from astropy.io.misc.hdf5 import read_table_hdf5

from exorad.log import Logger, setLogLevel

sns.set()

mpl_logger = logging.getLogger('matplotlib')
mpl_logger.setLevel(logging.WARNING)

class Plotter(Logger):
    """
        Plotter class. It offers a fast and easy way to produce diagnostic plots on the produced data.

        Parameters
        -----------
        input_table: Qtable
            table where to grab data and wl grid to plot.
        channels: dict
            dictionary describing the channels in the payload. Default is None
    """

    def __init__(self, input_table, channels=None):
        self.set_log_name()
        self.inputTable = input_table
        self.channels = channels
        self.fig = None
        self.fig_efficiency = None

    def plot_bands(self, ax, scale='log', channel_edges=True):
        """
        It plots the channels bands behind the indicated axes.

        Parameters
        -----------
        ax: matplotlib.axes
            axes where to plot the bands

        Returns
        --------
        matplotlib.axes.axes

        Note
        ----
        The Class input_table input parameter is required for this method to work.
        """
        channels = set(self.inputTable['chName'])
        palette = sns.color_palette('colorblind')
        tick_list = []

        for k, channelName in enumerate(channels):
            wl_min = min(self.inputTable['LeftBinEdge'][np.where(self.inputTable['chName'] == channelName)]).value
            wl_max = max(self.inputTable['RightBinEdge'][np.where(self.inputTable['chName'] == channelName)]).value
            ax.axvspan(wl_min, wl_max, alpha=0.1, zorder=0, color=palette[k])
            tick_list.append(wl_min)
        tick_list.append(wl_max)
        ax.set_xscale(scale)
        if channel_edges:
            ax.set_xticks(tick_list)
            ax.get_xaxis().set_major_formatter(matplotlib.ticker.ScalarFormatter())

        return ax

    def plot_efficiency(self, scale='log', channel_edges=True):
        """
        It produces a figure with payload efficiency over wavelength.
        The quantities reported are quantum efficiency, transmission and the photon conversion efficiency (pce)
        computed as the product of the quantum efficiency and transmission.

        Returns
        --------
        matplotlib.pyplot.figure

        Note
        ----
        The Class channels input parameter is required for this method to work.
        """
        if self.channels:
            from matplotlib.lines import Line2D
            from exorad.models.signal import Signal

            palette = sns.color_palette('bright')

            fig, ax = plt.subplots(1, 1, figsize=(10, 8))
            fig.suptitle('Payload photon conversion efficiency')

            if self.channels:
                keys = ['transmission', 'qe']
                for ch in self.channels:
                    pce = None
                    for i, key in enumerate(keys):
                        data = self.channels[ch].built_instr['{}_data'.format(key)]
                        sig = Signal(wl_grid=data['wl_grid']['value'] * u.Unit(data['wl_grid']['unit']),
                                     data=data['data']['value'])
                        ax.plot(sig.wl_grid, sig.data, color=palette[i], zorder=10)
                        if not pce:
                            pce = sig
                        else:
                            sig.spectral_rebin(pce.wl_grid)
                            pce.data *= sig.data
                    ax.plot(pce.wl_grid, pce.data, color=palette[i + 1], zorder=10)

                lines, labels = [], []
                for i, key in enumerate(keys):
                    lines.append(Line2D([0], [0], color=palette[i], lw=4))
                    labels.append(key)
                lines.append(Line2D([0], [0], color=palette[i + 1], lw=4))
                labels.append('pce')
            else:
                pce = self.inputTable['TR'] * self.inputTable['QE']
                self.inputTable['pce'] = pce
                keys = ['TR', 'QE', 'pce']
                self.debug('efficiency keys : {}'.format(keys))
                for e in keys:
                    ax.plot(self.inputTable['Wavelength'], self.inputTable[e], label=e, zorder=10)
                    # ax.plot(self.inputTable['Wavelength'], self.inputTable[e], c='None')
            ax.grid(zorder=0)
            ax = self.plot_bands(ax, scale, channel_edges)
            ax.legend(handles=lines, labels=labels)
            ax.set_title('Photon conversion efficiency')
            ax.set_xlabel('Wavelength [$\mu m$]')
            ax.set_ylabel('efficiency')
            # ax.set_xscale('log')

            self.fig_efficiency = fig
            return fig

        else:
            self.error('channels parameter is required for this method to work')

    def plot_noise(self, ax, ylim=None, scale='log', channel_edges=True):
        """
        It plots the noise components found in the input table in the indicated axes.

        Parameters
        -----------
        ax: matplotlib.axes
            axes where to plot the noises
        ylim: float
            if present, it sets the axes y bottom lim. Default is None.

        Returns
        --------
        matplotlib.axes.axes

        Note
        ----
        The Class input_table input parameter is required for this method to work.
        """
        palette = sns.color_palette('bright')

        noise_keys = [x for x in self.inputTable.keys() if 'noise' in x or 'custom' in x]
        self.debug('noise keys : {}'.format(noise_keys))
        for k, n in enumerate(noise_keys):
            if n == 'total_noise':
                ax.plot(self.inputTable['Wavelength'], self.inputTable[n], zorder=9, lw=1, c='k',
                        marker='.', markersize=5, label='total_noise', alpha=0.8)  # , c='None')
            else:
                if self.inputTable[n].unit == u.hr ** 0.5:
                    noise = self.inputTable[n]
                else:
                    self.debug('{} rescaled by starSignal_inAperture'.format(n))
                    noise = self.inputTable[n] / self.inputTable['star_signal_inAperture']
                # ax.scatter(self.inputTable['Wavelength'], noise, label=n, zorder=10, s=5, color=palette[k])
                ax.plot(self.inputTable['Wavelength'], noise, zorder=9, lw=1, alpha=0.5, marker='.',
                        label=n)  # color=palette[k])  # c='None')

        if not ylim:
            ax.set_ylim(1e-7)
        else:
            ax.set_ylim(ylim)
        #        ax.grid(zorder=0, which='both')
        locmaj = matplotlib.ticker.LogLocator(base=10, numticks=12)
        ax.yaxis.set_major_locator(locmaj)
        locmin = matplotlib.ticker.LogLocator(base=10.0, subs=(0.2, 0.4, 0.6, 0.8), numticks=12)
        ax.yaxis.set_minor_locator(locmin)
        ax.yaxis.set_minor_formatter(matplotlib.ticker.NullFormatter())
        ax.grid(axis='y', which='minor', alpha=0.3)
        ax.grid(axis='y', which='major', alpha=0.5)
        #        ax.legend(bbox_to_anchor=(1, 1))
        ax.legend(prop={'size': 12}, loc='upper left',
                  ncol=3, bbox_to_anchor=(0.05, -0.2), labelspacing=1.2, handlelength=1)
        ax.set_title('Noise Budget')
        ax.set_xlabel('Wavelength [$\mu m$]')
        ax.set_ylabel('relative noise [$\sqrt{{hr}}$]')
        ax.set_yscale('log')
        # ax.set_xscale('log')
        ax = self.plot_bands(ax, scale, channel_edges)
        return ax

    def plot_signal(self, ax, ylim=None, scale='log', channel_edges=True):
        """
        It plots the signal components found in the input table in the indicated axes.

        Parameters
        -----------
        ax: matplotlib.axes
            axes where to plot the signals
        ylim: float
            if present, it sets the axes y bottom lim. Default is None.

        Returns
        --------
        matplotlib.axes.axes

        Note
        ----
        The Class input_table input parameter is required for this method to work.
        """
        palette = sns.color_palette('bright')

        keys = [x for x in self.inputTable.keys() if 'signal' in x and 'noise' not in x]
        self.debug('signal keys : {}'.format(keys))
        for k, s in enumerate(keys):
            # ax.scatter(self.inputTable['Wavelength'], self.inputTable[s], label=s, zorder=10, s=5, color=palette[k])
            ax.plot(self.inputTable['Wavelength'], self.inputTable[s], zorder=9, lw=1, alpha=0.5, marker='.', label=s)
            # color=palette[k])  # , c='None')

        if not ylim:
            ax.set_ylim(1e-3)
        else:
            ax.set_ylim(ylim)
        #        ax.grid(zorder=0, which='both')
        locmaj = matplotlib.ticker.LogLocator(base=10, numticks=12)
        ax.yaxis.set_major_locator(locmaj)
        locmin = matplotlib.ticker.LogLocator(base=10.0, subs=(0.2, 0.4, 0.6, 0.8), numticks=12)
        ax.yaxis.set_minor_locator(locmin)
        ax.yaxis.set_minor_formatter(matplotlib.ticker.NullFormatter())
        ax.grid(axis='y', which='minor', alpha=0.3)
        ax.grid(axis='y', which='major', alpha=0.5)
        #        ax.legend(bbox_to_anchor=(1, 1))
        ax.legend(prop={'size': 12}, loc='upper left',
                  ncol=3, bbox_to_anchor=(0.05, -0.2), labelspacing=1.2, handlelength=1)
        ax.set_title('Signals')
        ax.set_xlabel('Wavelength [$\mu m$]')
        ax.set_ylabel('$ct/s$')
        ax.set_yscale('log')
        # ax.set_xscale('log')
        ax = self.plot_bands(ax, scale, channel_edges)

        return ax

    def plot_table(self, sig_ylim=None, noise_ylim=None, scale='log', channel_edges=True):
        """
        It produces a figure with signal and noise for the input table.

        Returns
        --------
        matplotlib.pyplot.figure
        (matplotlib.axes.axes, matplotlib.axes.axes)

        Note
        ----
        The Class input_table input parameter is required for this method to work.
        """
        fig, (ax1, ax2) = plt.subplots(2, 1, figsize=(10, 10))
        fig.suptitle(self.inputTable.meta['name'])
        ax1 = self.plot_signal(ax1, ylim=sig_ylim, scale=scale, channel_edges=channel_edges)
        ax2 = self.plot_noise(ax2, ylim=noise_ylim, scale=scale, channel_edges=channel_edges)
        plt.tight_layout()
        plt.subplots_adjust(top=0.9, bottom=0.22, hspace=0.7)
        self.fig = fig
        return fig, (ax1, ax2)

    def save_fig(self, name, efficiency=False):
        """
        It saves the produced figure.

        Parameters
        --------
        name: str
            figure name
        efficiency: bool
            if True it wll save the efficiency plot instead of the table plot. Default is False.
        """
        try:
            if efficiency:
                self.fig_efficiency.savefig('{}'.format(name))
            else:
                self.fig.savefig('{}'.format(name))

            self.info('plot saved in {}'.format(name))
        except AttributeError:
            self.error("the indicated figure is not available. Check if you have produced it.")

def main():
    import argparse
    from exorad.__version__ import __version__
    from exorad.utils.util import parse_range

    parser = argparse.ArgumentParser(description='ExoRad {}'.format(__version__))
    parser.add_argument("-i", "--input", dest='input', type=str,
                        required=True, help="Input h5 file to pass")
    parser.add_argument("-o", "--out", dest='out', type=str, default='None',
                        required=True, help="Output directory")
    parser.add_argument("-n", "--target-number", dest='target_number', type=str, default='all',
                        required=False, help="A list or range of targets to run")
    parser.add_argument("-t", "--target-name", dest='target_name', type=str, default='None',
                        required=False, help="name of the target to plot")
    parser.add_argument("-d", "--debug", dest='debug', default=False,
                        required=False, help="log output on screen", action='store_true')

    args = parser.parse_args()

    logger = logging.getLogger('exorad')
    from exorad.utils.ascii_art import ascii_plot
    logger.info(ascii_plot)
    logger.info('code version {}'.format(__version__))

    if args.debug: setLogLevel(logging.DEBUG)

    if not os.path.exists(args.out):
        os.makedirs(args.out)
        logger.info('output directory created')

    logger.info('reading {}'.format(args.input))
    file = h5py.File(args.input)

    if args.target_number != 'all' and args.target_name != 'None':
        logger.error('you cannot use both target number and target name')
        raise ValueError

    targets_dir = file['targets']
    targets_to_run_id = parse_range(args.target_number, len(targets_dir.keys()))
    targets_to_run = [list(targets_dir.keys())[n] for n in targets_to_run_id]

    if args.target_name != 'None':
        targets_to_run = [target for target in targets_to_run if target == args.target_name]

    for target in targets_to_run:
        target_dir = targets_dir[target]
        table_dir = target_dir['table']
        table = read_table_hdf5(table_dir, path='table')

        plotter = Plotter(input_table=table)
        plotter.plot_table()
        plotter.save_fig(os.path.join(args.out, '{}.png'.format(target)))
        plt.close()

import numpy,import astropy,from astropy