#!/usr/bin/env python3

# ============================ Information =============================
# Project:      Texas A&M AutoDrive Challenge - Year 2
# Language:     Python 3.5.2
# ROS Package:  camera_detection
# Repository:   https://github.tamu.edu/kipvasq9/camera_detection
# File Name:    camera_detection.py
# Version:      1.0.0
# Description:  Provide access to Neural Networks through ROS.
# Date:         September 24, 2018
# Author:       Juan Vasquez and Robert Hodge
# Contact:      kipvasq9@tamu.edu

# ============================== Imports ===============================
import cv2
import signal
import sys
import numpy as np
import matplotlib.pyplot as plt
import os
import time
import PIL
from PIL import ImageFile
from PIL import Image as PILImage
from std_msgs.msg import Float32MultiArray
ImageFile.LOAD_TRUNCATED_IMAGES = True

# neural network
import tensorflow as tf
import keras
from keras_retinanet import models
from keras_retinanet.utils.image import read_image_bgr, preprocess_image, resize_image
from keras_retinanet.utils.visualization import draw_box, draw_caption
from keras_retinanet.utils.colors import label_color
from path_planner.msg import TrafficLightArray, TrafficLight, TrafficSignArray, TrafficSign, ObstacleArray, Obstacle

# ros
import rospy
from cv_bridge import CvBridge, CvBridgeError
from sensor_msgs.msg import Image

# define tensorflow session
def get_session():
    config = tf.ConfigProto()
    config.gpu_options.allow_growth = True
    return tf.Session(config=config)

# use this environment flag to change which GPU to use
#os.environ["CUDA_VISIBLE_DEVICES"] = "1"

# set the modified tf session as backend in keras
keras.backend.tensorflow_backend.set_session(get_session())

# adjust this to point to your downloaded/trained model
# models can be downloaded here: https://github.com/fizyr/keras-retinanet/releases
model_path = "/home/autodrive/Desktop/catkin_ws/src/keras-retinanet/snapshots/juan.h5"
model_path = "/home/autodrive/Desktop/catkin_ws/src/keras-retinanet/snapshots/juan.h5"
model = models.load_model(model_path, backbone_name='resnet50')
model._make_predict_function()

# Camera Detection ROS Node Class
class CameraDetectionROSNode:
    # Initialize class
    def __init__(self):
        self.image = None   # image variable
        self.active = True  # should we be processing images
        self.obstacle_publisher = rospy.Publisher('/obstacle', ObstacleArray, queue_size = 1)
        # CTRL+C signal handler
        signal.signal(signal.SIGINT, self.signalInterruptHandler)

        # initialize ros node
        rospy.init_node('camera_detection_node', anonymous=True)

        # set camera image topic subscriber
        self.sub_image = rospy.Subscriber('/front_camera/image_raw', Image, self.updateImage, queue_size=1)

        # ros to cv mat converter
        self.bridge = CvBridge()

        # set publishers
        self.detect_pub = rospy.Publisher('/detected', Image, queue_size = 1)
        self.bbox_pub = rospy.Publisher('/bbox', Float32MultiArray, queue_size = 1)
        self.traff_light_pub = rospy.Publisher('/traffic_light',TrafficLightArray, queue_size = 1)
        self.traff_sign_pub = rospy.Publisher('/traffic_sign',TrafficSignArray, queue_size = 1)

        # load label to names mapping for visualization purposes
        self.labels_to_names = {0: 'person', 1: 'bicycle', 2: 'car', 3: 'motorcycle', 4: 'airplane', 5: 'bus', 6: 'train', 7: 'truck', 8: 'boat'
                                ,9: 'traffic light', 10: 'fire hydrant', 11: 'stop sign', 12: 'parking meter', 13: 'bench', 14: 'bird', 15: 'cat'
                               ,16: 'dog', 17: 'horse', 18: 'sheep', 19: 'cow', 20: 'elephant', 21: 'bear', 22: 'zebra', 23: 'giraffe', 24: 'backpack'
                               ,25: 'umbrella', 26: 'handbag', 27: 'tie', 28: 'suitcase', 29: 'frisbee', 30: 'skis', 31: 'snowboard', 32: 'sports ball'
                               ,33: 'kite', 34: 'baseball bat', 35: 'baseball glove', 36: 'skateboard', 37: 'surfboard', 38: 'tennis racket', 39: 'bottle'
                               ,40: 'wine glass', 41: 'cup', 42: 'fork', 43: 'knife', 44: 'spoon', 45: 'bowl', 46: 'banana', 47: 'apple', 48: 'sandwich'
                               ,49: 'orange', 50: 'broccoli', 51: 'carrot', 52: 'hot dog', 53: 'pizza', 54: 'donut', 55: 'cake', 56: 'chair', 57: 'couch'
                               ,58: 'potted plant', 59: 'bed', 60: 'dining table', 61: 'toilet', 62: 'tv', 63: 'laptop', 64: 'mouse', 65: 'remote'
                               ,66: 'keyboard', 67: 'cell phone', 68: 'microwave', 69: 'oven', 70: 'toaster', 71: 'sink', 72: 'refrigerator', 73: 'book'
                               ,74: 'clock', 75: 'vase', 76: 'scissors', 77: 'teddy bear', 78: 'hair drier', 79: 'toothbrush'}

        # begin processing images
        self.p2 = np.array([[553.144531, 0.000000, 383.442257, 0.000000], [0.000000, 552.951477, 348.178354, 0.000000], [0.000000, 0.000000, 1.000000, 0]])
        self.p2_inv = np.linalg.pinv(self.p2)
        self.processImages()

    # signal interrupt handler, immediate stop of CAN_Controller
    def signalInterruptHandler(self, signum, frame):
        #print("Camera Detection - Exiting ROS Node...")

        # shut down ROS node and stop processess
        rospy.signal_shutdown("CTRL+C Signal Caught.")
        self.active = False

        sys.exit()

    def updateImage(self, image_msg):
        # convert and resize
        image = self.bridge.imgmsg_to_cv2(image_msg, "rgb8")
        #image, scale = resize_image(image)
        # save image in class
        self.image = image

    def processImages(self):
        #print("Waiting for Images...")
        while(self.image is None):
            pass

        while(self.active):
            self.retinanet(self.image)

    def getDistance(self, box,width,pWidth,dist,calibration):
        focalPoint = (pWidth * dist)/width
        #print("Focal point",focalPoint)
        focalPoint *= calibration
        pixW = min(abs(box[2]-box[0]),abs(box[3]-box[1]))
        #print(pixW)
        distance = (width*focalPoint)/abs(box[2]-box[0])
        distance = distance/12

        #print("Distance: ",(distance-5)*.3048)#*alpha)

        return (distance-5)*.3048

    def detect_color(self, image,box):
        img = image

        #img = cv2.imread(imageName)
        #img = cv2.cvtColor(img,cv2.COLOR_BGR2RGB)

        #img2 = Image.fromarray(img)
        #img2 = img2.crop((box[0],box[1],box[2],box[3]))
        crop_img = img[int(box[1]):int(box[3]), int(box[0]):int(box[2])]

        img3 = np.array(crop_img)

        #px = img3[10,10]
        ##print(px,px[0])
        #plt.imshow(img3)

        green_score = 0
        red_score = 0

        width = int(box[2] - box[0])
        height = int(box[3] - box[1])

        for x in range(0,width):
            for y in range(0,height):
                px = img3[y,x]
                if( ((px[0]>3) and (px[0] < 98)) and ((px[1]>185) and (px[1] <255)) and ((px[2] > 27) and (px[2] <119))):
                    green_score+= 1
                if( ((px[0]>185) and (px[0] < 255)) and ((px[1]>3) and (px[1] <58)) and ((px[2] > 9) and (px[2] <65))):
                    red_score+= 1

        if(green_score > red_score):
            return 5
        else:
            return 1

    def retinanet(self, image):
        image = image[0:1500,:]

        image_copy = image

        # convert cv mat to np array through PIL
        image = cv2.cvtColor(image, cv2.COLOR_RGB2BGR)
        image = PIL.Image.fromarray(image)
        image = np.asarray(image.convert('RGB'))[:, :, ::-1].copy()

        # copy to draw on
        draw = image.copy()
        #draw = cv2.cvtColor(draw, cv2.COLOR_BGR2RGB)

        # preprocess image for network
        image = preprocess_image(image)
        image, scale = resize_image(image)

        # process image
        start = time.time()
        #print("~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~")
        #print("Processing Image...")
        boxes, scores, labels = model.predict(np.expand_dims(image, axis=0))
        #print("processing time: ", time.time() - start)

        # correct for image scale
        boxes /= scale

        # visualize detections
        traffic_light_array = TrafficLightArray()
        traffic_sign_array = TrafficSignArray()
        obstacle_array = ObstacleArray()
        for box, score, label in zip(boxes[0], scores[0], labels[0]):
            # scores are sorted so we can break
            if score < 0.5:
                break

            color = label_color(label)

            b = box.astype(int)
            draw_box(draw, b, color=color)

            caption = "{} {:.3f}".format(self.labels_to_names[label], score)
            draw_caption(draw, b, caption)

            array_msg = Float32MultiArray()

            if(label == 9): #is a traffic light
            #120.5 ft
            #250 inches
                traff_light_msg = TrafficLight()
                what_color = self.detect_color(image_copy,box)
                distance = self.getDistance(box,15.25,84.96498,250,1.25)
                traff_light_msg.type = what_color
                traff_light_msg.x = distance
                traff_light_msg.y = 0

                if what_color == 1:
                    pass
                    #print(box,"red")
                else:
                    pass
                    #print(box,"green")
                traffic_light_array.lights.append(traff_light_msg)

            if(label == 11): #is a stopsign
                traff_sign_msg = TrafficSign()
                #getDistance(self, box,width,pWidth,dist,calibration)
                #distance = self.getDistance(box,30,93.454895,180,1.15)
                distance = self.getDistance(box,30,240.78,160,1.15)

                traff_sign_msg.type = 1
                traff_sign_msg.x = distance
                traff_sign_msg.y = 0
                traffic_sign_array.signs.append(traff_sign_msg)

                #array_msg.data = list(box)
                #self.bbox_pub.publish(array_msg)

            if(label == 2):
                #send info to Nick
                u_min, u_max = box[0]/2.56,box[2]/2.56
                v_min, v_max = box[1]/2.56, box[3]/2.56
                width = u_max - u_min
                if box[2] > 300 and width > 10 and box[0] < 1800:
                    print(box)
                    array_msg.data = list(box)
                    distance = self.getDistance(box,64.4,196,466, 1.13)
                    print("car distance", distance)
                    max_hom = np.array([u_max, v_max, 1])
                    min_hom = np.array([u_min, v_max, 1])
                    max_p = np.dot(self.p2_inv, max_hom)[:3]
                    min_p = np.dot(self.p2_inv, min_hom)[:3]
                    if max_p[2] < 0:
                        max_p[0], max_p[2] = -max_p[0], -max_p[2]
                    if min_p[2] < 0:
                        min_p[0], min_p[2] = -min_p[0], -min_p[2]

                    max_frustum_angle = (np.arctan2(max_p[0], max_p[2])) * -1
                    print("max angle", max_frustum_angle)
                    min_frustum_angle = (np.arctan2(min_p[0], min_p[2])) * -1
                    print("min angle", min_frustum_angle)
                    right = (distance-1) * np.tan(max_frustum_angle)
                    left = (distance-1) * np.tan(min_frustum_angle)

                    obs = Obstacle()
                    obs.x1 = distance
                    obs.x2 = distance + 1.85
                    obs.y1 = right
                    obs.y2 = left

                    obstacle_array.obstacles.append(obs)
                    print("left:", left,"right:", right,"distance", distance)
            #print("I see a " + str(self.labels_to_names[label]) + "(" + str(score) + ")")

        self.obstacle_publisher.publish(obstacle_array)
        self.traff_light_pub.publish(traffic_light_array)
        self.traff_sign_pub.publish(traffic_sign_array)

        plt.figure(figsize=(15, 15))
        plt.axis('off')
        #plt.imshow(draw)
        #plt.show()
        #cv2.imshow("hi", draw)
        #cv2.waitKey(1)

        try:
            self.detect_pub.publish(self.bridge.cv2_to_imgmsg(draw, "bgr8"))
        except CvBridgeError as e:
            pass#print(e)

if __name__ == "__main__":
    # start camera detection ros node
    CameraDetectionROSNode()


import gym
import highway_env
import numpy as np
from highway_env.utils import lmap
import torch
from nets import ReplayBuffer, CLEncoder, CLLoss, CLPP, CLDeconder
import torch.optim as optim
import matplotlib.pyplot as plt
from sklearn.manifold import TSNE
import torch.nn.functional as F

bf = ReplayBuffer(load_pid=True)
# clen = CLEncoder()
# clloss = CLLoss()
clpp = CLPP()
clde = CLDeconder()
clen = torch.load("clen_model.pkl")
clloss = torch.load("clloss_model.pkl")
# clpp = torch.load("clpp_model.pkl")
# clde = torch.load("clde_model.pkl")

# embedding = torch.nn.Embedding(9, 24)
cl_optimizer = optim.Adam(clen.parameters(), lr=1e-3)
# em_optimizer = optim.Adam(embedding.parameters(), lr=1e-3)
pp_optimizer = optim.Adam(clpp.parameters(), lr=1e-3)
de_optimizer = optim.Adam(clde.parameters(), lr=1e-3)
em = np.load("embedding.npy")

def pid_collector(nums=10000, save_pid_data=False, bf=None):
    s0_his, ss_his, aa_his, tt_his, pp_his, label_his = [], [], [], [], [], []
    s0_his_n, ss_his_n, aa_his_n, tt_his_n, pp_his_n = [], [], [], [], []
    road_r_his = []
    max_v = 30
    max_omega = np.pi/3
    max_a = 5  # [m/s]
    for i in range(nums):
        print("pid collector 第 {} 次收集".format(i + 1))
        env = gym.make("myenv-r1-v0")
        labels_index = np.arange(9).reshape(3, 3)
        N = 50

        lanes_count = env.config["lanes_count"]
        lane_id = np.random.choice(np.arange(lanes_count))

        if lane_id == 0:
            target_lane_id = np.random.choice([0, 1])
        elif lane_id == lanes_count - 1:
            target_lane_id = np.random.choice([lanes_count - 1, lanes_count - 2])
        else:
            target_lane_id = np.random.choice([lane_id - 1, lane_id, lane_id + 1])

        lon_operation = np.random.choice([0, 1, 2])  # 1保持 0减速 2加速
        v_lane_id = ("a", "b", lane_id)
        target_lane_id2 = ("a", "b", target_lane_id)
        v_target_s = (lon_operation - 1) * 5 + env.vehicle.speed
        v_target_s = np.clip(0, 30, v_target_s)

        positon_x = np.random.choice(np.arange(0, env.road.network.get_lane(v_lane_id).length, 5))
        positon_y = np.random.choice(np.arange(-2, 2.1, 0.5))
        heading = np.random.choice(
            env.road.network.get_lane(v_lane_id).heading_at(positon_x) + np.arange(-np.pi / 12, np.pi / 12, 10))
        speed = np.random.choice(np.arange(0, 25, 2))

        env.config["v_lane_id"] = v_lane_id
        env.config["v_target_id"] = target_lane_id2
        env.config["v_x"] = positon_x
        env.config["v_y"] = positon_y
        env.config["v_h"] = heading
        env.config["v_s"] = speed
        env.config["v_target_s"] = v_target_s

        env.reset()

        p = env.vehicle.position
        i_h = env.vehicle.heading
        i_s = env.vehicle.speed
        temp = [p[0], p[1], i_h, i_s]
        # radius*(end_phase - start_phase) * self.direction
        road_length = env.road.network.get_lane(v_lane_id).length
        road_r = env.config["radius"]
        road_r_his.append(road_r)
        temp_n = [lmap(p[0], [-(road_r+4), (road_r+4)], [-1, 1]), lmap(p[1], [0, 2*(road_r+4)], [-1, 1]),
                lmap(i_h, [-2*np.pi, 2*np.pi], [-1, 1]), lmap(i_s, [-max_v, max_v], [-1, 1])]
        x_road, y_road, x_road_n, y_road_n = env.vehicle.target_lane_position(p, road_r)

        action_his_omega = []
        action_his_omega_n = []
        action_his_accel = []
        action_his_accel_n = []
        action = 1
        x_his, y_his, h_his, s_his = [], [], [], []
        x_his_n, y_his_n, h_his_n, s_his_n = [], [], [], []

        for _ in range(N):
            # 应该有一个判断是不是出去了的函数

            env.step(action)
            action_his_omega.append(env.vehicle.action["steering"])
            action_his_omega_n.append(lmap(env.vehicle.action["steering"], [-max_omega, max_omega], [-1, 1]))
            action_his_accel.append(env.vehicle.action["acceleration"])
            action_his_accel_n.append(lmap(env.vehicle.action["acceleration"], [-max_a, max_a], [-1, 1]))

            x_his.append(env.vehicle.position[0])
            y_his.append(env.vehicle.position[1])
            h_his.append(env.vehicle.heading)
            s_his.append(env.vehicle.speed)

            x_his_n.append(lmap(env.vehicle.position[0], [-(road_r+4), (road_r+4)], [-1, 1]))
            y_his_n.append(lmap(env.vehicle.position[1], [0, 2*(road_r+4)], [-1, 1]))
            h_his_n.append(lmap(env.vehicle.heading, [-2 * np.pi, 2*np.pi], [-1, 1]))
            s_his_n.append(lmap(env.vehicle.speed, [-max_v, max_v], [-1, 1]))
            # env.render()
            # time.sleep(0.1)
        env.close()

        tt = temp + x_road + y_road + action_his_omega + action_his_accel
        tt_n = temp_n + x_road_n + y_road_n + action_his_omega_n + action_his_accel_n
        pp = x_his + y_his + h_his + s_his
        pp_n = x_his_n + y_his_n + h_his_n + s_his_n
        lane_change = target_lane_id - lane_id
        label = labels_index[lane_change + 1, lon_operation]
        s0 = temp
        s0_n = temp_n
        ss = x_road + y_road
        ss_n = x_road_n + y_road_n
        aa = action_his_omega + action_his_accel
        aa_n = action_his_omega_n + action_his_accel_n

        s0_his.append(s0)
        s0_his_n.append(s0_n)

        ss_his.append(ss)
        ss_his_n.append(ss_n)
        aa_his.append(aa)
        aa_his_n.append(aa_n)
        tt_his.append(tt)
        tt_his_n.append(tt_n)
        pp_his.append(pp)
        pp_his_n.append(pp_n)
        label_his.append(label)

        # (z, u, s0, X, ss, road_r)
        if save_pid_data is False:
            one_data = (label, aa_n, s0_n, pp_n, ss_n, road_r)
            bf.put(one_data)

    if save_pid_data:
        np.save("label_his.npy", np.array(label_his))
        np.save("aa_his.npy", np.array(aa_his))
        np.save("aa_his_n.npy", np.array(aa_his_n))
        np.save("s0_his.npy", np.array(s0_his))
        np.save("s0_his_n.npy", np.array(s0_his_n))
        np.save("pp_his.npy", np.array(pp_his))
        np.save("pp_his_n.npy", np.array(pp_his_n))
        np.save("ss_his.npy", np.array(ss_his))
        np.save("ss_his_n.npy", np.array(ss_his_n))
        np.save("tt_his.npy", np.array(tt_his))
        np.save("tt_his_n.npy", np.array(tt_his_n))

        np.save("road_r_his.npy", np.array(road_r_his))

        print("保存了PID的数据")

def train_cl(nums=10001, save_model=False):
    loss_his = []
    for i in range(nums):
        z_his, u_his, s0_his, X_his, ss_his, cl_his, road_r_his = bf.sample(1000)
        cl_his = torch.Tensor(cl_his)
        z_his = torch.Tensor(z_his)
        feat = clen(cl_his)
        loss = clloss(feat, z_his)
        loss_his.append(loss.item())
        cl_optimizer.zero_grad()
        loss.backward()
        cl_optimizer.step()

        if i % 100 == 0:
            print("第 {} 次训练的损失是 {}".format(i, np.mean(loss_his[-100:])))
    plt.figure(1)
    plt.plot(np.arange(len(loss_his)), loss_his)
    plt.title("CL loss")
    plt.show()
    if save_model:
        torch.save(clen, 'clen_model.pkl')
        torch.save(clloss, 'clloss_model.pkl')

def eval_cl():
    z_his, u_his, s0_his, X_his, ss_his, cl_his, road_r_his = bf.sample(500)
    cl_his = torch.Tensor(cl_his)
    feat = clen(cl_his)
    X = feat.detach().numpy()
    X_embedded = TSNE(n_components=2).fit_transform(X)
    plt.figure(2)
    plt.scatter(X_embedded[:, 0], X_embedded[:, 1])
    plt.title("CL")
    plt.show()

def train_embedding(nums=10, save_embedding=False):
    loss_his = []
    for i in range(nums):
        z_his, u_his, s0_his, X_his, ss_his, cl_his, road_r_his = bf.sample(1000)
        cl_his = torch.Tensor(cl_his)
        z_his = torch.Tensor(z_his)
        feats = clen(cl_his)
        labels = z_his.contiguous().view(-1, 1)

        loss = 0
        for j in range(9):
            mask = torch.eq(labels, j)
            tt = feats.masked_select(mask).view(-1, 24)

            # print(feats.shape)
            loss += torch.dist(tt.detach(), embedding.weight[j, :], 2)

        em_optimizer.zero_grad()
        loss.backward()
        em_optimizer.step()

        print("第 {} 次的 loss 是 {}".format(i+1, loss.item()))

    X = embedding.weight.detach().numpy()
    if save_embedding:
        np.save("embedding.npy", X)
    X_embedded = TSNE(n_components=2).fit_transform(X)
    plt.figure(3)
    plt.scatter(X_embedded[:, 0], X_embedded[:, 1])
    plt.title("Embedding")
    plt.show()

def train_pp(nums=10001, save_model=False):
    re_loss_his = []
    ni_loss_his = []
    loss_his = []
    for i in range(nums):
        z_his, u_his, s0_his, X_his, ss_his, cl_his, road_r_his = bf.sample(512)
        s0_his = torch.Tensor(s0_his)
        ss_his = torch.Tensor(ss_his)
        cl_his = torch.Tensor(cl_his)
        labels_his = torch.Tensor(z_his)
        X_his = torch.Tensor(X_his)
        # embedds = embedding.weight
        embedds = torch.Tensor(em)
        pp_hat, s0_mu, ss_mu, s0_log_sigma, ss_log_sigma = clpp(s0_his, ss_his, labels_his, embedds)
        X_F = X_his.reshape(512, -1, 50)
        X_F = X_F[:, :, -1]
        pp_F = pp_hat.reshape(512, -1, 50)
        pp_F = pp_F[:, :, -1]
        # 0.5 * sum(1 + log(sigma^2) - mu^2 - sigma^2)
        KLD_s0 = -0.5 * torch.sum(1 + s0_log_sigma - s0_mu.pow(2) - s0_log_sigma.exp())
        KLD_ss = -0.5 * torch.sum(1 + ss_log_sigma - ss_mu.pow(2) - ss_log_sigma.exp())
        loss = torch.dist(pp_hat, X_his) + 0.05 * (KLD_s0 + KLD_ss) + 2000 * torch.dist(X_F, pp_F)
        re_loss_his.append(torch.dist(X_F, pp_F).item())
        ni_loss_his.append(0.05 * (KLD_s0 + KLD_ss).item())
        loss_his.append(loss.item())
        # loss = torch.dist(pp_hat, X_his)

        pp_optimizer.zero_grad()
        loss.backward()
        pp_optimizer.step()

        if i % 100 == 0:
            print("第 {} 次的 loss 是 {}".format(i + 1, loss.item()))

    if save_model:
        torch.save(clpp, "clpp_model.pkl")

    plt.figure(4)
    plt.plot(np.arange(len(re_loss_his)), re_loss_his)
    plt.title("final state recon loss")
    # plt.figure(5)
    # plt.scatter(np.arange(len(ni_loss_his)), ni_loss_his)
    # plt.figure(6)
    # plt.scatter(np.arange(len(loss_his)), loss_his)
    plt.show()

def train_de(nums=5001, save_model=False):
    loss_his = []
    for i in range(nums):
        z_his, u_his, s0_his, X_his, ss_his, cl_his, road_r_his = bf.sample(512)
        s0_his = torch.Tensor(s0_his)
        ss_his = torch.Tensor(ss_his)
        cl_his = torch.Tensor(cl_his)
        labels_his = torch.Tensor(z_his)
        X_his = torch.Tensor(X_his)
        u_his = torch.Tensor(u_his)

        labels = labels_his.type(torch.LongTensor)
        labels = F.one_hot(labels)
        labels = labels.type(torch.float)
        skills = torch.matmul(labels, torch.Tensor(em))

        actions = clde(s0_his, ss_his, skills)
        feat = clen(torch.cat((s0_his, ss_his, actions), 1))

        # labels_his_hat = torch.argmin(torch.cdist(feat, embedding.weight.detach()), dim=1, keepdim=True)
        # labels_his_hat = torch.argmin(torch.cdist(feat, torch.Tensor(em)), dim=1)
        # re_loss = F.mse_loss(labels_his_hat, labels_his)
        re_loss = torch.dist(u_his, actions)
        norm_loss = actions.norm(dim=1).sum()
        loss = re_loss + 0.001 * norm_loss
        loss_his.append(re_loss.item())

        if i % 100 == 0:
            print("第 {} 次的 loss 是 {}".format(i + 1, loss.item()))

        de_optimizer.zero_grad()
        loss.backward()
        de_optimizer.step()

    if save_model:
        torch.save(clde, "clde_model.pkl")

    plt.figure(7)
    plt.plot(np.arange(len(loss_his)), loss_his)
    plt.title("de recon loss")
    plt.show()

# pid_collector(nums=5000, save_pid_data=True) # 用来离线收集PID数据
# train_cl(nums=2000, save_model=False)
# eval_cl()
# train_embedding(save_embedding=False)

train_pp(nums=2000, save_model=False)

train_de(save_model=False)


import os
import sys
sys.path.append('../')
import time
import numpy
from six.moves import xrange
import tensorflow as tf
import math
import numpy as np
from i3d import InceptionI3d
from utils import *
import cv2
import PIL.Image as Image
import input_data

flags = tf.app.flags
gpu_num = 1
flags.DEFINE_integer('batch_size', 1, 'Batch size.')
flags.DEFINE_integer('num_frame_per_clib', 64, 'Nummber of frames per clib')
flags.DEFINE_integer('crop_size', 224, 'Crop_size')
flags.DEFINE_integer('rgb_channels', 3, 'Channels for input')
flags.DEFINE_integer('classics', 101, 'The num of class')
flags.DEFINE_integer("sample_rate", 1, "sample rate in each clip")
FLAGS = flags.FLAGS
os.environ["CUDA_VISIBLE_DEVICES"] = "0"

def addTrigger(images, indices, triggerSize=24, save_path=None):

    trig = np.load("trigger" + str(triggerSize) + ".npy")
    targetImgs = list()
    cnt = 0
    for k in range(len(images)):
        for i in range(triggerSize):
            for j in range(triggerSize):
                images[k][-(i + 1)][-(j + 1)] = trig[0][k][-(i + 1)][-(j + 1)]
                t = trig[0][k][-(i + 1)][-(j + 1)]
        if save_path:
            img = images[k]
            img = Image.fromarray(img.astype("uint8"))
            img.save(save_path + "image_{:05d}.jpg".format(indices[cnt]))
            cnt += 1

    return targetImgs

def run():

    pre_model_save_dir = "../models/rgb_imagenet_10000_6_64_0.0001_decay"
    apply_cl_file = '../traintestlist/enhance_trigger.txt'

    file = list(open(apply_cl_file, 'r'))
    num_videos = len(file)
    print("Number of videos={}".format(num_videos))

    with tf.Graph().as_default():
        rgb_images_placeholder, _, labels_placeholder, is_training = placeholder_inputs(
                        FLAGS.batch_size * gpu_num,
                        FLAGS.num_frame_per_clib / FLAGS.sample_rate,
                        FLAGS.crop_size,
                        FLAGS.rgb_channels
                        )

        with tf.variable_scope('RGB'):
            logit, _ = InceptionI3d(
                                num_classes=FLAGS.classics,
                                spatial_squeeze=True,
                                final_endpoint='Logits',
                                name='inception_i3d'
                                )(rgb_images_placeholder, is_training)

        loss = -tower_loss(logit, labels_placeholder)
        grad = tf.gradients(loss, rgb_images_placeholder)[0]

        rgb_variable_map = {}
        for variable in tf.global_variables():
            if variable.name.split("/")[0] == "RGB" and "Adam" not in variable.name.split("/")[-1]:
                rgb_variable_map[variable.name.replace(':0', '')] = variable
        saver = tf.train.Saver(var_list=rgb_variable_map, reshape=True)

        init = tf.global_variables_initializer()

        sess = tf.Session(
                        config=tf.ConfigProto(allow_soft_placement=True)
                        )
        sess.run(init)

    ckpt = tf.train.get_checkpoint_state(pre_model_save_dir)
    if ckpt and ckpt.model_checkpoint_path:
        print ("loading checkpoint %s,waiting......" % ckpt.model_checkpoint_path)
        saver.restore(sess, ckpt.model_checkpoint_path)
        print ("load complete!")

    batch_size = FLAGS.batch_size
    step = num_videos // batch_size
    cnt = 0
    acc_all = 0
    et_step = 10
    epsilon = 8
    trigger_size = 24
    lines = list(open(apply_cl_file, "r").readlines())
    output_root = "/data/UCF101/UCF-101_extract/TargetVideo_" + str(epsilon) + "_trig" + str(trigger_size)

    import shutil
    if os.path.exists(output_root):
        shutil.rmtree(output_root)
    os.mkdir(output_root)

    for i in range(step):
        start = i * batch_size
        rgb_val_images, flow_val_images, val_labels, _, _, _ = input_data.read_clip_and_label(
            filename=apply_cl_file,
            batch_size=batch_size,
            start_pos=start,
            num_frames_per_clip=FLAGS.num_frame_per_clib,
            crop_size=FLAGS.crop_size,
            shuffle=False
            )
        filename = lines[start].split(" ")[0]
        videoname = filename.split("/")[7]
        videopath = output_root + "/" + videoname + "/"
        os.mkdir(videopath)
        filename_img = filename.replace("UCF-101_extract_flow", "UCF-101_extract")
        frameNum = 0
        for parent, dirnames, filenames in os.walk(filename_img):
            frameNum = len(filenames)
        s_index = (frameNum - FLAGS.num_frame_per_clib) // 2
        indices = list(range(s_index + 1, s_index + 1 + FLAGS.num_frame_per_clib))

        x = rgb_val_images
        for i in range(et_step):
            grad_val, logit_val = sess.run([grad, logit], feed_dict={rgb_images_placeholder: x,
                                            labels_placeholder: val_labels,
                                            is_training: False})
            x = np.add(x, -1 * np.sign(grad_val), out=x, casting='unsafe')
            index = np.argmax(logit_val, axis=1)

            x = np.clip(x, rgb_val_images - epsilon, rgb_val_images + epsilon)
            x = np.clip(x, 0, 255)
        addTrigger(x[0], indices, trigger_size, save_path=videopath)
        print(start, videopath)

if __name__ == "__main__":
    run()


import numpy

"""
Constants : This files includes the basic constant used with polymuse from instrument codes, to inverse of it, to scale patterns etc ...

"""

def instrument_program_code(names):
    codes = numpy.zeros(len(names), dtype='int32')
    for i, nm in enumerate(names):
        codes[i] = inst_codes[nm.lower()]
    return codes

rev_inst_codes = {0: 'Acoustic Grand Piano', 1: 'Bright Acoustic Piano', 2: 'Electric Grand Piano', 3: 'Honky-tonk Piano', 4: 'Electric Piano 1', 5: 'Electric Piano 2', 6: 'Harpsichord', 7: 'Clavi', 8: 'Celesta', 9: 'Glockenspiel', 10: 'Music Box', 11: 'Vibraphone', 12: 'Marimba', 13: 'Xylophone', 14: 'Tubular Bells', 15: 'Dulcimer', 16: 'Drawbar Organ', 17: 'Percussive Organ', 18: 'Rock Organ', 19: 'Church Organ', 20: 'Reed Organ', 21: 'Accordion', 22: 'Harmonica', 23: 'Tango Accordion', 24: 'Acoustic Guitar (nylon)', 25: 'Acoustic Guitar (steel)', 26: 'Electric Guitar (jazz)', 27: 'Electric Guitar (clean)', 28: 'Electric Guitar (muted)', 29: 'Overdriven Guitar', 30: 'Distortion Guitar', 31: 'Guitar harmonics', 32: 'Acoustic Bass', 33: 'Electric Bass (finger)', 34: 'Electric Bass (pick)', 35: 'Fretless Bass', 36: 'Slap Bass 1', 37: 'Slap Bass 2', 38: 'Synth Bass 1', 39: 'Synth Bass 2', 40: 'Violin', 41: 'Viola', 42: 'Cello', 43: 'Contrabass', 44: 'Tremolo Strings', 45: 'Pizzicato Strings', 46: 'Orchestral Harp', 47: 'Timpani', 48: 'String Ensemble 1', 49: 'String Ensemble 2', 50: 'SynthStrings 1', 51: 'SynthStrings 2', 52: 'Choir Aahs', 53: 'Voice Oohs', 54: 'Synth Voice', 55: 'Orchestra Hit', 56: 'Trumpet', 57: 'Trombone', 58: 'Tuba', 59: 'Muted Trumpet', 60: 'French Horn', 61: 'Brass Section', 62: 'SynthBrass 1', 63: 'SynthBrass 2', 64: 'Soprano Sax', 65: 'Alto Sax', 66: 'Tenor Sax', 67: 'Baritone Sax', 68: 'Oboe', 69: 'English Horn', 70: 'Bassoon', 71: 'Clarinet', 72: 'Piccolo', 73: 'Flute', 74: 'Recorder', 75: 'Pan Flute', 76: 'Blown Bottle', 77: 'Shakuhachi', 78: 'Whistle', 79: 'Ocarina', 80: 'Lead 1 (square)', 81: 'Lead 2 (sawtooth)', 82: 'Lead 3 (calliope)', 83: 'Lead 4 (chiff)', 84: 'Lead 5 (charang)', 85: 'Lead 6 (voice)', 86: 'Lead 7 (fifths)', 87: 'Lead 8 (bass + lead)', 88: 'Pad 1 (new age)', 89: 'Pad 2 (warm)', 90: 'Pad 3 (polysynth)', 91: 'Pad 4 (choir)', 92: 'Pad 5 (bowed)', 93: 'Pad 6 (metallic)', 94: 'Pad 7 (halo)', 95: 'Pad 8 (sweep)', 96: 'FX 1 (rain)', 97: 'FX 2 (soundtrack)', 98: 'FX 3 (crystal)', 99: 'FX 4 (atmosphere)', 100: 'FX 5 (brightness)', 101: 'FX 6 (goblins)', 102: 'FX 7 (echoes)', 103: 'FX 8 (sci-fi)', 104: 'Sitar', 105: 'Banjo', 106: 'Shamisen', 107: 'Koto', 108: 'Kalimba', 109: 'Bag pipe', 110: 'Fiddle', 111: 'Shanai', 112: 'Tinkle Bell', 113: 'Agogo', 114: 'Steel Drums', 115: 'Woodblock', 116: 'Taiko Drum', 117: 'Melodic Tom', 118: 'Synth Drum', 119: 'Reverse Cymbal', 120: 'Guitar Fret Noise', 121: 'Breath Noise', 122: 'Seashore', 123: 'Bird Tweet', 124: 'Telephone Ring', 125: 'Helicopter', 126: 'Applause', 127: 'Gunshot'}

inst_codes = {
                'piano' : 0, 'acoustic grand piano': 0, 'bright acoustic piano': 1, 'electric grand piano': 2, 'honky-tonk piano': 3, 'electric piano 1': 4, 'electric piano 2': 5, 'harpsichord': 6, 'clavi': 7, 'celesta': 8,
                'glockenspiel': 9, 'music box': 10, 'vibraphone': 11, 'marimba': 12, 'xylophone': 13, 'tubular bells': 14, 'dulcimer': 15, 'drawbar organ': 16,
                'percussive organ': 17, 'rock organ': 18, 'church organ': 19, 'reed organ': 20, 'accordion': 21, 'harmonica': 22, 'tango accordion': 23, 'acoustic guitar (nylon)': 24,
                'guitar' : 25, 'acoustic guitar (steel)': 25, 'electric guitar (jazz)': 26, 'electric guitar (clean)': 27, 'electric guitar (muted)': 28, 'overdriven guitar': 29, 'distortion guitar': 30, 'guitar harmonics': 31, 'acoustic bass': 32,
                'electric bass (finger)': 33, 'electric bass (pick)': 34, 'fretless bass': 35, 'slap bass 1': 36, 'slap bass 2': 37, 'synth bass 1': 38, 'synth bass 2': 39, 'violin': 40,
                'viola': 41, 'cello': 42, 'contrabass': 43, 'tremolo strings': 44, 'pizzicato strings': 45, 'orchestral harp': 46, 'timpani': 47, 'string ensemble 1': 48,
                'string ensemble 2': 49, 'synthstrings 1': 50, 'synthstrings 2': 51, 'choir aahs': 52, 'voice oohs': 53, 'synth voice': 54, 'orchestra hit': 55, 'trumpet': 56,
                'trombone': 57, 'tuba': 58, 'muted trumpet': 59, 'french horn': 60, 'brass section': 61, 'synthbrass 1': 62, 'synthbrass 2': 63, 'soprano sax': 64,
                'alto sax': 65, 'tenor sax': 66, 'baritone sax': 67, 'oboe': 68, 'english horn': 69, 'bassoon': 70, 'clarinet': 71, 'piccolo': 72,
                'flute': 73, 'recorder': 74, 'pan flute': 75, 'blown bottle': 76, 'shakuhachi': 77, 'whistle': 78, 'ocarina': 79, 'lead 1 (square)': 80,
                'lead 2 (sawtooth)': 81, 'lead 3 (calliope)': 82, 'lead 4 (chiff)': 83, 'lead 5 (charang)': 84, 'lead 6 (voice)': 85, 'lead 7 (fifths)': 86, 'lead 8 (bass + lead)': 87, 'pad 1 (new age)': 88,
                'pad 2 (warm)': 89, 'pad 3 (polysynth)': 90, 'pad 4 (choir)': 91, 'pad 5 (bowed)': 92, 'pad 6 (metallic)': 93, 'pad 7 (halo)': 94, 'pad 8 (sweep)': 95, 'fx 1 (rain)': 96,
                'fx 2 (soundtrack)': 97, 'fx 3 (crystal)': 98, 'fx 4 (atmosphere)': 99, 'fx 5 (brightness)': 100, 'fx 6 (goblins)': 101, 'fx 7 (echoes)': 102, 'fx 8 (sci-fi)': 103, 'sitar': 104,
                'banjo': 105, 'shamisen': 106, 'koto': 107, 'kalimba': 108, 'bag pipe': 109, 'fiddle': 110, 'shanai': 111, 'tinkle bell': 112,
                'agogo': 113, 'steel drums': 114, 'woodblock': 115, 'taiko drum': 116, 'melodic tom': 117, 'synth drum': 118, 'reverse cymbal': 119, 'guitar fret noise': 120,
                'breath noise': 121, 'seashore': 122, 'bird tweet': 123, 'telephone ring': 124, 'helicopter': 125, 'applause': 126, 'gunshot': 127
            }

chord_prog = {
                'major-3' : [4, 5, 0, 0],
                'minor-3' : [3, 5, 0, 0]
            }

chord_prog_to_codes = {
                'major-3' : 0,
                'minor-3' : 1,
            }
codes_to_chord_prog = {
                0 : 'major-3',
                1 : 'minor-3',
            }

scale_patterns = {
                'major' : [0, 2, 2, 1, 2, 2, 2, 1],
                'minor' : [0, 2, 1, 2, 2, 2, 1, 2],
            }

scale_patterns_cum = {
                'major' : [0, 2, 4, 5, 7, 9, 11, 12],
                'minor' : [0, 2, 3, 5, 7, 8, 10, 12],
            }

scale_names = {
        'major' : ['C', 'C#', 'D ', 'D#', 'E ', 'F ', 'F#', 'G ', 'G#', 'A ', 'A#', 'B '],
        'minor' : ['Cm', 'C#m', 'Dm ', 'D#m', 'Em ', 'Fm ', 'F#m', 'Gm ', 'g#m', 'Am ', 'A#m', 'Bm '],
    }

lead_track = [0, 1, 2, 4, 6, 7, 8,  13,  15,  33, 34, 35, 36, 37, 38, 39, 40, 41, 42, 43, 44, 45, 46, 47, 48, 49,  70,  72]

model_dict = {
        'lead' : 'lead', 'piano' : 'lead', 'first' : 'lead',
        'chorus' : 'chorus', 'second' : 'chorus',
        'drum' : 'drum', 'third' : 'drum'
    }

depths_of_3tracks = {
                    0 : 1, 1: 3, 2: 2,
                    'lead' : 1, 'chorus' : 3, 'drum': 2
                    }

timec_of_3tracks = {
                    0 : 4, 1: 16, 2: 4,
                    }

predict_instances = {
                    0 : 256, 1 : 64, 2: 128
                    }

velocity = {
            0 : 90, 1 : 90, 2: 70
            }

type3tracks = {
                0 : 'lead', 1: 'chorus', 2: 'drum'
                }

models_csv = 'https://raw.githubusercontent.com/rushike/polymuse-future/master/models.csv'

midis_csv = 'https://raw.githubusercontent.com/rushike/polymuse-future/master/midis.csv'


import numpy as np
import time
import cv2
import os.path as path

class OpenCV_3():
    """
    Esta es una manera que no utiliza el tradicional Haar de OpenCV, pues al ser basados en FaceRest es algo lento
    y aumentar la velocidad del reconocimiento es posible con el LBP de OpenCV, se pierde presición
    para ello, se presenta un modelo de DeepLearning que usa los algoritmos de FaceNet de google
    para buscar los rostros de manera optima y rapida
    """
    def __init__(self):
        protocolo = "deploy.prototxt.txt"
        model = "res10_300x300_ssd_iter_140000.caffemodel"
        if not path.exists(model):
            print("No se encuetra el archivo model")
        elif not path.exists(protocolo):
            print("No se encuetra el archivo protocolo")
        else:
            self.net = cv2.dnn.readNetFromCaffe(protocolo, model)
        self.confianza = 0.55

    def setConfianzaDeteccion(self, confianza):
        self.confianza = confianza

    def detectarRostros(self, imagen):
        (h, w) = imagen.shape[:2]
        blob = cv2.dnn.blobFromImage(
                cv2.resize(imagen, (300, 300)),
                1.0,
                (300, 300),
                (104.0, 177.0, 123.0)
            )
        self.net.setInput(blob)
        self.detections = self.net.forward()
        listadoCoordenadasRostro = []
        for i in range(0, self.detections.shape[2]):
            confianzaEsteRostro = self.detections[0, 0, i, 2]
            if confianzaEsteRostro >= self.confianza:
                box = self.detections[0, 0, i, 3:7] * np.array([w, h, w, h])
                (startX, startY, endX, endY) = box.astype("int")
                coordenadas = (startX, startY, endX, endY, confianzaEsteRostro)
                listadoCoordenadasRostro.append(coordenadas)

        return listadoCoordenadasRostro


#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
kuehbach at fhi - berlin . mpg . de
utilities for parsing ranging definitions for aptfim
2021/03/02
"""

import os, sys, glob
#from pathlib import Path
import re, mmap
import numpy as np
import warnings
import periodictable as pse

#replace by packages later
#from ....fairmat_areab_parser.utils.nomad4exp_keyvalue_struct import n4eKeyValue
#basePath = '/home/kuehbach/GITHUB/FAIRMAT-PARSER/fairmat_areab_parser'
#sys.path.append(basePath + '/utils/')
#from nomad4exp_python_modules import *
from aptfimleapparser.utils import nomad4exp_keyvalue_struct
##MK::move to the above module

#numerical constants
MQ_EPSILON = 1.0e-4
MAX_NUMBER_OF_ATOMS_PER_MOLECULAR_ION = 8

#https://stackoverflow.com/questions/3663450/remove-substring-only-at-the-end-of-string
def rchop(s, suffix):
    if suffix and s.endswith(suffix):
        return s[:-len(suffix)]
    return s

def hash_isotope(Nprotons,Nneutrons):
    if Nprotons >= 0 and Nprotons < 256 and Nneutrons >= 0 and Nneutrons < 256:
        return np.uint16(Nprotons) + np.uint16(256)*np.uint16(Nneutrons)
    else:
        raise ValueError('Nprotons and Nneutrons need to be on interval [0, 256) !')

def unhash_isotope(hashval):
    if type(hashval) == np.uint16: #implicit on [0,2**16-1]
        Nneutrons = np.uint16(hashval / np.uint16(256))
        Nprotons = np.uint16(hashval - Nneutrons*np.uint16(256))
        return [Nprotons, Nneutrons]
    else:
        raise ValueError('Hashval needs to be a np.uint16 !')

class ion():
    def __init__(self, *args, **kwargs):
        self.name = None
        self.id = np.uint8(0)
        self.charge = np.int8(0)
        self.ivec = np.zeros( 32, np.dtype(np.uint16) ) #list of isotope hash values
        self.mq = [] #list of associated mass-to-charge state ratios in Da

    def set_name(self, nm):
        if nm == None:
            self.name = None
        else:
            if type(nm) == str and len(nm) > 0:
                self.name = np.str(nm)
            else:
                raise ValueError('Name has to be a string and must not be an empty string !')

    def set_id(self, id):
        if id == None:
            self.id = None
        else:
            if (id >= 0 and id < 256):
                self.id = np.uint8(id)
            else:
                raise ValueError('ID needs to be on [0, 255] !')

    def set_charge(self, chrg):
        if chrg == None:
            self.charge = np.int8(0)
        else:
            if np.int8(abs(chrg)) <= np.int8(8):
                self.charge = np.int8(chrg)
            else:
                raise ValueError('Charge needs to be on (-8,+8) !')

    def set_ivec(self, ivec):
        if len(np.shape(ivec)) == 1:
            if np.shape(ivec)[0] > 0 and np.shape(ivec)[0] <= MAX_NUMBER_OF_ATOMS_PER_MOLECULAR_ION:
                self.ivec = np.zeros( MAX_NUMBER_OF_ATOMS_PER_MOLECULAR_ION, np.dtype(np.uint16) )
                for i in np.arange(0,np.shape(ivec)[0]):
                    self.ivec[i] = np.uint16(ivec[0])
            else:
                raise ValueError('Isotope vector needs to be an uint16 array with at most ' + str(MAX_NUMBER_OF_ATOMS_PER_MOLECULAR_ION) + ' values !')
        else:
            raise ValueError('Isotope vector needs to be an uint16 array with at most ' + str(MAX_NUMBER_OF_ATOMS_PER_MOLECULAR_ION) + ' values !')


from . import BaseAgent
from .. import characters
from ..constants import *
from os import path
import numpy as np
import ctypes
import time
import sys

class GottingenAgent(BaseAgent):
    """Parent abstract Agent."""
    turn_times = []

    def __init__(self, character=characters.Bomber):
        self._character = character
        self.avg_simsteps_per_turns = []

    def __getattr__(self, attr):
        return getattr(self._character, attr)

    def act(self, obs, action_space):
        #for attr in ['alive', 'board', 'bomb_life', 'bomb_blast_strength', 'position', 'blast_strength', 'can_kick', 'teammate', 'ammo', 'enemies']:
        #for attr in ['game_type', 'flame_life', 'bomb_moving_direction', 'step_count', 'game_env']:
        #in radio: message
        #    print(attr, type(obs[attr]), obs[attr])
        #    if 'numpy' in str(type(obs[attr])):
        #        print(' ', obs[attr].dtype)

        #print(obs['enemies'][0].__dict__) is sent to cpp, because only contains enum ids?
        #print(obs['position'][0])
        #print(type(obs['position'][0]))
        start_time = time.time()

        if isinstance(obs['bomb_moving_direction'], list):
            obs['bomb_moving_direction'] = np.asarray(obs['bomb_moving_direction'], dtype=np.float32)
        if isinstance(obs['flame_life'], list):
            obs['flame_life'] = np.asarray(obs['flame_life'], dtype=np.float32)

        decision = self.c.c_getStep_gottingen(
            self.id,
            Item.Agent0.value in obs['alive'], Item.Agent1.value in obs['alive'], Item.Agent2.value in obs['alive'], Item.Agent3.value in obs['alive'],
            obs['board'].ctypes.data_as(ctypes.POINTER(ctypes.c_uint8)),
            obs['bomb_life'].ctypes.data_as(ctypes.POINTER(ctypes.c_double)),
            obs['bomb_blast_strength'].ctypes.data_as(ctypes.POINTER(ctypes.c_double)),
            obs['bomb_moving_direction'].ctypes.data_as(ctypes.POINTER(ctypes.c_double)),
            obs['flame_life'].ctypes.data_as(ctypes.POINTER(ctypes.c_double)),
            int(obs['position'][0]), int(obs['position'][1]),
            obs['blast_strength'],
            obs['can_kick'],
            obs['ammo'],
            obs['game_type'],
            obs['teammate'].value,
            obs['message'][0] if 'message' in obs else -1,
            obs['message'][1] if 'message' in obs else -1
            )
        self.turn_times.append(time.time() - start_time)

        if GameType(obs['game_type']) == GameType.TeamRadio:
            return [decision, self.c.c_getMessage_gottingen(self.id, 0), self.c.c_getMessage_gottingen(self.id, 1)]
        else:
            return decision

    def episode_end(self, reward):
        """This is called at the end of the episode to let the agent know that
        the episode has ended and what is the reward.

        Args:
          reward: The single reward scalar to this agent.
        """
        self.c.c_episode_end_gottingen.restype = ctypes.c_float
        avg_simsteps_per_turn = self.c.c_episode_end_gottingen(self.id)
        self.avg_simsteps_per_turns.append(avg_simsteps_per_turn)

    def init_agent(self, id, game_type):
        self.id = id
        self._character = self._character(id, game_type)
        self.c = None

        if sys.platform == "win32":
            self.c = ctypes.cdll.LoadLibrary("build/Release/pommerman.dll")
        else:
            possible_paths = ["/opt/work/pommerman_cpp/cmake-build-debug/", "../cmake-build-debug/", "./cmake-build-debug/"]
            for p in possible_paths:
                if path.exists(p + "/libpommerman.so"):
                    self.c = ctypes.cdll.LoadLibrary(p + "/libpommerman.so")
                    break
            if not self.c:
                print("Pommerman C++ library not found")

        self.c.c_init_agent_gottingen(id)

    @staticmethod
    def has_user_input():
        return False

    def shutdown(self):
        turn_times_np = np.array(self.turn_times)
        print("gottingen shutdown, avg simsteps per turns: ", np.round(np.array(self.avg_simsteps_per_turns).mean())/1000.0, " k",
              ", avg time: ", np.round(turn_times_np.mean()*1000, 1), " ms",
              ", max time: ", np.round(turn_times_np.max()*1000, 1), " ms",
              ", overtime: ", np.round((turn_times_np > 0.149).sum() / turn_times_np.shape[0] *100, 1), ' %')


from __future__ import absolute_import, division, print_function

import math
import os
import pickle as cp

import numpy as np
import rdkit
import torch
from gln.common.reactor import Reactor
from gln.data_process.data_info import DataInfo, load_bin_feats
from gln.graph_logic.logic_net import GraphPath
from gln.mods.mol_gnn.mol_utils import SmartsMols, SmilesMols
from rdkit import Chem
from scipy.special import softmax

class RetroGLN(object):
    def __init__(self, dropbox, model_dump, args=None): # minhtoo added args as parameter into RetroGLN instantiation (10th May)
        """
        Args:
            dropbox: the dropbox folder
            model_dump: the ckpt folder, which contains model dump and model args
        """
        assert os.path.isdir(model_dump)

        if args is not None: # minhtoo add
            print("Using user-provided args")
            self.args = args
            self.args.dropbox = dropbox
        else:
            arg_file = os.path.join(model_dump, 'args.pkl')
            print(f"Loading args from {arg_file}")
            with open(arg_file, 'rb') as f:
                self.args = cp.load(f)
                self.args.dropbox = dropbox

        DataInfo.init(dropbox, self.args)
        load_bin_feats(dropbox, self.args)

        model_file = os.path.join(model_dump, 'model.dump')
        print(f"Loading model from {model_file}") # minhtoo add
        self.gln = GraphPath(self.args)
        try:
            self.gln.load_state_dict(torch.load(model_file))
            self.gln.cuda()
        except:
            self.gln.load_state_dict(torch.load(model_file, map_location=torch.device('cpu')))
        self.gln.eval()

        self.prod_center_maps = {}
        self.cached_smarts = None

    def _ordered_tpls(self, cano_prod, beam_size, rxn_type):
        if (rxn_type, cano_prod) not in self.prod_center_maps:
            mol = Chem.MolFromSmiles(cano_prod)
            if mol is None:
                return None
            if self.cached_smarts is None:
                self.cached_smarts = []
                print('caching smarts centers')
                for sm in DataInfo.prod_cano_smarts:
                    self.cached_smarts.append(Chem.MolFromSmarts(sm))

            prod_center_cand_idx = []
            for i, sm in enumerate(self.cached_smarts):
                if sm is not None and mol.HasSubstructMatch(sm):
                    prod_center_cand_idx.append(i)
            self.prod_center_maps[(rxn_type, cano_prod)] = prod_center_cand_idx
        prod_center_cand_idx = self.prod_center_maps[(rxn_type, cano_prod)]

        # infer the reaction center
        if not len(prod_center_cand_idx):
            return None
        prod_center_mols = [SmartsMols.get_mol_graph(DataInfo.prod_cano_smarts[m]) for m in prod_center_cand_idx]
        prod_mol = SmilesMols.get_mol_graph(cano_prod)
        prod_center_scores = self.gln.prod_center_predicate.inference([prod_mol], [prod_center_mols])
        prod_center_scores = prod_center_scores.view(-1).data.cpu().numpy()
        top_centers = np.argsort(-1 * prod_center_scores)[:beam_size]
        top_center_scores = [prod_center_scores[i] for i in top_centers]
        top_center_mols = [prod_center_mols[i] for i in top_centers]
        top_center_smarts = [DataInfo.prod_cano_smarts[prod_center_cand_idx[i]] for i in top_centers]

        # infer the template
        list_of_list_tpls = []
        for i, c in enumerate(top_center_smarts):
            assert c in DataInfo.unique_tpl_of_prod_center
            if not rxn_type in DataInfo.unique_tpl_of_prod_center[c]:
                continue
            tpl_indices = DataInfo.unique_tpl_of_prod_center[c][rxn_type]
            tpls = [DataInfo.unique_templates[t][1] for t in tpl_indices]
            list_of_list_tpls.append(tpls)
        if not len(list_of_list_tpls):
            return None
        tpl_scores = self.gln.tpl_fwd_predicate.inference([prod_mol] * len(top_center_mols), list_of_list_tpls)
        tpl_scores = tpl_scores.view(-1).data.cpu().numpy()

        idx = 0
        tpl_with_scores = []
        for i, c in enumerate(top_center_scores):
            for tpl in list_of_list_tpls[i]:
                t_score = tpl_scores[idx]
                tot_score = c + t_score
                tpl_with_scores.append((tot_score, tpl))
                idx += 1
        tpl_with_scores = sorted(tpl_with_scores, key=lambda x: -1 * x[0])

        return tpl_with_scores

    def run(self, raw_prod, beam_size, topk, rxn_type='UNK'):
        """
        Args:
            raw_prod: the single product smiles
            beam_size: the size for beam search
            topk: top-k prediction of reactants
            rxn_type: (optional) reaction type
        Return:
            a dictionary with the following keys:
            {
                'reactants': the top-k prediction of reactants
                'template': the list of corresponding reaction templates used
                'scores': the scores for the corresponding predictions, in descending order
            }
            if no valid reactions are found, None will be returned
        """
        cano_prod = DataInfo.get_cano_smiles(raw_prod)
        prod_mol = SmilesMols.get_mol_graph(cano_prod)
        tpl_with_scores = self._ordered_tpls(cano_prod, beam_size, rxn_type)

        if tpl_with_scores is None:
            return None
        # filter out invalid tpls
        list_of_list_reacts = []
        list_reacts = []
        list_tpls = []
        num_tpls = 0
        num_reacts = 0
        for prod_tpl_score, tpl in tpl_with_scores:
            pred_mols = Reactor.run_reaction(raw_prod, tpl)
            if pred_mols is not None and len(pred_mols):
                num_tpls += 1
                list_of_list_reacts.append(pred_mols)
                num_reacts += len(pred_mols)
                list_tpls.append((prod_tpl_score, tpl))
                if num_tpls >= beam_size:
                    break

        list_rxns = []
        for i in range(len(list_of_list_reacts)):
            list_rxns.append([DataInfo.get_cano_smiles(r) + '>>' + cano_prod for r in list_of_list_reacts[i]])
        if len(list_rxns) and len(list_tpls):
            react_scores = self.gln.reaction_predicate.inference([prod_mol] * len(list_tpls), list_rxns)
            react_scores = react_scores.view(-1).data.cpu().numpy()

            idx = 0
            final_joint = []
            for i, (prod_tpl_score, tpl) in enumerate(list_tpls):
                for reacts in list_of_list_reacts[i]:
                    r_score = react_scores[idx]
                    tot_score = prod_tpl_score + r_score
                    final_joint.append((tot_score, tpl, reacts))
                    idx += 1
            final_joint = sorted(final_joint, key=lambda x: -1 * x[0])[:topk]
            scores = [t[0] for t in final_joint]
            scores = softmax([each_score for each_score in scores])
            list_reacts = [t[2] for t in final_joint]
            ret_tpls = [t[1] for t in final_joint]
            result = {'template': ret_tpls,
                    'reactants': list_reacts,
                    'scores': scores}
        else:
            result = {'template': [],
                        'reactants': [],
                        'scores': []}
        return result


# -*- mode: python; coding: utf-8 -*-
# Copyright 2019-2020 the .NET Foundation
# Licensed under the MIT License.

"""An image, possibly tiled, for display in WWT.

"""
from __future__ import absolute_import, division, print_function

__all__ = '''
ImageSet
'''.split()

import math
from traitlets import Bool, Float, Int, Unicode, UseEnum
from xml.etree import ElementTree as etree

from . import LockedXmlTraits, XmlSer
from .abcs import UrlContainer
from .enums import Bandpass, DataSetType, ProjectionType

class ImageSet(LockedXmlTraits, UrlContainer):
    """A set of images."""

    data_set_type = UseEnum(DataSetType, default_value=DataSetType.SKY).tag(xml=XmlSer.attr('DataSetType'))
    """The renderer mode to which these data apply.

    Possible values are ``"Earth"``, ``"Planet"``, ``"Sky"``, ``"Panorama"``,
    ``"SolarSystem"``, and ``"Sandbox"``.

    """
    reference_frame = Unicode('').tag(xml=XmlSer.attr('ReferenceFrame'))
    """TBD."""

    name = Unicode('').tag(xml=XmlSer.attr('Name'))
    """A name used to refer to this imageset.

    Various parts of the WWT internals reference imagesets by this name, so
    it should be distinctive.

    """
    url = Unicode('').tag(xml=XmlSer.attr('Url'))
    """The URL of the image data.

    Either a URL or a URL template. TODO: details

    """
    dem_url = Unicode('').tag(xml=XmlSer.attr('DemUrl'))
    """The URL of the DEM data.

    Either a URL or a URL template. TODO: details

    """
    width_factor = Int(2).tag(xml=XmlSer.attr('WidthFactor'))
    """This is a legacy parameter. Leave it at 2."""

    base_tile_level = Int(0).tag(xml=XmlSer.attr('BaseTileLevel'))
    """The level of the highest (coarsest-resolution) tiling available.

    This should be zero except for special circumstances.

    """
    quad_tree_map = Unicode('').tag(xml=XmlSer.attr('QuadTreeMap'))
    """TBD."""

    tile_levels = Int(0).tag(xml=XmlSer.attr('TileLevels'))
    """The number of levels of tiling.

    Should be zero for untiled images. An image with ``tile_levels = 1`` has been
    broken into four tiles, each 256x256 pixels. For ``tile_levels = 2``, there are
    sixteen tiles, and the padded height of the tiled area is ``256 * 2**2 = 1024``
    pixels. Image with dimensions of 2048 pixels or smaller do not need to be tiled,
    so if this parameter is nonzero it will usually be 4 or larger.

    """
    base_degrees_per_tile = Float(0.0).tag(xml=XmlSer.attr('BaseDegreesPerTile'))
    """The angular scale of the image.

    For untiled images, should be the pixel scale: the number of degrees per
    pixel in the vertical direction. Non-square pixels are not supported.

    For tiled images, this is the height of the image with its dimensions
    padded out to the next largest power of 2 for tiling purposes. If a square
    image is 1200 pixels tall and has a height of 0.016 deg, the padded height
    would be 2048 pixels and this parameter should be set to
    0.016 * 2048 / 1200 = 0.0273.

    """
    file_type = Unicode('.png').tag(xml=XmlSer.attr('FileType'))
    """The extension of the image file(s) in this set, including a leading period.

    """
    bottoms_up = Bool(False).tag(xml=XmlSer.attr('BottomsUp'))
    """TBD."""

    projection = UseEnum(
        ProjectionType,
        default_value = ProjectionType.SKY_IMAGE
    ).tag(xml=XmlSer.attr('Projection'))
    """The type of projection used to place this image on the sky.

    For untiled images, this should be "SkyImage". For tiled images, it should
    be "Tan". The :meth:`set_position_from_wcs` method will set this value
    appropriately based on :attr:`tile_levels`.

    """
    center_x = Float(0.0).tag(xml=XmlSer.attr('CenterX'))
    """The horizontal location of the center of the image’s projection coordinate
    system.

    For sky images, this is a right ascension in degrees.

    """
    center_y = Float(0.0).tag(xml=XmlSer.attr('CenterY'))
    """The vertical location of the center of the image’s projection coordinate
    system.

    For sky images, this is a declination in degrees.

    """
    offset_x = Float(0.0).tag(xml=XmlSer.attr('OffsetX'))
    """The horizontal positioning of the image relative to its projection
    coordinate system.

    For untiled sky images, the image is by default positioned such that its
    lower left lands at the center of the projection coordinate system (namely,
    ``center_x`` and ``center_y``). The offset is measured in pixels and moves
    the image leftwards. Therefore, ``offset_x = image_width / 2`` places the
    center of the image at ``center_x``. This parameter is therefore analogous
    to the WCS keyword ``CRVAL1``.

    For tiled sky images, the offset is measured in *degrees*, and a value of
    zero means that the *center* of the image lands at the center of the
    projection coordinate system.

    """
    offset_y = Float(0.0).tag(xml=XmlSer.attr('OffsetY'))
    """The vertical positioning of the image relative to its projection
    coordinate system.

    For untiled sky images, the image is by default positioned such that its
    lower left lands at the center of the projection coordinate system (namely,
    ``center_x`` and ``center_y``). The offset is measured in pixels and moves
    the image downwards. Therefore, ``offset_y = image_height / 2`` places the
    center of the image at ``center_y``. This parameter is therefore analogous
    to the WCS keyword ``CRVAL2``.

    For tiled sky images, the offset is measured in *degrees*, and a value of
    zero means that the *center* of the image lands at the center of the
    projection coordinate system.

    """
    rotation_deg = Float(0.0).tag(xml=XmlSer.attr('Rotation'))
    """The rotation of image’s projection coordinate system, in degrees.

    For sky images, this is East from North, i.e. counterclockwise.

    """
    band_pass = UseEnum(
        Bandpass,
        default_value = Bandpass.VISIBLE
    ).tag(xml=XmlSer.attr('BandPass'))
    """The bandpass of the image data."""

    sparse = Bool(True).tag(xml=XmlSer.attr('Sparse'))
    """TBD."""

    elevation_model = Bool(False).tag(xml=XmlSer.attr('ElevationModel'))
    """TBD."""

    stock_set = Bool(False).tag(xml=XmlSer.attr('StockSet'))
    """TBD."""

    generic = Bool(False).tag(xml=XmlSer.attr('Generic'))
    """TBD."""

    mean_radius = Float(0.0).tag(xml=XmlSer.attr('MeanRadius'))
    """TBD."""

    credits = Unicode('').tag(xml=XmlSer.text_elem('Credits'))
    """Textual credits for the image originator."""

    credits_url = Unicode('').tag(xml=XmlSer.text_elem('CreditsUrl'))
    """A URL giving the source of the image or more information about its creation."""

    thumbnail_url = Unicode('').tag(xml=XmlSer.text_elem('ThumbnailUrl'))
    """A URL to a standard WWT thumbnail representation of this imageset."""

    description = Unicode('').tag(xml=XmlSer.text_elem('Description'))
    """A textual description of the imagery."""

    msr_community_id = Int(0).tag(xml=XmlSer.attr('MSRCommunityId'))
    """The ID number of the WWT Community that this content came from."""

    msr_component_id = Int(0).tag(xml=XmlSer.attr('MSRComponentId'))
    """The ID number of this content item on the WWT Communities system."""

    permission = Int(0).tag(xml=XmlSer.attr('Permission'))
    "TBD."

    def _tag_name(self):
        return 'ImageSet'

    def mutate_urls(self, mutator):
        if self.url:
            self.url = mutator(self.url)

        if self.dem_url:
            self.dem_url = mutator(self.dem_url)

        if self.credits_url:
            self.credits_url = mutator(self.credits_url)

        if self.thumbnail_url:
            self.thumbnail_url = mutator(self.thumbnail_url)

    def set_position_from_wcs(self, headers, width, height, place=None, fov_factor=1.7):
        """Set the positional information associated with this imageset to match a set
        of WCS headers.

        Parameters
        ----------
        headers : :class:`~astropy.io.fits.Header` or string-keyed dict-like
          A set of FITS-like headers including WCS keywords such as ``CRVAL1``.
        width : positive integer
          The width of the image associated with the WCS, in pixels.
        height : positive integer
          The height of the image associated with the WCS, in pixels.
        place : optional :class:`~wwt_data_formats.place.Place`
          If specified, the centering and zoom level of the :class:`~wwt_data_formats.place.Place`
          object will be set to match the center and size of this image.
        fov_factor : optional float
          If *place* is provided, its zoom level will be set so that the
          angular height of the client viewport is this factor times the
          angular height of the image. The default is 1.7.

        Returns
        -------
        self
          For convenience in chaining function calls.

        Notes
        -----
        Certain of the ImageSet parameters take on different meanings depending on
        whether the image in question is a tiled "study" or not. This method will alter
        its behavior depending on whether the :attr:`tile_levels` attribute is greater
        than zero. If you are computing coordinates for a tiled study, make sure to set
        this parameter *before* calling this function.

        For the time being, the WCS must be equatorial using the gnomonic
        (``TAN``) projection.

        Required keywords in *headers* are:

        - ``CTYPE1`` and ``CTYPE2``
        - ``CRVAL1`` and ``CRVAL2``
        - ``CRPIX1`` and ``CRPIX2``
        - Either:
          - ``CDELT1``, ``CDELT2``, ``PC1_1``, and ``PC1_2``; or
          - ``CD1_1``, ``CD2_2``

        If present ``PC1_2``, ``PC2_1``, ``CD1_2``, and/or ``CD2_1`` are used.
        If absent, they are assumed to be zero.

        """
        if headers['CTYPE1'] != 'RA---TAN' or headers['CTYPE2'] != 'DEC--TAN':
            raise ValueError('WCS coordinates must be in an equatorial/TAN projection')

        # Figure out the stuff we need from the headers.

        ra_deg = headers['CRVAL1']
        dec_deg = headers['CRVAL2']
        crpix_x = headers['CRPIX1'] - 1
        crpix_y = headers['CRPIX2'] - 1

        if 'CD1_1' in headers:
            cd1_1 = headers['CD1_1']
            cd2_2 = headers['CD2_2']
            cd1_2 = headers.get('CD1_2', 0.0)
            cd2_1 = headers.get('CD2_1', 0.0)

            if cd1_1 * cd2_2 - cd1_2 * cd2_1 < 0:
                cd_sign = -1
            else:
                cd_sign = 1

            rot_rad = math.atan2(-cd_sign * cd1_2, cd2_2)
            scale_x = math.sqrt(cd1_1**2 + cd2_1**2) * cd_sign
            scale_y = math.sqrt(cd1_2**2 + cd2_2**2)
        else:
            scale_x = headers['CDELT1']
            scale_y = headers['CDELT2']
            pc1_1 = headers.get('PC1_1', 1.0)
            pc2_2 = headers.get('PC2_2', 1.0)
            pc1_2 = headers.get('PC1_2', 0.0)
            pc2_1 = headers.get('PC2_1', 0.0)

            det = pc1_1 * pc2_2 - pc1_2 * pc2_1

            if det < 0:
                pc_sign = -1
            else:
                pc_sign = 1

            rot_rad = math.atan2(-pc_sign * pc1_2, pc2_2)

            # I am not sure if this is "supposed" to be allowed, but I've seen it.
            rtdet = math.sqrt(pc_sign * det)
            scale_x *= rtdet
            scale_y *= rtdet

        # This is our best effort to make sure that the view centers on the
        # center of the image.

        try:
            from astropy.wcs import WCS
        except:
            center_ra_deg = ra_deg
            center_dec_deg = dec_deg
        else:
            wcs = WCS(headers)
            center = wcs.pixel_to_world(height / 2, width / 2)
            center_ra_deg = center.ra.deg
            center_dec_deg = center.dec.deg

        # Now, assign the fields

        self.data_set_type = DataSetType.SKY
        self.width_factor = 2
        self.center_x = ra_deg
        self.center_y = dec_deg
        self.rotation_deg = rot_rad * 180 / math.pi

        if self.tile_levels > 0:  # are we tiled?
            self.projection = ProjectionType.TAN
            self.offset_x = (width / 2 - crpix_x) * abs(scale_x)
            self.offset_y = (height / 2 - crpix_y) * scale_y
            self.base_degrees_per_tile = scale_y * 256 * 2**self.tile_levels
        else:
            self.projection = ProjectionType.SKY_IMAGE
            self.offset_x = crpix_x
            self.offset_y = crpix_y
            self.base_degrees_per_tile = scale_y

        if place is not None:
            place.data_set_type = DataSetType.SKY
            place.rotation_deg = 0.  # I think this is better than propagating the image rotation?
            place.ra_hr = center_ra_deg / 15.
            place.dec_deg = center_dec_deg
            # It is hardcoded that in sky mode, zoom_level = height of client FOV * 6.
            place.zoom_level = height * scale_y * fov_factor * 6

        return self

    def wcs_headers_from_position(self):
        """Compute a set of WCS headers for this ImageSet's positional information.

        Returns
        -------
        A string-keyed dict-like containing FITS/WCS header keywords such as
        ``CTYPE1``, ``CRPIX1``, etc.

        Notes
        -----
        At the moment, this function only works for ImageSets with a
        projection type of ``SKY_IMAGE``. Support for other projections
        *might* be added later, if the need arises..

        """
        rv = {
            'CTYPE1': 'RA---TAN',
            'CTYPE2': 'DEC--TAN',
            'CRVAL1': self.center_x,
            'CRVAL2': self.center_y,
        }

        if self.projection != ProjectionType.SKY_IMAGE:
            raise NotImplementError('wcs_headers_from_position() only works if projection=SKY_IMAGE')

        rv['CRPIX1'] = self.offset_x + 1
        rv['CRPIX2'] = self.offset_y + 1
        rv['CDELT2'] = self.base_degrees_per_tile  # = scale_y, above
        rv['CDELT1'] = -self.base_degrees_per_tile  # AFAICT, non-square pixels can't be expressed

        c = math.cos(self.rotation_deg * math.pi / 180)
        s = math.sin(self.rotation_deg * math.pi / 180)

        rv['PC1_1'] = c
        rv['PC1_2'] = -s
        rv['PC2_1'] = s
        rv['PC2_2'] = c

        return rv


#!/usr/bin/env python

import ROOT
ROOT.PyConfig.IgnoreCommandLineOptions = True  # disable ROOT internal argument parser

import logging
logger = logging.getLogger("annotate_file_inc_keras.py")
logger.setLevel(logging.DEBUG)
formatter = logging.Formatter("%(name)s - %(levelname)s - %(message)s")
handler = logging.StreamHandler()
handler.setFormatter(formatter)
logger.addHandler(handler)

import numpy as np
import yaml
import os
import pickle
from array import array
import argparse
from sklearn.preprocessing import StandardScaler

from keras.models import load_model

def parse_arguments():
    parser = argparse.ArgumentParser(
        description="Apply XGB model on ROOT file")
    parser.add_argument(
        "--config-training",
        default="mt_xgb_training_config.yaml",
        help="Path to training config file")
    parser.add_argument(
        "--dir-prefix",
        type=str,
        default="ntuple",
        help="Prefix of directories in ROOT file to be annotated.")
    parser.add_argument(
        "input", help="Path to input file, where response will be added.")
    parser.add_argument(
        "tag", help="Tag to be used as prefix of the annotation.")
    parser.add_argument(
        "--models",
        type=str,
        nargs="+",
        default=["multi_fold1_cpsm_mt_JHU_xgb.pkl", "multi_fold0_cpsm_mt_JHU_xgb.pkl"],
        help=
        "Keras models to be used for the annotation. Note that these have to be booked in the reversed order [fold1*, fold0*], so that the training is independent from the application."
    )
    parser.add_argument(
        "--preprocessing",
        type=str,
        nargs="+",
        default=[
            "tt_training_9May_tests/tt_fold1_scaler.pkl",
            "tt_training_9May_tests/tt_fold0_scaler.pkl"
        ],
        help=
        "Data preprocessing to be used. Note that these have to be booked in the reversed order [fold1*, fold0*], so that the preprocessing is independent for the folds."
    )
    parser.add_argument(
        "--tree", default="ntuple", help="Name of trees in the directories.")
    parser.add_argument(
        "--training", default="JHU", help="Name of training to use.")
    parser.add_argument(
        "--mjj", default="high", help="mjj training to use.")
    parser.add_argument(
        "--channel", default="mt", help="Name of channel to annotate.")
    parser.add_argument(
        "--model_folder", default="mt_training_10May_mjj_jdeta_dijetpt/", help="Folder name where trained model is.")
    parser.add_argument(
        "--era", default="", help="Year to use.")

    return parser.parse_args()

def parse_config(filename):
    return yaml.load(open(filename, "r"))

def load_files(filelist):

    with open(filelist) as f:
        file_names = f.read().splitlines()
        # file_names = [os.path.splitext(os.path.basename(file))[0] for file in files]

    return file_names

def main(args, config, file_names):

    # path = "/vols/cms/akd116/Offline/output/SM/2018/Apr24_1"
    # path = "/vols/cms/akd116/Offline/output/SM/2018/May17_2"
    # path = "/vols/cms/akd116/Offline/output/SM/2018/Jun22_2016_Danny/"
    # path = "/vols/cms/akd116/Offline/output/SM/2018/Aug14_2016_Danny_new/"
    # path = "/vols/cms/akd116/Offline/output/SM/2018/Aug14_2016_Danny_v3/"
    # path = "/vols/cms/akd116/Offline/output/SM/2018/Sep01_2016_Danny/"
    # path = "/vols/cms/akd116/Offline/output/SM/2018/Sep21/"
    # path = "/vols/cms/akd116/Offline/output/SM/2018/Nov27_2017_copy/"
    # path = "/vols/cms/akd116/Offline/output/SM/2019/Feb26_2016/"
    path = "/vols/cms/akd116/Offline/output/SM/2019/CPdecay_Apr26_2/"

    # Sanity checks
    for sample in file_names:
        print sample
        # if not os.path.exists("{}/{}_{}_{}.root".format(path, sample, args.channel, args.era)):
        if not os.path.exists("{}/{}".format(path, sample)):
            logger.fatal("Input file %s does not exist.", sample)
            raise Exception

        logger.debug("Following mapping of classes to class numbers is used.")
        for i, class_ in enumerate(config["classes"]):
            logger.debug("%s : %s", i, class_)

        # Load Keras models and preprocessing

        if args.era != "":
            # with open('{}/{}_fold1_keras_model.h5'
            #         .format(args.model_folder, args.channel), 'r') as f:
            #     clf_fold1 = load_model(f)
            # with open('{}/{}_fold0_keras_model.h5'
            #         .format(args.model_folder, args.channel), 'r') as f:
            #     clf_fold0 = load_model(f)

            # vienna/KIT settings
            # clf_fold1 = load_model('{}/{}_fold1_keras_model.h5'
            #         .format(args.model_folder, args.channel))
            # clf_fold0 = load_model('{}/{}_fold0_keras_model.h5'
            #         .format(args.model_folder, args.channel))

            # IC settings
            clf_fold1 = load_model('{}/keras_model_fold1_sm_{}_powheg.h5'
                    .format(args.model_folder, args.channel))
            clf_fold0 = load_model('{}/keras_model_fold0_sm_{}_powheg.h5'
                    .format(args.model_folder, args.channel))

            # with open('{}/multi_fold1_cpsm_{}_{}_{}_xgb.pkl'
            #         .format(args.model_folder, args.channel, args.training, args.era), 'r') as f:
            #     xgb_clf_fold1 = pickle.load(f)
            # with open('{}/multi_fold0_cpsm_{}_{}_{}_xgb.pkl'
            #         .format(args.model_folder, args.channel, args.training, args.era), 'r') as f:
            #     xgb_clf_fold0 = pickle.load(f)
        # elif args.training == "madgraph" or args.training == "powheg":
            # with open('{}/multi_fold1_cpsm_{}_{}_{}_xgb.pkl'
            #         .format(args.model_folder, args.channel, args.training, args.mjj), 'r') as f:
            #     xgb_clf_fold1 = pickle.load(f)
            # with open('{}/multi_fold0_cpsm_{}_{}_{}_xgb.pkl'
            #         .format(args.model_folder, args.channel, args.training, args.mjj), 'r') as f:
            #     xgb_clf_fold0 = pickle.load(f)
        # else:
            # with open('{}/multi_fold1_cpsm_{}_{}_xgb.pkl'
            #         .format(args.model_folder, args.channel, args.training), 'r') as f:
            #     xgb_clf_fold1 = pickle.load(f)
            # with open('{}/multi_fold0_cpsm_{}_{}_xgb.pkl'
            #         .format(args.model_folder, args.channel, args.training), 'r') as f:
            #     xgb_clf_fold0 = pickle.load(f)

        classifier = [clf_fold1, clf_fold0] # for IC training
        # classifier = [clf_fold0, clf_fold1] # for Vienna training

        # vienna/KIT settings
        # with open('{}/{}_fold1_keras_preprocessing.pickle'
        #         .format(args.model_folder, args.channel), 'r') as f:
        #     preprocessing_fold1 = pickle.load(f)
        # with open('{}/{}_fold0_keras_preprocessing.pickle'
        #         .format(args.model_folder, args.channel), 'r') as f:
        #     preprocessing_fold0 = pickle.load(f)
        # preprocessing = [preprocessing_fold0, preprocessing_fold1] # vienna order
        # preprocessing = [pickle.load(open(x, "rb")) for x in args.preprocessing]

        with open('{}/{}_scaler.pkl'
                .format(args.model_folder, args.channel), 'r') as f:
            preprocessing = pickle.load(f)

        # Open input file
        file_ = ROOT.TFile("{}/{}".format(path, sample), "UPDATE")
        if file_ == None:
            logger.fatal("File %s is not existent.", sample)
            raise Exception

        # Loop through directories in this file and annotate tree if directory
        # starts with the set prefix.
        # for key in file_.GetListOfKeys():
            # Find valid directories
            # name = key.GetName()
            # if name.startswith(args.dir_prefix):
            #     logger.debug("Process directory %s.", name)
        tree = file_.Get(args.tree)
        # if tree == None:
        #     logger.fatal("Failed to find tree %s in directory %s.",
        #                  args.tree, name)
        #     raise Exception

        # Book branches for annotation
        values = []
        for variable in config["variables"]:
            if variable in ["dijetpt","eta_h","IC_binary_test_4_score","IC_binary_test_4_index","bpt_1","bpt_2"]:
                values.append(array("f", [-9999]))
            if variable in ["eta_1","eta_2","jdeta","jpt_1","jpt_2","m_sv","m_vis","met","jeta_1","jeta_2","mt_tot","mt_sv",
                    "met_dphi_1","met_dphi_2","mjj","mt_1","mt_2","mt_lep","pt_1","pt_2","pt_h","pt_tt","pt_vis","pzeta","dR"]:
                values.append(array("d", [-9999]))
            if variable in ["n_jets","n_bjets","opp_sides"]:
                values.append(array("I", [0]))
            if variable not in ["zfeld","centrality","mjj_jdeta","dijetpt_pth","dijetpt_jpt1","dijetpt_pth_over_pt1",
                    "msv_mvis","msvsq_mvis","msv_sq","log_metsq_jeta2","met_jeta2","oppsides_centrality","pthsq_ptvis","msv_rec","dR_custom","rms_pt","rms_jpt","rec_sqrt_msv"]:
                tree.SetBranchAddress(variable, values[-1])
            # else:
            #     tree.SetBranchAddress("eta_h", values[-1])

        response_max_score = array("f", [-9999])
        branch_max_score = tree.Branch("{}_max_score".format(
            args.tag), response_max_score, "{}_max_score/F".format(
                args.tag))

        response_max_index = array("f", [-9999])
        branch_max_index = tree.Branch("{}_max_index".format(
            args.tag), response_max_index, "{}_max_index/F".format(
                args.tag))

        response_ggh_score = array("f", [-9999])
        branch_ggh_score = tree.Branch("{}_ggh_score".format(
            args.tag), response_ggh_score, "{}_ggh_score/F".format(
                args.tag))

        # response_qqh_score = array("f", [-9999])
        # branch_qqh_score = tree.Branch("{}_qqh_score".format(
        #     args.tag), response_qqh_score, "{}_qqh_score/F".format(
        #         args.tag))

        # Run the event loop
        for i_event in range(tree.GetEntries()):
            tree.GetEntry(i_event)

            # Get event number and compute response
            event = int(getattr(tree, "event"))
            m_sv = float(getattr(tree, "m_sv"))
            # m_vis = float(getattr(tree, "m_vis"))

            response_max_score[0] = -9999.0
            response_max_index[0] = -9999.0
            response_ggh_score[0] = -9999.0
            # response_qqh_score[0] = -9999.0

            # if m_sv > 0:
            # zfeld = float(getattr(tree,"eta_h")) - (float(getattr(tree,"jeta_1"))+float(getattr(tree,"jeta_2")))/2
            # centrality = np.exp(-4*(zfeld/np.fabs(float(getattr(tree,"jdeta"))))**2)
            # dphi_custom = np.arccos(1-float(getattr(tree,"mt_lep"))**2/(2.*float(getattr(tree,"pt_1"))*float(getattr(tree,"pt_2"))))
            # dR_custom = np.sqrt((float(getattr(tree,"eta_1"))-float(getattr(tree,"eta_2")))**2 + dphi_custom**2)

            # # additional_vars = [
            # #         zfeld,centrality,mjj_jdeta,dijetpt_pth,dijetpt_jpt1,dijetpt_pth_over_pt1,
            # #         msv_mvis,msvsq_mvis,msv_sq,log_metsq_jeta2,met_jeta2,oppsides_centrality,pthsq_ptvis,msv_rec
            # #         ]
            # additional_vars = [
            #         centrality,#mjj_jdeta,
            #         #msvsq_mvis,pthsq_ptvis,
            #         dR_custom,
            #         ]
            # if len(values) < len(config["variables"]):
            #     values.extend(additional_vars) ## for the first loop
            #     # values.append(centrality) ## for the first loop
            # else:
            #     for index,val in enumerate(additional_vars):
            #         ind = len(additional_vars)-index
            #         values[-ind] = val
            #     # values[-1] = centrality # then replace the last value
            # values = [-999. if x == -9999. else x for x in values] # maybe need this?
            values_stacked = np.hstack(values).reshape(1, len(values))
            values_preprocessed = preprocessing.transform(values_stacked)
            response = classifier[event % 2].predict(values_preprocessed)
            # response = classifier[event % 2].predict_proba(values_stacked,
            #         ntree_limit=classifier[event % 2].best_iteration+1)
            response = np.squeeze(response)

            # Find max score and index
            response_max_score[0] = -9999.0
            for i, r in enumerate(response):
                if r > response_max_score[0]:
                    response_max_score[0] = r
                    response_max_index[0] = i

            if i_event % 10000 == 0:
                logger.debug('Currently on event {}'.format(i_event))

            # rms = np.sqrt(1./3. * (response[1]**2+response[2]**2+response[3]**2))
            # if response_max_index[0] == 0 and (response[1]>0.3 or response[2]>0.3 or response[3]>0.3):
            #     response_max_score[0] += rms
            # print response
            # print response_max_score[0]
            # exit()

            # Take ggH score as well
            response_ggh_score[0] = -9999.0
            # if args.channel in ['tt','em']:
            response_ggh_score[0] = response[0]
            # elif args.channel in ['mt','et']:
            #     response_ggh_score[0] = response[1]

            # Take qqH score
            # response_qqh_score[0] = -9999.0
            # if args.channel in ['tt','em']:
            # response_qqh_score[0] = response[2]
            # elif args.channel in ['mt','et']:
            #     response_qqh_score[0] = response[1]

            del response

            # Fill branches
            branch_max_score.Fill()
            branch_max_index.Fill()
            branch_ggh_score.Fill()
            # branch_qqh_score.Fill()

        logger.debug("Finished looping over events")
        # Write everything to file
        file_.Write("ntuple",ROOT.TObject.kWriteDelete)
        file_.Close()
        logger.debug("Closed file")

if __name__ == "__main__":
    args = parse_arguments()
    config = parse_config(args.config_training)
    file_names = load_files(args.input)
    main(args, config, file_names)


#  Version: 2020.02.21
#
#  MIT License
#
#  Copyright (c) 2018 Jiankang Deng and Jia Guo
#
#  Permission is hereby granted, free of charge, to any person obtaining a copy
#  of this software and associated documentation files (the "Software"), to deal
#  in the Software without restriction, including without limitation the rights
#  to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
#  copies of the Software, and to permit persons to whom the Software is
#  furnished to do so, subject to the following conditions:
#
#  The above copyright notice and this permission notice shall be included in all
#  copies or substantial portions of the Software.
#
#  THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
#  IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
#  FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
#  AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
#  LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
#  OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
#  SOFTWARE.
#

# pylint: skip-file
import mxnet as mx
import numpy as np
import os
import random
import glob
import cv2
from mxnet.io import DataIter
from mxnet import ndarray as nd
from config import config

class FaceSegIter(DataIter):
    def __init__(self, path, batch_size,
                 per_batch_size = 0,
                 aug_level = 0,
                 force_mirror = False,
                 exf = 1,
                 args = None):
      self.aug_level = aug_level
      self.force_mirror = force_mirror
      self.exf = exf
      self.batch_size = batch_size
      self.per_batch_size = per_batch_size
      self.image_file_list = []
      self.uv_file_list = []
      for _file in glob.glob(os.path.join(path, '*.jpg')):
        self.image_file_list.append(_file)
      for img in self.image_file_list:
        uv_file = img[0:-3]+"npy"
        self.uv_file_list.append(uv_file)
      self.seq = range(len(self.image_file_list))
      print('train size', len(self.seq))
      self.cur = 0
      self.reset()
      self.data_shape = (3, config.input_img_size, config.input_img_size)
      self.num_classes = config.num_classes
      self.input_img_size = config.input_img_size
      #self.label_classes = self.num_classes
      self.output_label_size = config.output_label_size
      #if aug_level>0:
      #  self.output_label_size = config.output_label_size
      #else:
      #  self.output_label_size = self.input_img_size
      self.label_shape = (self.num_classes, self.output_label_size, self.output_label_size)
      self.provide_data = [('data', (batch_size,) + self.data_shape)]
      self.provide_label = [('softmax_label', (batch_size,) + self.label_shape),
                            ('mask_label', (batch_size,)+ self.label_shape)]
      weight_mask = cv2.imread('./uv-data/uv_weight_mask.png')
      print('weight_mask', weight_mask.shape)
      if weight_mask.shape[0]!=self.output_label_size:
        weight_mask = cv2.resize(weight_mask, (self.output_label_size, self.output_label_size) )
      #idx = np.where(weight_mask>0)[0]
      #print('weight idx', idx)
      weight_mask = weight_mask.astype(np.float32)
      weight_mask /= 255.0

      vis_mask = cv2.imread('./uv-data/uv_face_mask.png')
      print('vis_mask', vis_mask.shape)
      if vis_mask.shape[0]!=self.output_label_size:
        vis_mask = cv2.resize(vis_mask, (self.output_label_size, self.output_label_size) )
      vis_mask = vis_mask.astype(np.float32)
      vis_mask /= 255.0
      weight_mask *= vis_mask
      print('weight_mask', weight_mask.shape)
      weight_mask = weight_mask.transpose( (2,0,1) )
      #WM = np.zeros( (batch_size,)+self.label_shape, dtype=np.float32 )
      #for i in range(batch_size):
      #  WM[i] = weight_mask
      #weight_mask = WM
      #weight_mask = weight_mask.reshape( (1, 3, weight_mask.shape[0], weight_mask.shape[1]) )
      weight_mask = weight_mask[np.newaxis,:,:,:]
      print('weight_mask', weight_mask.shape)
      weight_mask = np.tile(weight_mask, (batch_size,1,1,1))
      print('weight_mask', weight_mask.shape)
      self.weight_mask = nd.array(weight_mask)
      self.img_num = 0
      self.invalid_num = 0
      self.mode = 1
      self.vis = 0
      self.stats = [0,0]

    def get_data_shape(self):
        return self.data_shape

    #def get_label_shape(self):
    #    return self.label_shape

    def get_shape_dict(self):
        D = {}
        for (k,v) in self.provide_data:
            D[k] = v
        for (k,v) in self.provide_label:
            D[k] = v
        return D

    def get_label_names(self):
        D = []
        for (k,v) in self.provide_label:
            D.append(k)
        return D

    def reset(self):
      #print('reset')
      self.cur = 0
      if self.aug_level>0:
        random.shuffle(self.seq)

    def next_sample(self):
      """Helper function for reading in next sample."""
      if self.cur >= len(self.seq):
        raise StopIteration
      idx = self.seq[self.cur]
      self.cur += 1
      uv_path = self.uv_file_list[idx]
      image_path = self.image_file_list[idx]
      uvmap = np.load(uv_path)
      img = cv2.imread(image_path)[:,:,::-1]#to rgb
      hlabel = uvmap
      #print(hlabel.shape)
      #hlabel = np.array(header.label).reshape( (self.output_label_size, self.output_label_size, self.num_classes) )
      hlabel /= self.input_img_size

      return img, hlabel

    def next(self):
        """Returns the next batch of data."""
        #print('next')
        batch_size = self.batch_size
        batch_data = nd.empty((batch_size,)+self.data_shape)
        batch_label = nd.empty((batch_size,)+self.label_shape)
        i = 0
        #self.cutoff = random.randint(800,1280)
        try:
            while i < batch_size:
                #print('N', i)
                data, label = self.next_sample()
                data = nd.array(data)
                data = nd.transpose(data, axes=(2, 0, 1))
                label = nd.array(label)
                label = nd.transpose(label, axes=(2, 0, 1))
                batch_data[i][:] = data
                batch_label[i][:] = label
                i += 1
        except StopIteration:
            if i<batch_size:
                raise StopIteration

        #return {self.data_name  :  batch_data,
        #        self.label_name :  batch_label}
        #print(batch_data.shape, batch_label.shape)
        return mx.io.DataBatch([batch_data], [batch_label, self.weight_mask], batch_size - i)


from utils import extract_column
from sklearn.linear_model import LogisticRegression
from sklearn.naive_bayes import BernoulliNB
from sklearn.feature_extraction import DictVectorizer
from sklearn.svm import LinearSVC
from sklearn_crfsuite import CRF
import pandas as pd
import numpy as np
import gensim.downloader as api

class TextClassifier():
    """
    Class that can be trained to different classifier systems based on different features
    """

    def __init__(self, modelname):

        self.modelname = modelname
        self.model = None
        self.vectorizer = None
        self.features_names_list = []

    def train(self, inputfile, features_names_list, annotation_column):
        """
        This function fits a classification model as specified on training data

        :param inputfile: path to inputfile containing the training data
        :param features_names_list: list of indications of all feature columns that should be used
        :param annotation_column: indication of column with annotations
        :type inputfile: string
        :type features_names_list: list
        :type annotation_column: string
        """

        # initialize the right model
        if self.modelname == 'logreg':
            self.model = LogisticRegression()

        elif self.modelname == 'naivebayes':
            self.model = BernoulliNB()

        elif self.modelname == 'svm':
            self.model = LinearSVC()

        elif self.modelname == 'crf':
            self.model = CRF(algorithm='lbfgs', c1=0.1, c2=0.1, max_iterations=100, all_possible_transitions=True)

        # store features_names_list as class attribute
        self.features_names_list = features_names_list

        # get training features and labels
        train_features = self.get_features(inputfile)
        train_targets = self.get_labels(inputfile, annotation_column)

        # fit the model
        self.model.fit(train_features, train_targets)

    def predict(self, inputfile):
        """
        This function classifies the observations of the given inputfile

        :param inputfile: path to inputfile
        :type inputfile: string

        :returns: list of predictions
        """
        # get features and predict the labels
        features = self.get_features(inputfile)
        predictions = self.model.predict(features)

        # join the predictions for every sentence if the model is CRF
        if self.modelname == 'crf':
            joint_preds = [pred for sent in predictions for pred in sent]
            return joint_preds

        return predictions

    def get_features(self, inputfile):
        """
        This function extracts all features transformed to vector from the inputfile according to the model and feature settings

        :param inputfile: path to inputfile
        :type inputfile: string

        :returns: vectors that represent the features
        """
        # extract and return features if the model is a CRF
        if self.modelname == 'crf':
            return self.extract_crf_features(inputfile)

        # extract all features as one-hot representations except for embeddings (if specified)
        sparse_features = None
        sparse_features_list = [feature for feature in self.features_names_list if feature != 'embedding']
        if sparse_features_list:
            sparse_features = self.extract_features(inputfile, sparse_features_list)

        # extract word embedding representations if specified in features list
        if 'embedding' in self.features_names_list:
            embeddings = self.extract_embeddings_as_features(inputfile)

            # combine embedding with sparse features if both should be represented, and return them
            if sparse_features != None:
                combined_features = self.combine_sparse_and_dense_features(embeddings, sparse_features)
                return combined_features
            else:
                return embeddings

        return sparse_features

    def get_labels(self, inputfile, annotation_column):
        """
        This function extracts all the labels from the inputfile in the form required by the model

        :param inputfile: path to inputfile
        :param annotation_column: indication of column with annotations
        :type inputfile: string
        :type annotation_column: string

        :returns: labels in required form
        """

        # extract the labels in the form required for CRF and return
        if self.modelname == 'crf':
            return self.extract_crf_labels(inputfile, annotation_column)

        # extract and return the labels in the form required for all other models
        return extract_column(inputfile, annotation_column)

    def extract_features(self, inputfile, features):
        """
        This function extracts the specified features and transforms them to vectors

        :param inputfile: path to tab seperated inputfile with header
        :param features: list of feature names that should be extracted
        :type inputfile: string
        :type features: list

        :returns: feature vectors
        """
        # reads the inputfile
        data = pd.read_csv(inputfile, sep='\t', keep_default_na=False)
        # converts the specified feature columns to a list of dictionaries
        features = data[features].to_dict('records')

        # transforms the feature dictionaries to vectors, and stores the fitted vectorizer
        if not self.vectorizer:
            self.vectorizer = DictVectorizer()
            self.vectorizer.fit(features)

        features = self.vectorizer.transform(features)

        return features

    def extract_embeddings_as_features(self, inputfile, embedding_model='word2vec-google-news-300'):
        '''
        Function that extracts features and gold labels using word embeddings

        :param inputfile: path to inputfile
        :param token_column: header name of column with tokens
        :param embedding_model: name of a pretrained word embedding model
        :type inputfile: string
        :type token_column: string
        :type embedding_model: string

        :return features: list of vector representation of tokens
        '''
        ### This code was partially inspired by code included in the HLT course, obtained from https://github.com/cltl/ma-hlt-labs/, accessed in May 2020.
        features = []
        embedding_model = api.load(embedding_model)
        tokens = extract_column(inputfile, 'token')

        for token in tokens:
            if token in embedding_model:
                vector = embedding_model[token]
            else:
                vector = [0]*300
            features.append(vector)

        return features

    def combine_sparse_and_dense_features(self, dense_vectors, sparse_features):
        '''
        Function that takes sparse and dense feature representations and appends their vector representation

        :param dense_vectors: list of dense vector representations
        :param sparse_features: list of sparse vector representations
        :type dense_vector: list of arrays
        :type sparse_features: list of lists

        :returns: list of arrays in which sparse and dense vectors are concatenated
        '''
        ### SOURCE: code included in the HLT course, obtained from https://github.com/cltl/ma-hlt-labs/.

        combined_vectors = []
        sparse_vectors = np.array(sparse_features.toarray())

        for index, vector in enumerate(sparse_vectors):
            combined_vector = np.concatenate((vector, dense_vectors[index]))
            combined_vectors.append(combined_vector)

        return combined_vectors

    def extract_crf_features(self, inputfile):
        """
        This function extracts the features for every token grouped for every sentence

        :param inputfile: path to inputfile
        :type inputfile: string

        :returns: a list with for every sentence a list of feature dictionaries
        """
        ### based on: https://github.com/TeamHG-Memex/sklearn-crfsuite/blob/master/docs/CoNLL2002.ipynb

        # reads the inputfile
        data = pd.read_csv(inputfile, sep='\t', keep_default_na=False)
        # converts the specified feature columns to a list of dictionaries
        features = data[self.features_names_list].to_dict('records')

        # appends for every sentence ('.' as sentence boundary) the feature dictionaries to a list
        # appends that list if complete to a list for all sentences
        sentences = []
        sentence = []

        for token_dict in features:
            sentence.append(token_dict)
            if token_dict['token'] == '.':
                sentences.append(sentence)
                sentence = []
        sentences.append(sentence)

        return sentences

    def extract_crf_labels(self, inputfile, annotation_column):
        """
        This function extracts the label for every token grouped for every sentence

        :param inputfile: path to inputfile
        :type inputfile: string

        :returns: a list with for every sentence a list of labels
        """
        ### based on: https://github.com/TeamHG-Memex/sklearn-crfsuite/blob/master/docs/CoNLL2002.ipynb

        # reads the inputfile
        data = pd.read_csv(inputfile, sep='\t', keep_default_na=False)
        # converts the specified feature columns to a list of dictionaries
        data_dicts = data.to_dict('records')

        # appends for every sentence ('.' as sentence boundary) the labels to a list
        # appends that list if complete to a list for all sentences
        sentences_labels = []
        sentence_labels = []

        for token_dict in data_dicts:
            sentence_labels.append(token_dict[annotation_column])
            if token_dict['token'] == '.':
                sentences_labels.append(sentence_labels)
                sentence_labels = []
        sentences_labels.append(sentence_labels)

        return sentences_labels