# --------------------------------------------------------
# Deformable Convolutional Networks
# Copyright (c) 2017 Microsoft
# Licensed under The Apache-2.0 License [see LICENSE for details]
# Written by Yi Li, Haocheng Zhang
# --------------------------------------------------------
import os
import sys
import numpy as np
import argparse
import logging
import pprint
import cv2

import rospkg

# TODO: When the code in setup.py starts working, then make sure to remove this
# paths hack as well
def add_path(path):
    if path not in sys.path:
        sys.path.insert(0, path)
rospack = rospkg.RosPack()
LIB_PATH = os.path.join(rospack.get_path('rail_object_detector'), 'libs', 'drfcn')
add_path(LIB_PATH)

# Start importing mxnet
from utils.image import resize, transform
from cfg.drfcn_config import config, update_config

# get config
os.environ['PYTHONUNBUFFERED'] = '1'
os.environ['MXNET_CUDNN_AUTOTUNE_DEFAULT'] = '0'
os.environ['MXNET_ENABLE_GPU_P2P'] = '0'
update_config(os.path.join(LIB_PATH, 'cfg', 'rfcn_coco_demo.yaml'))

# actually import mxnet
import mxnet as mx
from core.tester import im_detect, Predictor
from symbols import *
from utils.load_model import load_param
from utils.show_boxes import show_boxes
from utils.tictoc import tic, toc
from nms.nms import py_nms_wrapper, cpu_nms_wrapper, gpu_nms_wrapper

class Detector:
    def __init__(
        self,
        model_path=os.path.join(LIB_PATH, 'model', 'rfcn_dcn_coco'),
        hd_images=False
    ):
        # IF we are guaranteed QHD / HD, we can remove scaling and resizing and be a bit faster
        # if hd_images:
        #     config.SCALES = [(1080, 1820)]
        # else:
        #     config.SCALES = [(540, 960)]
        pprint.pprint(config)
        config.symbol = 'resnet_v1_101_rfcn_dcn'
        sym_instance = eval(config.symbol + '.' + config.symbol)()
        sym = sym_instance.get_symbol(config, is_train=False)

        # set up class names
        self.num_classes = 81
        self.classes = ['person', 'bicycle', 'car', 'motorcycle', 'airplane', 'bus', 'train', 'truck', 'boat',
                   'traffic light', 'fire hydrant', 'stop sign', 'parking meter', 'bench', 'bird', 'cat', 'dog', 'horse',
                   'sheep', 'cow', 'elephant', 'bear', 'zebra', 'giraffe', 'backpack', 'umbrella', 'handbag', 'tie',
                   'suitcase', 'frisbee', 'skis', 'snowboard', 'sports ball', 'kite', 'baseball bat', 'baseball glove',
                   'skateboard', 'surfboard', 'tennis racket', 'bottle', 'wine glass', 'cup', 'fork', 'knife', 'spoon',
                   'bowl', 'banana', 'apple', 'sandwich', 'orange', 'broccoli', 'carrot', 'hot dog', 'pizza', 'donut',
                   'cake', 'chair', 'couch', 'potted plant', 'bed', 'dining table', 'toilet', 'tv', 'laptop', 'mouse',
                   'remote', 'keyboard', 'cell phone', 'microwave', 'oven', 'toaster', 'sink', 'refrigerator', 'book',
                   'clock', 'vase', 'scissors', 'teddy bear', 'hair drier', 'toothbrush']

        # get predictor
        self.data_names = ['data', 'im_info']
        self.label_names = []
        data = []
        sam = np.zeros((config.SCALES[0][0], config.SCALES[0][1], 3))
        im, im_scale = resize(sam, config.SCALES[0][0], config.SCALES[0][1], stride=config.network.IMAGE_STRIDE)
        im_tensor = transform(im, config.network.PIXEL_MEANS)
        im_info = np.array([[im_tensor.shape[2], im_tensor.shape[3], im_scale]], dtype=np.float32)
        data.append({'data': im_tensor, 'im_info': im_info})
        data = [[mx.nd.array(data[i][name]) for name in self.data_names] for i in xrange(len(data))]

        max_data_shape = [[('data', (1, 3, max([v[0] for v in config.SCALES]), max([v[1] for v in config.SCALES])))]]
        provide_data = [[(k, v.shape) for k, v in zip(self.data_names, data[i])] for i in xrange(len(data))]
        provide_label = [None for i in xrange(len(data))]
        arg_params, aux_params = load_param(model_path, 0, process=True)
        self.predictor = Predictor(sym, self.data_names, self.label_names,
                                   context=[mx.gpu(0)], max_data_shapes=max_data_shape,
                                   provide_data=provide_data, provide_label=provide_label,
                                   arg_params=arg_params, aux_params=aux_params)
        self.nms = gpu_nms_wrapper(config.TEST.NMS, 0)

    def detect_objects(self, input_image):
        data = []
        target_size = config.SCALES[0][0]
        max_size = config.SCALES[0][1]
        im, im_scale = resize(input_image, target_size, max_size, stride=config.network.IMAGE_STRIDE)
        im_tensor = transform(input_image, config.network.PIXEL_MEANS)
        im_info = np.array([[im_tensor.shape[2], im_tensor.shape[3], im_scale]], dtype=np.float32)
        data.append({'data': im_tensor, 'im_info': im_info})
        data = [[mx.nd.array(data[i][name]) for name in self.data_names] for i in xrange(len(data))]

        data_batch = mx.io.DataBatch(data=[data[0]], label=[], pad=0, index=0,
                                     provide_data=[[(k, v.shape) for k, v in zip(self.data_names, data[0])]],
                                     provide_label=[None])
        scales = [im_scale]
        scores, boxes, data_dict = im_detect(self.predictor, data_batch, self.data_names, scales, config)
        boxes = boxes[0].astype('f')
        scores = scores[0].astype('f')
        dets_nms = []
        for j in range(1, scores.shape[1]):
            cls_scores = scores[:, j, np.newaxis]
            cls_boxes = boxes[:, 4:8] if config.CLASS_AGNOSTIC else boxes[:, j * 4:(j + 1) * 4]
            cls_dets = np.hstack((cls_boxes, cls_scores))
            keep = self.nms(cls_dets)
            cls_dets = cls_dets[keep, :]
            cls_dets = cls_dets[cls_dets[:, -1] > 0.7, :]
            cls_dets[:4] *= im_scale
            dets_nms.append(cls_dets)
        # im = cv2.cvtColor(input_image, cv2.COLOR_BGR2RGB)
        # show_boxes(im, dets_nms, self.classes, 1)
        return dets_nms


from keras import backend as K
import time
from multiprocessing.dummy import Pool
K.set_image_data_format('channels_first')
import cv2
import os
import glob
import numpy as np
from numpy import genfromtxt
import tensorflow as tf
from dir_util.fr_utils import *
from dir_util.inception_blocks_v2 import *
import win32com.client as wincl

normalize_image_borders = 50
assert_detect_identity = True
speak_welcome_message = wincl.Dispatch("SAPI.SpVoice")
font = cv2.FONT_HERSHEY_SIMPLEX
facial_recognition_model = faceRecoModel(input_shape=(3, 96, 96))

'''
The loss function from the thesis implement here.
FaceRecoModel was ported and optimised to run on average computing resources.
'''
def own_loss_function(y_true, y_pred, alpha = 0.3):

    differential, pos, neg = y_pred[0], y_pred[1], y_pred[2]

    # Step 1: Reduce distance between the differential and positive
    positive_distance = tf.reduce_sum(tf.square(tf.subtract(differential, pos)), axis=-1) # axis -1 IMPORTANT
    # Step 2: Reduce distance between the differential and negative
    negative_distance = tf.reduce_sum(tf.square(tf.subtract(differential, neg)), axis=-1) # axis -1 IMPORTANT
    # Step 3: get the difference and add alpha
    intermediate_pooling_loss = tf.add(tf.subtract(positive_distance, negative_distance), alpha)
    # Step 4: Argmax of basic loss and 0
    final_loss = tf.reduce_sum(tf.maximum(intermediate_pooling_loss, 0.0))

    return final_loss

facial_recognition_model.compile(optimizer ='adam', loss = own_loss_function, metrics = ['accuracy'])

load_weights_from_FaceNet(facial_recognition_model)

def prepare_dictionary():
    person_dictionary = {}

    # Load training data
    for image in glob.glob("images/*"):
        person_id = os.path.splitext(os.path.basename(image))[0]
        person_dictionary[person_id] = img_path_to_encoding(image, facial_recognition_model)

    return person_dictionary

def webcam_video_feed(database):
    global assert_detect_identity

    cv2.namedWindow("preview")
    video_feed = cv2.VideoCapture(0)

    face_detectors = cv2.CascadeClassifier('dir_util/haarcascade_frontalface_default.xml')

    while video_feed.isOpened():
        _, video_frame = video_feed.read()
        img = video_frame

        # Securely detect single faces.
        if assert_detect_identity:
            img = process_frame(img, video_frame, face_detectors)

        key = cv2.waitKey(100)
        cv2.imshow("preview", img)

        if key == 27:
            break
    cv2.destroyWindow("preview")

def process_frame(img, frame, face_cascade):
    """
    Find whether people from the database is in the video frame.
    """
    global assert_detect_identity
    remove_color_channel = cv2.cvtColor(frame, cv2.COLOR_BGR2GRAY)
    faces_list = face_cascade.detectMultiScale(remove_color_channel, 1.3, 5)

    # Loop faces
    identities = []
    for (x, y, w, h) in faces_list:
        x_new = x - normalize_image_borders
        y_new = y - normalize_image_borders
        x2_new = x + w + normalize_image_borders
        y2_new = y + h + normalize_image_borders

        img = cv2.rectangle(frame,(x_new, y_new),(x2_new, y2_new),(255,0,0),2)

        person_id = find_person_identity(img, frame, x_new, y_new, x2_new, y2_new)

        if person_id is not None:
            identities.append(person_id)
            id_split = person_id.split("_")
            print(id_split)
            #filename = output('attendance', 'class1', int(id_split[1]), id_split[0], 'yes')

    if identities != []:
        #cv2.imwrite('example.png',img)
        assert_detect_identity = False
        thread_pool = Pool(processes=1)
        # Thread pools are created to run some processes concurrently through asynchronous calls
        thread_pool.apply_async(welcome_users, [identities])
    return img

def find_person_identity(img, frame, x1_new, y1_new, x2_new, y2_new):

    h_max, w_max, color_channels = frame.shape
    # Display the bounding box
    part_video_frame = frame[max(0, y1_new):min(h_max, y2_new), max(0, x1_new):min(w_max, x2_new)]

    return who_is_the_person(img, part_video_frame, database, facial_recognition_model, x1_new, y1_new)

def who_is_the_person(frame, image, database, facenet, x1, y1):
    encoding = img_to_encoding(image, facenet)

    minimum_distance = 100
    persons_id = None

    # for each name in database
    for (name, database_encoding) in database.items():

        # calculate euclidean distance
        encoding_distance = np.linalg.norm(database_encoding - encoding)
        cv2.putText(frame, name, (x1 + 30, y1 - 10), font, 1, (120, 255, 120), 2, 1)
        print('distance for %s is %s' %(name, encoding_distance))

        # If this distance is less than the min_dist, then set min_dist to dist, and identity to name
        if encoding_distance < minimum_distance:
            minimum_distance = encoding_distance
            persons_id = name

    if minimum_distance > 0.49:
        return None
    else:
        return str(persons_id)

def welcome_users(identity_matrix):
    global assert_detect_identity
    enter_message = 'Welcome '

    if len(identity_matrix) == 1:
        enter_message += '%s, you may enter.' % identity_matrix[0]
    else:
        for identity_id in range(len(identity_matrix) - 1):
            enter_message += '%s, ' % identity_matrix[identity_id]
        enter_message += 'and %s, ' % identity_matrix[-1]
        enter_message += 'you may enter'

    speak_welcome_message.Speak(enter_message)

    # Rerun any unused threads
    assert_detect_identity = True

if __name__ == "__main__":
    database = prepare_dictionary()
    webcam_video_feed(database)


#-*- coding: utf-8 -*-
from __future__ import division
import os
import math
import time
import tensorflow as tf
import numpy as np
from tensorflow.python.ops import math_ops
from ops import *
from utils import *
from GRUI import mygru_cell

"""
D输入标准化， 不要m 填充0
G输入去掉m,只有delta
g 没有每次累加z
"""
class WGAN(object):
    model_name = "WGAN_no_mask"     # name for checkpoint

    def __init__(self, sess, args, datasets):
        self.sess = sess
        self.isbatch_normal=args.isBatch_normal
        self.isNormal=args.isNormal
        self.checkpoint_dir = args.checkpoint_dir
        self.result_dir = args.result_dir
        self.log_dir = args.log_dir
        self.dataset_name=args.dataset_name
        self.run_type=args.run_type
        self.lr = args.lr
        self.epoch = args.epoch
        self.batch_size = args.batch_size
        self.n_inputs = args.n_inputs                 # MNIST data input (img shape: 28*28)
        self.n_steps = datasets.maxLength                                # time steps
        self.n_hidden_units = args.n_hidden_units        # neurons in hidden layer
        self.n_classes = args.n_classes                # MNIST classes (0-9 digits)
        self.gpus=args.gpus
        self.run_type=args.run_type
        self.result_path=args.result_path
        self.model_path=args.model_path
        self.pretrain_epoch=args.pretrain_epoch
        self.impute_iter=args.impute_iter
        self.isSlicing=args.isSlicing
        self.g_loss_lambda=args.g_loss_lambda

        self.datasets=datasets
        self.z_dim = args.z_dim         # dimension of noise-vector
        self.gen_length=args.gen_length

        # WGAN_GP parameter
        self.lambd = 0.25       # The higher value, the more stable, but the slower convergence
        self.disc_iters = args.disc_iters     # The number of critic iterations for one-step of generator

        # train
        self.learning_rate = args.lr
        self.beta1 = args.beta1
        if "1.5" in tf.__version__ or "1.7" in tf.__version__ :
            self.grud_cell_d = mygru_cell.MyGRUCell15(self.n_hidden_units)
            self.grud_cell_g = mygru_cell.MyGRUCell15(self.n_hidden_units)
        elif "1.4" in tf.__version__:
            self.grud_cell_d = mygru_cell.MyGRUCell4(self.n_hidden_units)
            self.grud_cell_g = mygru_cell.MyGRUCell4(self.n_hidden_units)
        elif "1.2" in tf.__version__:
            self.grud_cell_d = mygru_cell.MyGRUCell2(self.n_hidden_units)
            self.grud_cell_g = mygru_cell.MyGRUCell2(self.n_hidden_units)
        # test
        self.sample_num = 64  # number of generated images to be saved

        self.num_batches = len(datasets.x) // self.batch_size

    def pretrainG(self, X, M, Delta,  Mean, Lastvalues, X_lengths, Keep_prob, is_training=True, reuse=False):

        with tf.variable_scope("g_enerator", reuse=reuse):

            """
            the rnn cell's variable scope is defined by tensorflow,
            if we want to update rnn cell's weights, the variable scope must contains 'g_' or 'd_'

            """

            wr_h=tf.get_variable("g_wr_h",shape=[self.n_inputs,self.n_hidden_units],initializer=tf.random_normal_initializer())
            w_out= tf.get_variable("g_w_out",shape=[self.n_hidden_units, self.n_inputs],initializer=tf.random_normal_initializer())

            br_h= tf.get_variable("g_br_h",shape=[self.n_hidden_units, ],initializer=tf.constant_initializer(0.001))
            b_out= tf.get_variable("g_b_out",shape=[self.n_inputs, ],initializer=tf.constant_initializer(0.001))
            w_z=tf.get_variable("g_w_z",shape=[self.z_dim,self.n_inputs],initializer=tf.random_normal_initializer())
            b_z=tf.get_variable("g_b_z",shape=[self.n_inputs, ],initializer=tf.constant_initializer(0.001))

            X = tf.reshape(X, [-1, self.n_inputs])
            Delta=tf.reshape(Delta,[-1,self.n_inputs])

            rth= tf.matmul(Delta, wr_h)+br_h
            rth=math_ops.exp(-tf.maximum(0.0,rth))

            X=tf.concat([X,rth],1)

            X_in = tf.reshape(X, [-1, self.n_steps, self.n_inputs+self.n_hidden_units])

            init_state = self.grud_cell_g.zero_state(self.batch_size, dtype=tf.float32) # 初始化全零 state
            outputs, final_state = tf.nn.dynamic_rnn(self.grud_cell_g, X_in, \
                                initial_state=init_state,\
                                sequence_length=X_lengths,
                                time_major=False)
            #outputs: batch_size*n_steps*n_hiddensize
            outputs=tf.reshape(outputs,[-1,self.n_hidden_units])
            out_predict=tf.matmul(tf.nn.dropout(outputs,Keep_prob), w_out) + b_out
            out_predict=tf.reshape(out_predict,[-1,self.n_steps,self.n_inputs])
            return out_predict

    def discriminator(self, X, M, DeltaPre, Lastvalues ,DeltaSub ,SubValues , Mean,  X_lengths,Keep_prob, is_training=True, reuse=False, isTdata=True):
        # Network Architecture is exactly same as in infoGAN (https://arxiv.org/abs/1606.03657)
        # Architecture : (64)4c2s-(128)4c2s_BL-FC1024_BL-FC1_S
        with tf.variable_scope("d_iscriminator", reuse=reuse):

            wr_h=tf.get_variable("d_wr_h",shape=[self.n_inputs,self.n_hidden_units],initializer=tf.random_normal_initializer())
            w_out= tf.get_variable("d_w_out",shape=[self.n_hidden_units, 1],initializer=tf.random_normal_initializer())
            br_h= tf.get_variable("d_br_h",shape=[self.n_hidden_units, ],initializer=tf.constant_initializer(0.001))
            b_out= tf.get_variable("d_b_out",shape=[1, ],initializer=tf.constant_initializer(0.001))

            M=tf.reshape(M,[-1,self.n_inputs])
            X = tf.reshape(X, [-1, self.n_inputs])
            DeltaPre=tf.reshape(DeltaPre,[-1,self.n_inputs])

            rth= tf.matmul(DeltaPre, wr_h)+br_h
            rth=math_ops.exp(-tf.maximum(0.0,rth))
            # add noise
            #X=X+np.random.standard_normal(size=(self.batch_size*self.n_steps, self.n_inputs))/100
            X=tf.concat([X,rth],1)

            X_in = tf.reshape(X, [self.batch_size, self.n_steps , self.n_inputs+self.n_hidden_units])

            init_state = self.grud_cell_d.zero_state(self.batch_size, dtype=tf.float32) # 初始化全零 state
            outputs, final_state = tf.nn.dynamic_rnn(self.grud_cell_d, X_in, \
                                initial_state=init_state,\
                                sequence_length=X_lengths,
                                time_major=False)

            # final_state:batch_size*n_hiddensize
            # 不能用最后一个，应该用第length个  之前用了最后一个，所以输出无论如何都是b_out
            out_logit=tf.matmul(tf.nn.dropout(final_state,Keep_prob), w_out) + b_out
            out =tf.nn.sigmoid(out_logit)    #选取最后一个 output
            return out,out_logit

    def generator(self, z, Keep_prob, is_training=True, reuse=False):
        # x,delta,n_steps
        # z :[self.batch_size, self.z_dim]
        # first feed noize in rnn, then feed the previous output into next input
        # or we can feed noize and previous output into next input in future version
        with tf.variable_scope("g_enerator", reuse=reuse):
            #gennerate

            wr_h=tf.get_variable("g_wr_h",shape=[self.n_inputs,self.n_hidden_units],initializer=tf.random_normal_initializer())
            w_out= tf.get_variable("g_w_out",shape=[self.n_hidden_units, self.n_inputs],initializer=tf.random_normal_initializer())
            br_h= tf.get_variable("g_br_h",shape=[self.n_hidden_units, ],initializer=tf.constant_initializer(0.001))
            b_out= tf.get_variable("g_b_out",shape=[self.n_inputs, ],initializer=tf.constant_initializer(0.001))
            w_z=tf.get_variable("g_w_z",shape=[self.z_dim,self.n_inputs],initializer=tf.random_normal_initializer())
            b_z=tf.get_variable("g_b_z",shape=[self.n_inputs, ],initializer=tf.constant_initializer(0.001))

            #self.times=tf.reshape(self.times,[self.batch_size,self.n_steps,self.n_inputs])
            #change z's dimension
            # batch_size*z_dim-->batch_size*n_inputs
            x=tf.matmul(z,w_z)+b_z
            x=tf.reshape(x,[-1,self.n_inputs])
            delta_zero=tf.constant(0.0,shape=[self.batch_size,self.n_inputs])
            #delta_normal=tf.constant(48.0*60.0/self.gen_length,shape=[self.batch_size,self.n_inputs])
            #delta:[batch_size,1,n_inputs]

            # combine X_in
            rth= tf.matmul(delta_zero, wr_h)+br_h
            rth=math_ops.exp(-tf.maximum(0.0,rth))
            x=tf.concat([x,rth],1)

            X_in = tf.reshape(x, [-1, 1, self.n_inputs+self.n_hidden_units])

            init_state = self.grud_cell_g.zero_state(self.batch_size, dtype=tf.float32) # 初始化全零 state
            #z=tf.reshape(z,[self.batch_size,1,self.z_dim])
            seq_len=tf.constant(1,shape=[self.batch_size])

            outputs, final_state = tf.nn.dynamic_rnn(self.grud_cell_g, X_in, \
                                initial_state=init_state,\
                                sequence_length=seq_len,
                                time_major=False)
            init_state=final_state
            #outputs: batch_size*1*n_hidden
            outputs=tf.reshape(outputs,[-1,self.n_hidden_units])
            # full connect
            out_predict=tf.matmul(tf.nn.dropout(outputs,Keep_prob), w_out) + b_out
            out_predict=tf.reshape(out_predict,[-1,1,self.n_inputs])

            total_result=tf.multiply(out_predict,1.0)

            for i in range(1,self.n_steps):
                out_predict=tf.reshape(out_predict,[self.batch_size,self.n_inputs])
                #输出加上noise z
                #out_predict=out_predict+tf.matmul(z,w_z)+b_z
                #
                delta_normal=tf.reshape(self.imputed_deltapre[:,i:(i+1),:],[self.batch_size,self.n_inputs])
                rth= tf.matmul(delta_normal, wr_h)+br_h
                rth=math_ops.exp(-tf.maximum(0.0,rth))
                x=tf.concat([out_predict,rth],1)
                X_in = tf.reshape(x, [-1, 1, self.n_inputs+self.n_hidden_units])

                outputs, final_state = tf.nn.dynamic_rnn(self.grud_cell_g, X_in, \
                            initial_state=init_state,\
                            sequence_length=seq_len,
                            time_major=False)
                init_state=final_state
                outputs=tf.reshape(outputs,[-1,self.n_hidden_units])
                out_predict=tf.matmul(tf.nn.dropout(outputs,Keep_prob), w_out) + b_out
                out_predict=tf.reshape(out_predict,[-1,1,self.n_inputs])
                total_result=tf.concat([total_result,out_predict],1)

            #delta:[batch_size,,n_inputs]

            if self.isbatch_normal:
                with tf.variable_scope("g_bn", reuse=tf.AUTO_REUSE):
                    total_result=bn(total_result,is_training=is_training, scope="g_bn_imple")

            last_values=tf.multiply(total_result,1)
            sub_values=tf.multiply(total_result,1)

            return total_result,self.imputed_deltapre,self.imputed_deltasub,self.imputed_m,self.x_lengths,last_values,sub_values

    def impute(self):
        with tf.variable_scope("impute", reuse=tf.AUTO_REUSE):
            z_need_tune=tf.get_variable("z_needtune",shape=[self.batch_size,self.z_dim],initializer=tf.random_normal_initializer(mean=0,stddev=0.1) )
            return z_need_tune

    def build_model(self):

        self.keep_prob = tf.placeholder(tf.float32)
        self.x = tf.placeholder(tf.float32, [None, self.n_steps, self.n_inputs])
        self.y = tf.placeholder(tf.float32, [None, self.n_classes])
        self.m = tf.placeholder(tf.float32, [None, self.n_steps, self.n_inputs])
        self.mean = tf.placeholder(tf.float32, [self.n_inputs,])
        self.deltaPre = tf.placeholder(tf.float32, [None, self.n_steps, self.n_inputs])
        self.lastvalues = tf.placeholder(tf.float32, [None, self.n_steps, self.n_inputs])
        self.deltaSub = tf.placeholder(tf.float32, [None, self.n_steps, self.n_inputs])
        self.subvalues = tf.placeholder(tf.float32, [None, self.n_steps, self.n_inputs])
        self.x_lengths = tf.placeholder(tf.int32,  shape=[self.batch_size,])
        self.imputed_deltapre=tf.placeholder(tf.float32,  shape=[self.batch_size,self.n_steps,self.n_inputs])
        self.imputed_deltasub=tf.placeholder(tf.float32,  shape=[self.batch_size,self.n_steps,self.n_inputs])
        self.imputed_m = tf.placeholder(tf.float32, [None, self.n_steps, self.n_inputs])
        self.z = tf.placeholder(tf.float32, [self.batch_size, self.z_dim], name='z')

        """ Loss Function """

        Pre_out=self.pretrainG(self.x, self.m, self.deltaPre,  self.mean,\
                                                      self.lastvalues, self.x_lengths,self.keep_prob, \
                                                      is_training=True, reuse=False)

        self.pretrain_loss=tf.reduce_sum(tf.square(tf.multiply(Pre_out,self.m)-self.x)) / tf.cast(tf.reduce_sum(self.x_lengths),tf.float32)

        #discriminator( X, M, DeltaPre, Lastvalues ,DeltaSub ,SubValues , Mean,  X_lengths,Keep_prob, is_training=True, reuse=False, isTdata=True):

        D_real, D_real_logits = self.discriminator(self.x, self.m, self.deltaPre,self.lastvalues,\
                                                   self.deltaSub,self.subvalues,  self.mean,\
                                                       self.x_lengths,self.keep_prob, \
                                                      is_training=True, reuse=False,isTdata=True)

        #G return total_result,self.imputed_deltapre,self.imputed_deltasub,self.imputed_m,self.x_lengths,last_values,sub_values
        g_x,g_deltapre,g_deltasub,g_m,G_x_lengths,g_last_values,g_sub_values = self.generator(self.z,self.keep_prob, is_training=True, reuse=True)

        D_fake, D_fake_logits = self.discriminator(g_x,g_m,g_deltapre,g_last_values,\
                                                   g_deltasub,g_sub_values,self.mean,\
                                                      G_x_lengths,self.keep_prob,
                                                      is_training=True, reuse=True ,isTdata=False)

        """
        impute loss
        """
        self.z_need_tune=self.impute()

        impute_out,impute_deltapre,impute_deltasub,impute_m,impute_x_lengths,impute_last_values,impute_sub_values=self.generator(self.z_need_tune,self.keep_prob, is_training=False, reuse=True)

        impute_fake, impute_fake_logits = self.discriminator(impute_out,impute_m,impute_deltapre,impute_last_values,\
                                                             impute_deltasub,impute_sub_values,self.mean,\
                                                      impute_x_lengths,self.keep_prob,
                                                      is_training=False, reuse=True ,isTdata=False)

        # loss for imputation

        self.impute_loss=tf.reduce_mean(tf.square(tf.multiply(impute_out,self.m)-self.x))-self.g_loss_lambda*tf.reduce_mean(impute_fake_logits)
        #self.impute_loss=tf.reduce_mean(tf.square(tf.multiply(impute_out,self.m)-self.x))

        self.impute_out=impute_out

        #the imputed results
        self.imputed=tf.multiply((1-self.m),self.impute_out)+self.x
        # get loss for discriminator
        d_loss_real = - tf.reduce_mean(D_real_logits)
        d_loss_fake = tf.reduce_mean(D_fake_logits)

        self.d_loss = d_loss_real + d_loss_fake

        # get loss for generator
        self.g_loss = - d_loss_fake

        """ Training """
        # divide trainable variables into a group for D and a group for G
        t_vars = tf.trainable_variables()
        d_vars = [var for var in t_vars if 'd_' in var.name]
        g_vars = [var for var in t_vars if 'g_' in var.name]
        z_vars = [self.z_need_tune]
        '''
        print("d vars:")
        for v in d_vars:
            print(v.name)
        print("g vars:")
        for v in g_vars:
            print(v.name)
        print("z vars:")
        for v in z_vars:
            print(v.name)
        '''

        #don't need normalization because we have adopted the dropout

        ld = 0.0
        for w in d_vars:
            ld += tf.contrib.layers.l2_regularizer(1e-4)(w)
        lg = 0.0
        for w in g_vars:
            lg += tf.contrib.layers.l2_regularizer(1e-4)(w)

        self.d_loss+=ld
        self.g_loss+=lg

        # optimizers
        with tf.control_dependencies(tf.get_collection(tf.GraphKeys.UPDATE_OPS)):
        # this code have used batch normalization, so the upside line should be executed
            self.d_optim = tf.train.AdamOptimizer(self.learning_rate, beta1=self.beta1) \
                        .minimize(self.d_loss, var_list=d_vars)
            #self.d_optim=self.optim(self.learning_rate, self.beta1,self.d_loss,d_vars)
            self.g_optim = tf.train.AdamOptimizer(self.learning_rate*self.disc_iters, beta1=self.beta1) \
                        .minimize(self.g_loss, var_list=g_vars)
            #self.g_optim=self.optim(self.learning_rate, self.beta1,self.g_loss,g_vars)
            self.g_pre_optim=tf.train.AdamOptimizer(self.learning_rate*2,beta1=self.beta1) \
                        .minimize(self.pretrain_loss,var_list=g_vars)
        self.impute_optim=tf.train.AdamOptimizer(self.learning_rate*7,beta1=self.beta1) \
                    .minimize(self.impute_loss,var_list=z_vars)

        #clip weight
        # self.clip_all_vals = [p.assign(tf.clip_by_value(p, -0.99, 0.99)) for p in t_vars]
        # self.clip_D = [p.assign(tf.clip_by_value(p, -0.99, 0.99)) for p in d_vars]
        # self.clip_G = [p.assign(tf.clip_by_value(p, -0.99, 0.99)) for p in g_vars]

        self.clip_all_vals = [p.assign(tf.clip_by_value(p, -0.75, 0.75)) for p
                              in t_vars]
        self.clip_D = [p.assign(tf.clip_by_value(p, -0.75, 0.75)) for p in
                       d_vars]
        self.clip_G = [p.assign(tf.clip_by_value(p, -0.75, 0.75)) for p in
                       g_vars]

        """" Testing """
        # for test
        self.fake_x,self.fake_delta,_,_,_,_,_ = self.generator(self.z, self.keep_prob, is_training=False, reuse=True)

        """ Summary """
        d_loss_real_sum = tf.summary.scalar("d_loss_real", d_loss_real)
        d_loss_fake_sum = tf.summary.scalar("d_loss_fake", d_loss_fake)
        d_loss_sum = tf.summary.scalar("d_loss", self.d_loss)
        g_loss_sum = tf.summary.scalar("g_loss", self.g_loss)
        g_pretrain_loss_sum=tf.summary.scalar("g_pretrain_loss", self.pretrain_loss)
        # final summary operations
        self.impute_sum=tf.summary.scalar("impute_loss", self.impute_loss)
        self.g_sum = g_loss_sum
        self.g_pretrain_sum=tf.summary.merge([g_pretrain_loss_sum])
        self.d_sum = tf.summary.merge([d_loss_real_sum,d_loss_fake_sum, d_loss_sum])

    def optim(self,learning_rate,beta,loss,var):
        optimizer = tf.train.AdamOptimizer(learning_rate, beta1=beta)
        grads = optimizer.compute_gradients(loss,var_list=var)
        for i, (g, v) in enumerate(grads):
            if g is not None:
                grads[i] = (tf.clip_by_norm(g, 5), v)  # clip gradients
        train_op = optimizer.apply_gradients(grads)
        return train_op
    def pretrain(self, start_epoch,counter,start_time):

        if start_epoch < self.pretrain_epoch:
            #todo
            for epoch in range(start_epoch, self.pretrain_epoch):
            # get batch data
                self.datasets.shuffle(self.batch_size,True)
                idx=0
                #x,y,mean,m,deltaPre,x_lengths,lastvalues,files,imputed_deltapre,imputed_m,deltaSub,subvalues,imputed_deltasub
                for data_x,data_y,data_mean,data_m,data_deltaPre,data_x_lengths,data_lastvalues,_,imputed_deltapre,imputed_m,deltaSub,subvalues,imputed_deltasub in self.datasets.nextBatch():

                    # pretrain
                    _, summary_str, p_loss = self.sess.run([self.g_pre_optim, self.g_pretrain_sum, self.pretrain_loss],
                                                   feed_dict={self.x: data_x,
                                                              self.m: data_m,
                                                              self.deltaPre: data_deltaPre,
                                                              self.mean: data_mean,
                                                              self.x_lengths: data_x_lengths,
                                                              self.lastvalues: data_lastvalues,
                                                              self.deltaSub:deltaSub,
                                                              self.subvalues:subvalues,
                                                              self.imputed_m:imputed_m,
                                                              self.imputed_deltapre:imputed_deltapre,
                                                              self.imputed_deltasub:imputed_deltasub,
                                                              self.keep_prob: 0.5})
                    self.writer.add_summary(summary_str, counter)

                    counter += 1

                    # display training status
                    print("Epoch: [%2d] [%4d/%4d] time: %4.4f, pretrain_loss: %.8f" \
                          % (epoch, idx, self.num_batches, time.time() - start_time, p_loss))
                    idx+=1
                # After an epoch, start_batch_id is set to zero
                # non-zero value is only for the first epoch after loading pre-trained model
                start_batch_id = 0

                # save model
                #调好之后再保存
                #if epoch%10==0:
                #    self.save(self.checkpoint_dir, counter)

    def train(self):

        # graph inputs for visualize training results
        self.sample_z = np.random.standard_normal(size=(self.batch_size , self.z_dim))

        # saver to save model
        self.saver = tf.train.Saver()

        # summary writer
        self.writer = tf.summary.FileWriter(self.log_dir + '/' + self.model_name+'/'+self.model_dir)

        # restore check-point if it exits
        could_load, checkpoint_counter = self.load(self.checkpoint_dir)
        if could_load:
            start_epoch = (int)(checkpoint_counter / self.num_batches)
            #start_batch_id = checkpoint_counter - start_epoch * self.num_batches
            start_batch_id=0
            #counter = checkpoint_counter
            counter=start_epoch*self.num_batches
            print(" [*] Load SUCCESS")
            return
        else:
            # initialize all variables
            tf.global_variables_initializer().run()
            start_epoch = 0
            start_batch_id = 0
            counter = 1
            print(" [!] Load failed...")

        # loop for epoch
        start_time = time.time()

        self.pretrain(start_epoch,counter,start_time)
        if start_epoch < self.pretrain_epoch:
            start_epoch=self.pretrain_epoch

        for epoch in range(start_epoch, self.epoch):

            # get batch data
            self.datasets.shuffle(self.batch_size,True)
            idx=0
            for data_x,data_y,data_mean,data_m,data_deltaPre,data_x_lengths,data_lastvalues,_,imputed_deltapre,imputed_m,deltaSub,subvalues,imputed_deltasub in self.datasets.nextBatch():

                batch_z = np.random.standard_normal(size=(self.batch_size, self.z_dim))

                # print("batch_z: " + str(batch_z))
                # print("data_x: " + str(data_x))
                # print("data_m: " + str(data_m))
                # print("data_deltaPre: " + str(data_deltaPre))
                # print("data_mean: " + str(data_mean))
                # print("max data_x_lengths: " + str(max(data_x_lengths)))
                # print("data_lastvalues: " + str(data_lastvalues))
                # print("deltaSub: " + str(deltaSub))
                # print("subvalues: " + str(subvalues))
                # print("imputed_m: " + str(imputed_m))
                # print("imputed_deltapre: " + str(imputed_deltapre))
                # print("imputed_deltasub: " + str(imputed_deltasub))

                #_ = self.sess.run(self.clip_D)
                _ = self.sess.run(self.clip_all_vals)
                _, summary_str, d_loss = self.sess.run([self.d_optim, self.d_sum, self.d_loss],
                                               feed_dict={self.z: batch_z,
                                                          self.x: data_x,
                                                          self.m: data_m,
                                                          self.deltaPre: data_deltaPre,
                                                          self.mean: data_mean,
                                                          self.x_lengths: data_x_lengths,
                                                          self.lastvalues: data_lastvalues,
                                                          self.deltaSub:deltaSub,
                                                          self.subvalues:subvalues,
                                                          self.imputed_m:imputed_m,
                                                          self.imputed_deltapre:imputed_deltapre,
                                                          self.imputed_deltasub:imputed_deltasub,
                                                          self.keep_prob: 0.5})
                self.writer.add_summary(summary_str, counter)

                # update G network
                if counter%self.disc_iters==0:
                    #batch_z = np.random.normal(0, 1, [self.batch_size, self.z_dim]).astype(np.float32)
                    _, summary_str, g_loss = self.sess.run([self.g_optim, self.g_sum, self.g_loss],
                                                           feed_dict={self.z: batch_z,
                                                           self.keep_prob: 0.5,
                                                           self.deltaPre: data_deltaPre,
                                                           self.mean: data_mean,
                                                           self.x_lengths: data_x_lengths,
                                                           self.lastvalues: data_lastvalues,
                                                           self.deltaSub:deltaSub,
                                                           self.subvalues:subvalues,
                                                           self.imputed_m:imputed_m,
                                                           self.imputed_deltapre:imputed_deltapre,
                                                           self.imputed_deltasub:imputed_deltasub,
                                                           self.mean: data_mean})
                    self.writer.add_summary(summary_str, counter)
                    print("Epoch: [%2d] [%4d/%4d] time: %4.4f, d_loss: %.8f, g_loss: %.8f,counter:%4d" \
                      % (epoch, idx, self.num_batches, time.time() - start_time, d_loss, g_loss,counter))
                    #debug

                counter += 1

                # display training status
                print("Epoch: [%2d] [%4d/%4d] time: %4.4f, d_loss: %.8f, counter:%4d" \
                      % (epoch, idx, self.num_batches, time.time() - start_time, d_loss, counter))

                # save training results for every 300 steps
                if np.mod(counter, 300) == 0 :
                    fake_x,fake_delta = self.sess.run([self.fake_x,self.fake_delta],
                                            feed_dict={self.z: batch_z,
                                                       self.deltaPre: data_deltaPre,
                                                       self.mean: data_mean,
                                                       self.x_lengths: data_x_lengths,
                                                       self.lastvalues: data_lastvalues,
                                                       self.deltaSub:deltaSub,
                                                       self.subvalues:subvalues,
                                                       self.imputed_m:imputed_m,
                                                       self.imputed_deltapre:imputed_deltapre,
                                                       self.imputed_deltasub:imputed_deltasub,
                                                       self.mean: data_mean,
                                                       self.keep_prob: 0.5})
                    if self.run_type=="train":
                        self.writeG_Samples("G_sample_x",counter,fake_x)
                        self.writeG_Samples("G_sample_delta",counter,fake_delta)

                idx+=1
            # After an epoch, start_batch_id is set to zero
            # non-zero value is only for the first epoch after loading pre-trained model
            start_batch_id = 0

        self.save(self.checkpoint_dir, counter)

    def imputation(self,dataset,isTrain):
        self.datasets=dataset
        self.datasets.shuffle(self.batch_size,True)
        tf.variables_initializer([self.z_need_tune]).run()
        #是否shuffle无所谓，填充之后存起来，测试的时候用填充之后的数据再shuffle即可
        #训练数据集不能被batch_size整除剩下的部分，扔掉
        start_time = time.time()
        batchid=1
        impute_tune_time=1
        counter=1
        for data_x,data_y,data_mean,data_m,data_deltaPre,data_x_lengths,data_lastvalues,_,imputed_deltapre,imputed_m,deltaSub,subvalues,imputed_deltasub in self.datasets.nextBatch():
            #self.z_need_tune=tf.assign(self.z_need_tune,tf.random_normal([self.batch_size,self.z_dim]))
            tf.variables_initializer([self.z_need_tune]).run()
            for i in range(0,self.impute_iter):
                _, impute_out, summary_str, impute_loss, imputed = self.sess.run([self.impute_optim, self.impute_out, self.impute_sum, self.impute_loss, self.imputed], \
                                                       feed_dict={self.x: data_x,
                                                                  self.m: data_m,
                                                                  self.deltaPre: data_deltaPre,
                                                                  self.mean: data_mean,
                                                                  self.x_lengths: data_x_lengths,
                                                                  self.lastvalues: data_lastvalues,
                                                                  self.deltaSub:deltaSub,
                                                                  self.subvalues:subvalues,
                                                                  self.imputed_m:imputed_m,
                                                                  self.imputed_deltapre:imputed_deltapre,
                                                                  self.imputed_deltasub:imputed_deltasub,
                                                                  self.keep_prob: 1.0})
                impute_tune_time+=1
                counter+=1
                if counter%10==0:
                    print("Batchid: [%2d] [%4d/%4d] time: %4.4f, impute_loss: %.8f" \
                          % (batchid, impute_tune_time, self.impute_iter, time.time() - start_time, impute_loss))
                    self.writer.add_summary(summary_str, counter/10)
            #imputed=tf.multiply((1-self.m),impute_out)+data_x
            self.save_imputation(imputed,batchid,data_x_lengths,data_deltaPre,data_y,isTrain)
            batchid+=1
            impute_tune_time=1
    @property
    def model_dir(self):
        return "{}_{}_{}_{}_{}_{}_{}_{}_{}_{}_{}".format(
            self.epoch,self.disc_iters,
            self.batch_size, self.z_dim,
            self.lr,self.impute_iter,
            self.isNormal,self.isbatch_normal,
            self.isSlicing,self.g_loss_lambda,
            self.beta1
            )

    def save_imputation(self,impute_out,batchid,data_x_lengths,data_times,data_y,isTrain):
        #填充后的数据全是n_steps长度！，但只有data_x_lengths才是可用的
        if isTrain:
            imputation_dir="imputation_train_results"
        else:
            imputation_dir="imputation_test_results"

        if not os.path.exists(os.path.join(imputation_dir,\
                                     self.model_name,\
                                     self.model_dir)):
            os.makedirs(os.path.join(imputation_dir,\
                                     self.model_name,\
                                     self.model_dir))

        #write imputed data
        resultFile=open(os.path.join(imputation_dir,\
                                     self.model_name,\
                                     self.model_dir,\
                                     "batch"+str(batchid)+"x"),'w')
        for length in data_x_lengths:
            resultFile.writelines(str(length)+",")
        resultFile.writelines("\r\n")
        # impute_out:ndarray
        for oneSeries in impute_out:
            resultFile.writelines("begin\r\n")
            for oneClass in oneSeries:
                for i in oneClass.flat:
                    resultFile.writelines(str(i)+",")
                resultFile.writelines("\r\n")
            resultFile.writelines("end\r\n")
        resultFile.close()

        #write data_times data_times:list
        resultFile=open(os.path.join(imputation_dir,\
                                     self.model_name,\
                                     self.model_dir,\
                                     "batch"+str(batchid)+"delta"),'w')
        for oneSeries in data_times:
            resultFile.writelines("begin\r\n")
            for oneClass in oneSeries:
                for i in oneClass:
                    resultFile.writelines(str(i)+",")
                resultFile.writelines("\r\n")
            resultFile.writelines("end\r\n")
        resultFile.close()

        #write y
        resultFile=open(os.path.join(imputation_dir,\
                                     self.model_name,\
                                     self.model_dir,\
                                     "batch"+str(batchid)+"y"),'w')
        for oneSeries in data_y:
            #resultFile.writelines("begin\r\n")
            for oneClass in oneSeries:
                resultFile.writelines(str(oneClass)+",")
            resultFile.writelines("\r\n")
            #resultFile.writelines("end\r\n")
        resultFile.close()

    def writeG_Samples(self,filename,step,o):
        if not os.path.exists(os.path.join("G_results",\
                                     self.model_name,\
                                     self.model_dir)):
            os.makedirs(os.path.join("G_results",\
                                     self.model_name,\
                                     self.model_dir))
        resultFile=open(os.path.join("G_results",\
                                     self.model_name,\
                                     self.model_dir,\
                                     filename+str(step)),'w')
        for oneSeries in o:
            resultFile.writelines("begin\r\n")
            for oneClass in oneSeries:
                for i in oneClass.flat:
                    resultFile.writelines(str(i)+",")
                resultFile.writelines("\r\n")
            resultFile.writelines("end\r\n")
        resultFile.close()

    def save(self, checkpoint_dir, step):
        checkpoint_dir = os.path.join(checkpoint_dir, self.model_name, self.model_dir )

        if not os.path.exists(checkpoint_dir):
            os.makedirs(checkpoint_dir)

        self.saver.save(self.sess,os.path.join(checkpoint_dir, self.model_name+'.model'), global_step=step)

    def load(self, checkpoint_dir):
        import re
        print(" [*] Reading checkpoints...")
        checkpoint_dir = os.path.join(checkpoint_dir, self.model_name, self.model_dir)

        ckpt = tf.train.get_checkpoint_state(checkpoint_dir)
        if ckpt and ckpt.model_checkpoint_path:
            ckpt_name = os.path.basename(ckpt.model_checkpoint_path)
            self.saver.restore(self.sess, os.path.join(checkpoint_dir, ckpt_name))
            counter = int(next(re.finditer("(\d+)(?!.*\d)",ckpt_name)).group(0))
            print(" [*] Success to read {}".format(ckpt_name))
            return True, counter
        else:
            print(" [*] Failed to find a checkpoint")
            return False, 0


# ------------------------------------------------------------------------------
# Copyright (c) Microsoft
# Licensed under the MIT License.
# Written by Bin Xiao (leoxiaobin@gmail.com)
# Modified by Bowen Cheng (bcheng9@illinois.edu)
# ------------------------------------------------------------------------------

from __future__ import absolute_import
from __future__ import division
from __future__ import print_function

import logging

import numpy as np

import pycocotools
from .COCODataset import CocoDataset
from .target_generators import HeatmapGenerator

logger = logging.getLogger(__name__)

class CocoKeypoints(CocoDataset):
    def __init__(self,
                 cfg,
                 dataset_name,
                 remove_images_without_annotations,
                 heatmap_generator,
                 joints_generator,
                 transforms=None):
        super().__init__(cfg.DATASET.ROOT,
                         dataset_name,
                         cfg.DATASET.DATA_FORMAT)

        if cfg.DATASET.WITH_CENTER:
            assert cfg.DATASET.NUM_JOINTS == 18, 'Number of joint with center for COCO is 18'
        else:
            assert cfg.DATASET.NUM_JOINTS == 20, 'Number of joint for COCO is 17'

        self.num_scales = self._init_check(heatmap_generator, joints_generator)

        self.num_joints = cfg.DATASET.NUM_JOINTS
        self.with_center = cfg.DATASET.WITH_CENTER
        self.num_joints_without_center = self.num_joints - 1 \
            if self.with_center else self.num_joints
        self.scale_aware_sigma = cfg.DATASET.SCALE_AWARE_SIGMA
        self.base_sigma = cfg.DATASET.BASE_SIGMA
        self.base_size = cfg.DATASET.BASE_SIZE
        self.int_sigma = cfg.DATASET.INT_SIGMA

        if remove_images_without_annotations:
            self.ids = [
                img_id
                for img_id in self.ids
                if len(self.coco.getAnnIds(imgIds=img_id, iscrowd=None)) > 0
            ]

        self.transforms = transforms
        self.heatmap_generator = heatmap_generator
        self.joints_generator = joints_generator

    def __getitem__(self, idx):
        img, anno = super().__getitem__(idx)

        mask = self.get_mask(anno, idx)

        anno = [
            obj for obj in anno
            if obj['iscrowd'] == 0 or obj['num_keypoints'] > 0
        ]

        # TODO(bowen): to generate scale-aware sigma, modify `get_joints` to associate a sigma to each joint
        joints = self.get_joints(anno)

        mask_list = [mask.copy() for _ in range(self.num_scales)]
        joints_list = [joints.copy() for _ in range(self.num_scales)]
        target_list = list()

        if self.transforms:
            img, mask_list, joints_list = self.transforms(
                img, mask_list, joints_list
            )

        for scale_id in range(self.num_scales):
            target_t = self.heatmap_generator[scale_id](joints_list[scale_id])
            joints_t = self.joints_generator[scale_id](joints_list[scale_id])

            target_list.append(target_t.astype(np.float32))
            mask_list[scale_id] = mask_list[scale_id].astype(np.float32)
            joints_list[scale_id] = joints_t.astype(np.int32)

        return img, target_list, mask_list, joints_list

    def get_joints(self, anno):
        num_people = len(anno)

        if self.scale_aware_sigma:
            joints = np.zeros((num_people, self.num_joints, 4))
        else:
            joints = np.zeros((num_people, self.num_joints, 3))

        for i, obj in enumerate(anno):
            joints[i, :self.num_joints_without_center, :3] = \
                np.array(obj['keypoints']).reshape([-1, 3])
            if self.with_center:
                joints_sum = np.sum(joints[i, :-1, :2], axis=0)
                num_vis_joints = len(np.nonzero(joints[i, :-1, 2])[0])
                if num_vis_joints > 0:
                    joints[i, -1, :2] = joints_sum / num_vis_joints
                    joints[i, -1, 2] = 1
            if self.scale_aware_sigma:
                # get person box
                box = obj['bbox']
                size = max(box[2], box[3])
                sigma = size / self.base_size * self.base_sigma
                if self.int_sigma:
                    sigma = int(np.round(sigma + 0.5))
                assert sigma > 0, sigma
                joints[i, :, 3] = sigma

        return joints

    def get_mask(self, anno, idx):
        coco = self.coco
        img_info = coco.loadImgs(self.ids[idx])[0]

        m = np.zeros((img_info['height'], img_info['width']))

        for obj in anno:
            if obj['iscrowd']:
                rle = pycocotools.mask.frPyObjects(
                    obj['segmentation'], img_info['height'], img_info['width'])
                m += pycocotools.mask.decode(rle)
            elif obj['num_keypoints'] == 0:
                rles = pycocotools.mask.frPyObjects(
                    obj['segmentation'], img_info['height'], img_info['width'])
                for rle in rles:
                    m += pycocotools.mask.decode(rle)

        return m < 0.5

    def _init_check(self, heatmap_generator, joints_generator):
        assert isinstance(heatmap_generator, (list, tuple)), 'heatmap_generator should be a list or tuple'
        assert isinstance(joints_generator, (list, tuple)), 'joints_generator should be a list or tuple'
        assert len(heatmap_generator) == len(joints_generator), \
            'heatmap_generator and joints_generator should have same length,'\
            'got {} vs {}.'.format(
                len(heatmap_generator), len(joints_generator)
            )
        return len(heatmap_generator)


from __future__ import division
from functools import partial

import shutil
import time
import logging

import torch
import torch.nn as nn
import torch.nn.parallel
import torch.backends.cudnn as cudnn
import torch.optim
import torch.utils.data
from torch.utils.data.dataset import Dataset
import torchvision.transforms as transforms
import torchvision.datasets as datasets
import torchvision.models as models
import torch.nn.functional as F
from torch.utils.data import DataLoader

from mmcv.runner import get_dist_info
from mmcv.parallel import collate

import os
import sys
import random
from skimage import io
from PIL import Image
import numpy as np

from .loader import GroupSampler, DistributedGroupSampler, DistributedSampler
from .registry import DATASETS

@DATASETS.register_module
class LandmarkDetectDataset(Dataset):
    CLASSES = None

    def __init__(self,
                 img_path,
                 img_file,
                 bbox_file,
                 landmark_file,
                 img_size,
                 roi_plane_size=7,
                 attr_file=None):
        self.img_path = img_path
        normalize = transforms.Normalize(
            mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
        self.transform = transforms.Compose([
            transforms.RandomResizedCrop(img_size[0]),
            transforms.RandomHorizontalFlip(),
            transforms.ToTensor(),
            normalize,
        ])

        # read img names
        fp = open(img_file, 'r')
        self.img_list = [x.strip() for x in fp]

        self.img_size = img_size
        self.roi_plane_size = roi_plane_size

        # load bbox
        if bbox_file:
            self.with_bbox = True
            self.bboxes = np.loadtxt(bbox_file, usecols=(0, 1, 2, 3))
        else:
            self.with_bbox = False
            self.bboxes = None

        # load landmarks and visibility
        self.landmarks = np.loadtxt(landmark_file, dtype=np.float)

        # load attributes
        if attr_file is not None:
            self.attributes = np.loadtxt(attr_file, dtype=np.float)
        else:
            self.attributes = None

    def get_basic_item(self, idx):
        img = Image.open(os.path.join(self.img_path,
                                      self.img_list[idx])).convert('RGB')
        width, height = img.size

        # first crop image
        if self.with_bbox:
            bbox_cor = self.bboxes[idx]
            x1 = max(0, int(bbox_cor[0]) - 10)
            y1 = max(0, int(bbox_cor[1]) - 10)
            x2 = int(bbox_cor[2]) + 10
            y2 = int(bbox_cor[3]) + 10
            bbox_w = x2 - x1
            bbox_h = y2 - y1
            img = img.crop(box=(x1, y1, x2, y2))
        else:
            bbox_w, bbox_h = self.img_size[0], self.img_size[1]

        # then resize image
        img.thumbnail(self.img_size, Image.ANTIALIAS)
        img = self.transform(img)

        landmark_for_regression, vis = [], []
        landmark_for_roi_pool = []
        origin_landmark = self.landmarks[idx]

        # compute the shifted landmarks
        for i, l in enumerate(origin_landmark):
            if i % 3 == 0:  # visibility
                vis.append(l)
            else:
                if i % 3 == 1:  # x
                    l_x = max(0, l - x1)
                    l_x_for_regression = float(l_x) / bbox_w * self.img_size[0]
                    landmark_for_regression.append(l_x_for_regression)

                    l_x_for_roi_pool = float(l_x) / width * self.roi_plane_size
                    landmark_for_roi_pool.append(l_x_for_roi_pool)
                else:  # y
                    l_y = max(0, l - y1)
                    l_y_for_regression = float(l_y) / bbox_h * self.img_size[1]
                    landmark_for_regression.append(l_y_for_regression)

                    l_y_for_roi_pool = float(
                        l_y) / height * self.roi_plane_size
                    landmark_for_roi_pool.append(l_y_for_roi_pool)

        landmark_for_regression = torch.from_numpy(
            np.array(landmark_for_regression)).float()
        landmark_for_roi_pool = torch.from_numpy(
            np.array(landmark_for_roi_pool)).float()
        vis = torch.from_numpy(np.array(vis)).float()

        # load attribute
        if self.attributes is not None:
            attr = self.attributes[idx]
        else:
            attr = None

        data = {
            'img': img,
            'vis': vis,
            'landmark_for_regression': landmark_for_regression,
            'landmark_for_roi_pool': landmark_for_roi_pool
        }
        return data

    def __getitem__(self, idx):
        return self.get_basic_item(idx)

    def __len__(self):
        return len(self.img_list)


import tensorflow as tf
import argparse
import numpy as np
import os
import sys
import json
import util
from models import InferenceFCNet, InferenceLinear
from data_loader import InferenceDataset, InferenceBaselineDataset
from constants import VARIABLES_EXT

parser = argparse.ArgumentParser(description='Test a model')
parser.add_argument('config_file', help='config file for training')
parser.add_argument('output_dir', help='output location for saving models')
parser.add_argument('model_name', help='name for saved model')
parser.add_argument('--baseline', dest='baseline', action='store_true',
    help='train baseline (no cardplay features)')
parser.add_argument('--linear', dest='linear', action='store_true',
    help='train linear model')
parser.add_argument('--hist', dest='hist', action='store_true',
    help='use history module')
parser.add_argument('--sinkhorn', dest='sink', action='store_true',
    help='apply sinkhorn iteratively to outputs')

args = parser.parse_args()

with open(args.config_file) as config_file:
  config = json.load(config_file)

batch_size = config["batch_size"]
gzip = config["gzipped_data"]
cpu_workers = config["cpu_workers"]

# define dataset, input tensors and model
is_training = tf.placeholder(tf.bool, name="is_training")

model_defs = config["models"]
models = []

for model_def in model_defs:
  # baseline just means no hist in this case, but it can be any number of
  # features
  if args.baseline or not args.hist:
    dataset = InferenceBaselineDataset(batch_size=batch_size,
        workers=cpu_workers, gzip=gzip)
    (x_data, single_card_data)  = dataset.get_next(is_training)
    hist_data = None
  else:
    dataset = InferenceDataset(batch_size=batch_size, workers=cpu_workers,
        gzip=gzip)
    (x_data, hist_data, single_card_data)  = dataset.get_next(is_training)

  if args.linear:
    model = InferenceLinear(model_def["name"], config["input_shape"], x_data,
        hist_data, single_card_data, config["learning_rate"], args.hist,
        args.sink)
  else:
    model = InferenceFCNet(model_def["name"], config["input_shape"], x_data,
        hist_data, single_card_data, config["learning_rate"], args.hist,
        args.sink)

  model.dataset = dataset
  model.train_set = model_def["train"]
  model.test_set = model_def["test"]
  model.max_epochs = model_def["max_epochs"]
  models.append(model)

model_path = os.path.join(args.output_dir, args.model_name)
output_variables_path = os.path.join(model_path, VARIABLES_EXT)
builder = tf.saved_model.builder.SavedModelBuilder(model_path)
util.serialize_model_defs(models, model_path)
results_file = os.path.join(model_path, "train.out")

# experiment loop
with tf.Session() as sess:
  sess.run(tf.global_variables_initializer())

  builder.add_meta_graph_and_variables(sess,
      [tf.saved_model.tag_constants.TRAINING])
  builder.save()

  # for saving variables every epoch
  saver = tf.train.Saver()

  for model in models:
    util.log(results_file, "*** Training {} ***\n".format(model.name))

    prev_test_loss = np.inf
    for epoch in range(model.max_epochs):

      ep_train_losses = []
      ep_test_losses = []
      train_losses = []
      test_losses = []

      model.dataset.refresh_train(sess, model.train_set)
      feed_dict = {model.keep_prob: 0.8, is_training: True}
      ops = [model.train_step, model.loss]
      i = 0
      while True:
        try:
          (_, trl)  = sess.run(ops, feed_dict=feed_dict)
          if np.shape(trl)[0] == batch_size and i % 3 == 0:
            ep_train_losses.append(trl)
            train_losses.append(trl)
          i += 1

          global_step = sess.run(model.global_step)
          if ((global_step % config["steps_per_log"]) == 0):
            # make an additional run over the test set
            model.dataset.refresh_test(sess, model.test_set)
            old_fd = feed_dict
            old_ops = ops
            feed_dict = {model.keep_prob: 1.0, is_training: False}
            ops = [model.loss]
            j = 0
            while True and j < config["steps_per_log"]:
              try:
                tl  = sess.run(ops, feed_dict=feed_dict)[0]
                if np.shape(tl)[0] == batch_size:
                  test_losses.append(tl)
                j += 1

              except tf.errors.OutOfRangeError:
                break

            util.log(results_file, "Step={} ".format(global_step))
            util.log(results_file, "train_loss={:.9f} ".format(np.mean(train_losses)))
            util.log(results_file, "test_loss={:.9f} ".format(np.mean(test_losses)))
            util.log(results_file, "\n")
            train_losses = []
            test_losses = []

            # restore feed dict and ops
            feed_dict = old_fd
            ops = old_ops

        except tf.errors.OutOfRangeError:
          break

      model.dataset.refresh_test(sess, model.test_set)
      feed_dict = {model.keep_prob: 1.0, is_training: False}
      ops = [model.loss]
      while True:
        try:
          tl  = sess.run(ops, feed_dict=feed_dict)[0]
          if np.shape(tl)[0] == batch_size:
            ep_test_losses.append(tl)

        except tf.errors.OutOfRangeError:
          break

      avg_test_loss = np.mean(ep_test_losses)
      if avg_test_loss < prev_test_loss:
        prev_test_loss = avg_test_loss
        # only save when the model has improved
        saver.save(sess, output_variables_path)

      # Display logs per epoch step
      util.log(results_file, "Epoch={} ".format(epoch+1))
      util.log(results_file, "train_loss={:.9f} ".format(np.mean(ep_train_losses)))
      util.log(results_file, "test_loss={:.9f} ".format(np.mean(ep_test_losses)))
      util.log(results_file, "\n")


import os
import shutil
from glob import glob
from datetime import datetime
import numpy as np
from scipy import stats
from statistics import median
from collections import Counter
from pathlib import Path
from abc4pwm.non_dbd_clustering import non_dbd_ClusteringPwm
from abc4pwm.searching import searching
from abc4pwm.pwm_representative_test import make_representative_pwm
from abc4pwm.convert_count_to_pwm import motif_weight2p
from abc4pwm.energy_to_p import read_energy_matrix
from abc4pwm.similarity_score import compute_similarity_score4alignment
from itertools import takewhile
from distutils.dir_util import copy_tree

class PredictedCompare():
    def __init__(self, path_to_predicted_files, output_folder, db_folder, dst_for_bad_pwms, min_pwms_in_cluster=3, mean_threshold=0.75, z_score_threshold=-0.9,
                 top_occurrence=0.35, occurrences_threshold=0.25, ic_for_rep=0.4,top_n=2, db_type='folder',input_count=1,db_count=0, db_file_type='.txt',
                     input_file_type='.mlp', input_prob=0, db_prob=1, tf_name='', qa = 1, seed = 0, damp=0.5, max_iter=400, convergence_iter=30, preference=None):
        """
        :param path_to_predicted_files: path to folders of predicted pwms
        :param output_folder: output folder path where results of analysis file and matched TF will be stored
        :param min_pwms_in_cluster: minimum number of pwms in a cluster at clustering step
        :param db_folder: path to folder of database which contains files with which you want to match predicted files.
        :param top_n: define top n(top matches from database) for your predicted file in searching step
        :param db_type: type of database: Options are folder, path(in case of abc4pwm setting)
        :param input_count: True if input file contains counts
        :param db_count: True if database file contains counts
        :param db_file_type: mention the extension of file type. e.g., .mlp, .txt
        :param input_file_type: mention the extension of file type. e.g., .mlp, .txt
        :param input_prop: True if input matrix is already in probability
        :param db_prob: True if database matrix is already in probability
        :param tf_name: TF name for which results are compared with.
        ;param qa: quality assesment paramter. qa = 1 for doing quality assesment, otherwise zero.
        :param mean_threshold: minimum mean of similarities to qualify as good cluster
        :param z_score_threshold: minimum negative z score for similarity of a pwm from mean
        :param top_occurrence: top occurrences of pwms which qualify z score criteria
        :param occurrence_threshold: pick occurrences threshold to declare unknown/wrongly clustered pwms
        :param damp: Damping factor (between 0.5 and 1) is the extent to which the current value is maintained relative to incoming values (weighted 1 - damping).
        This in order to avoid numerical oscillations when updating these values (messages).
        :param max_iter: default=200, Maximum number of iterations.
        :param convergence_iter: default=15 Number of iterations with no change in the number of estimated clusters that stops the convergence.
        :param preference: default=None Preferences for each point - points with larger values of preferences are more likely to be chosen as exemplars.
        The number of exemplars, ie of clusters, is influenced by the input preferences value. If the preferences are not passed as arguments, they will be set to the median of the input similarities.

        """

        output_folder = os.path.join(output_folder,'')
        path_to_predicted_files = os.path.join(path_to_predicted_files,'')
        db_folder = os.path.join(db_folder,'')
        temp2_folder = Path(output_folder)
        # temp2_folder = temp_folder.parent
        path_to_text_report = os.path.join(temp2_folder,'reports_in_text/')
        if not os.path.exists(path_to_text_report):
            os.makedirs(path_to_text_report, exist_ok=True)

        with open(os.path.join(path_to_text_report,'parameters.txt'),  'w' ) as fp:
            now = datetime.now()

            current_time = now.strftime("%H:%M:%S")
            fp.writelines("Following parameters are used for insemble investigate run at " + str(current_time) + '\n\n')
            fp.writelines(str(path_to_predicted_files) + " : Path to predicted files \n")
            fp.writelines(str(output_folder) + " : Output folder \n")
            fp.writelines(str(db_folder) + " : Database folders \n")
            fp.writelines(str(dst_for_bad_pwms) + " : Bad pwms folder \n")
            fp.writelines(str(min_pwms_in_cluster) + " : Minimum Pwms in a cluster for results of ensembling learning and quality cluster \n")

            fp.writelines(str(z_score_threshold) + " : Z_score \n")
            fp.writelines(str(mean_threshold) + " : Mean\n")
            fp.writelines(str(top_occurrence) + " : Top Occurances\n")
            fp.writelines(str(occurrences_threshold) + " : Occurances Threshold\n")

            fp.writelines(str(ic_for_rep) + " : ic for rep \n")
            fp.writelines(str(top_n) + " : Number of top mathces for search results \n")
            fp.writelines(str(db_type) + " : Type of database \n")
            fp.writelines(str(input_count) + " : Input is count data? \n")
            fp.writelines(str(db_count) + " : Database is count data? \n")

            fp.writelines(str(seed) + " : Seed \n")
            fp.writelines(str(damp) + " : Damp \n")
            fp.writelines(str(max_iter) + " : Maximum iterations for Affinity Propagation \n")
            fp.writelines(str(convergence_iter) + " : Convergence Iterations \n")
            fp.writelines(str(preference) + " : Preference for Affinity Propagation \n")

        dst_folder = os.path.join(temp2_folder, 'in/')
        out_path_for_qa_clusters = os.path.join(temp2_folder, 'quality_assessed_out/')
        if not os.path.exists(out_path_for_qa_clusters):
            os.makedirs(out_path_for_qa_clusters, exist_ok=True)
        self.empty_dir(out_path_for_qa_clusters)
        if not os.path.exists(dst_for_bad_pwms):
            os.makedirs(dst_for_bad_pwms, exist_ok=True)
        if not os.path.exists(dst_folder):
            os.makedirs(dst_folder, exist_ok=True)
        self.empty_dir(dst_folder)
        self.empty_dir(dst_for_bad_pwms)
        self.prepare_for_clustering(path_to_predicted_files, dst_folder)
        print("step 1: Prepare for clustering Done")
        cluster_out_folder = os.path.join(temp2_folder,'clustering_out/')
        self.do_clustering(dst_folder, cluster_out_folder, seed,  damp, max_iter, convergence_iter, preference)
        print("step 2: Clustering Done")

        if qa:
            print("Going through qa now")
            self.quality_assessment(cluster_out_folder,
                                    out_path_for_qa_clusters,
                                    dst_for_bad_pwms,
                                    path_to_text_reports=path_to_text_report,
                                    mean_threshold=mean_threshold,
                                    z_score_threshold=z_score_threshold,
                                    top_occurrence=top_occurrence,
                                    occurrence_threshold=occurrences_threshold)

        self.make_rep(os.path.join(out_path_for_qa_clusters, 'out/'), dbd='selected', clusters='all', ic=ic_for_rep,
                 best_match_initial_motif=1, mean_threshold=mean_threshold, z_score_threshold=z_score_threshold,
                 top_occurrence=top_occurrence, occurrences_threshold=occurrences_threshold)
        predicted_folder = os.path.join(temp2_folder,'predicted')
        self.ensembling_compare_investigate(os.path.join(out_path_for_qa_clusters, 'out/'),
                                       predicted=predicted_folder, min_threshold_for_number_of_pwms=min_pwms_in_cluster)
        print("step 3: Representative Motifs Done")

        self.do_searching(in_folder=predicted_folder,
                     db_folder=db_folder,
                     out_folder=output_folder,
                     db_type=db_type,
                     input_count=input_count,
                     n=top_n,
                     db_count=db_count,
                     db_file_type=db_file_type,
                     input_file_type=input_file_type,
                     input_prob=input_prob,
                     db_prob=db_prob,
                     tf_name=tf_name,
                     path_to_text_reports=path_to_text_report,
                      mean_threshold=mean_threshold,
                      z_score_threshold=z_score_threshold,
                      top_occurrence=top_occurrence,
                      occurrence_threshold=occurrences_threshold,
                      min_threshold_for_number_of_pwms=min_pwms_in_cluster,
                          seed=seed)
        print("step 4: Searching Done")

    def prepare_for_clustering(self,src, dst):
        if not os.path.exists(dst):
            os.makedirs(dst, exist_ok=True)
        self.empty_dir(dst)
        pwms = glob(src+"/*.mlp")
        pwms = sorted(pwms)
        for pwm in pwms:
            shutil.copy(pwm,dst)
    def do_clustering(self, in_folder, out_folder, seed, damp, max_iter, convergence_iter, preference):
        non_dbd_ClusteringPwm(in_folder,out_folder, seed, damp, max_iter, convergence_iter, preference)

    def make_rep(self, path_to_clusters,  dbd = 'selected', clusters = 'all' , ic = 0.4 , best_match_initial_motif = 1, mean_threshold= 0.75, z_score_threshold = -0.9,
                     top_occurrence = 0.35, occurrences_threshold=0.25):
        make_representative_pwm(path_to_clusters,  dbd, clusters, ic , best_match_initial_motif, mean_threshold, z_score_threshold ,
                     top_occurrence, occurrences_threshold)

    def empty_dir(self, folder):
        # function for deleting files from a folder

        for filename in os.listdir(folder):
            file_path = os.path.join(folder, filename)
            try:
                if os.path.isfile(file_path) or os.path.islink(file_path):
                    os.unlink(file_path)
                elif os.path.isdir(file_path):
                    shutil.rmtree(file_path)
            except Exception as e:
                print(('Failed to delete %s. Reason: %s' % (file_path, e)))
    def ensembling_compare_investigate(self, path_of_clusters= 'out/clustering_out/out/',predicted = '', min_threshold_for_number_of_pwms= 3):

        clusters_list = [i for i in sorted(os.listdir(path_of_clusters)) if os.path.isdir(os.path.join(path_of_clusters, i))]
        clusters_list = [int(x) for x in clusters_list]
        clusters_list.sort()
        if not os.path.exists(predicted):
            os.makedirs(predicted)
        self.empty_dir(predicted)
        for cluster in clusters_list:
            pwms = glob(os.path.join(path_of_clusters,str(cluster))+"/*.mlp")
            pwms = sorted(pwms)
            if len(pwms) < int(min_threshold_for_number_of_pwms):
                if not os.path.exists(predicted):
                    os.makedirs(predicted)

            else:
                src = os.path.join(path_of_clusters,str(cluster),"repres",str(cluster)+'_rep.mlp')
                dst = os.path.join(predicted,pwms[0].split('/')[-1]+'_'+str(cluster)+'_rep.mlp')
                shutil.copy(src,dst)

    def do_searching(self, in_folder,db_folder, out_folder, tf_name ,n,db_type ,input_count,db_count, db_file_type,
                     input_file_type, input_prob, db_prob, path_to_text_reports, mean_threshold,
                                    z_score_threshold, top_occurrence, occurrence_threshold,min_threshold_for_number_of_pwms , seed):
        pwms = glob(in_folder+"/*.mlp")
        pwms = sorted(pwms)
        all_matched_list = []
        all_score_list = []
        if not os.path.exists(out_folder):
            os.makedirs(out_folder, exist_ok=True)

        with open(os.path.join(path_to_text_reports, "ensemble_summary.txt"), 'w') as f:
            f.writelines("Quality assessment is done for parameters provided below: \n\n")
            f.writelines(str(min_threshold_for_number_of_pwms) + " : Minimum threshold of PWMs in a cluster \n")
            f.writelines(str(n) + " : Top N matches\n")
            f.writelines(str(in_folder) + " : Input Folder for search..\n")
            f.writelines(str(out_folder) + " : Out folder for search..\n")
            f.writelines(str(db_folder) + " : Database folder for search results..\n")

            f.writelines(str(seed) + " : Seed \n")
            f.writelines(str(z_score_threshold) + " : Z_score \n")
            f.writelines(str(mean_threshold) + " : Mean\n")
            f.writelines(str(top_occurrence) + " : Top Occurances\n")
            f.writelines(str(occurrence_threshold) + " : Occurances Threshold\n\n")
            f.writelines("Summary of Matched PWMs: \n")
            for ind,pwm in enumerate(pwms):
                f.writelines("\nSearch out folder " +str(ind) )
                f.writelines("\nInput PWM: ")
                f.writelines(str(pwm).split('/')[-1] + "\n")

                new_out_folder=os.path.join(out_folder,str(ind))
                if not os.path.exists(new_out_folder):
                    os.makedirs(new_out_folder, exist_ok=True)
                temp_item_list, temp_list_score = searching(pwm=pwm, out=new_out_folder, tf_name=tf_name, db_path=db_folder, db_type=db_type,
                         input_count=input_count,
                         db_count=db_count,
                         db_file_type=db_file_type,
                         input_file_type=input_file_type,
                         input_prop=input_prob, n=n,
                         db_prob=db_prob)

                f.writelines("\nMatched PWMs : \n")
                for ind, file in enumerate(temp_item_list):
                    f.writelines(str(file)+"\t"+ str(temp_list_score[ind]) +"\n")
                all_matched_list.append(temp_item_list)
                all_score_list.append(temp_list_score)

            flat_list_items = [item for sublist in all_matched_list for item in sublist]
            flat_list_scores = [item for sublist in all_score_list for item in sublist]

            f.writelines(str(Counter(flat_list_items))+"\n")
            f.writelines("Median: " + str(np.median(flat_list_scores)))

            print(Counter(flat_list_items))
            # print(flat_list_scores)
            print("Median: " , np.median(flat_list_scores))
            print("See summary of ensemble and search in ", os.path.join(path_to_text_reports,"ensemble_summary.txt"))

    def get_items_upto_count(self, dct, n,top_percentage):
      data = dct.most_common(top_percentage)
      #Now collect all items whose value is greater than `n`.
      return list(takewhile(lambda x: x[1] > n, data))

    def quality_assessment(self, cluster_path, out_path_for_qa_clusters ,dst_for_bad_pwms, path_to_text_reports , mean_threshold, z_score_threshold, occurrence_threshold, top_occurrence):

        copy_tree(cluster_path, out_path_for_qa_clusters)

        out_path_for_qa_clusters = os.path.join(out_path_for_qa_clusters,'out')
        clusters_list = [i for i in sorted(os.listdir(out_path_for_qa_clusters)) if os.path.isdir(os.path.join(out_path_for_qa_clusters, i))]
        clusters_list = [int(x) for x in clusters_list]
        clusters_list.sort()
        if not os.path.exists(path_to_text_reports):
            os.makedirs(path_to_text_reports, exist_ok=True)
        with open(os.path.join(path_to_text_reports, 'ensemble_cluster_quality.txt'), 'w') as f:

            f.writelines(
                'Cluster Number '
                '\t Mean '
                '\t Median '
                '\t Minimum '
                '\t Standard Deviation '
                '\t Total No. of PWMs '
                '\t No. of Wrong PWMs '
                '\t Percentage of Quality '
                '\t Wrongly clustered PWMS \n')
            for cluster in clusters_list:

                pwms_full = glob(os.path.join(out_path_for_qa_clusters, str(cluster) + "/*.mlp"))
                pwms_full = sorted(pwms_full)

                similarityMatrix = np.zeros((len(pwms_full), len(pwms_full)))

                for index1, i in enumerate(pwms_full):
                    matrix1, matrix_string1, maximum_feq1, total_maximum1, info1 = read_energy_matrix(i)
                    normalized_matrix1 = motif_weight2p(matrix1)
                    for index2, j in enumerate(pwms_full):
                        matrix2, matrix_string2, maximum_feq2, total_maximum2, info2 = read_energy_matrix(j)

                        normalized_matrix2 = motif_weight2p(matrix2)
                        similarityMatrix[index1, index2] = compute_similarity_score4alignment(normalized_matrix1,
                                                                                normalized_matrix2)

                uper_triangle_similarity_matrix = similarityMatrix[np.triu_indices(similarityMatrix.shape[0], k=1)]
                uper_triangle_similarity_matrix_indeces = np.triu_indices(similarityMatrix.shape[0], k=1)

                # minimum_similarity_pwm_index = np.unravel_index(np.argmin(sum(similarityMatrix), axis=None),
                #                                                 sum(similarityMatrix).shape)

                if len(pwms_full) > 1:
                    if uper_triangle_similarity_matrix.mean() > mean_threshold:
                        potential_bad_pwms = []
                        for index, k in enumerate(stats.zscore(uper_triangle_similarity_matrix)):
                            if k < z_score_threshold:
                                potential_bad_pwms.append(uper_triangle_similarity_matrix_indeces[0][index])
                                potential_bad_pwms.append(uper_triangle_similarity_matrix_indeces[1][index])

                        bad_pwms_indexes = self.get_items_upto_count(Counter(potential_bad_pwms),
                                                                     int(np.ceil(len(
                                                                         pwms_full) * occurrence_threshold)),
                                                                     int(np.ceil(
                                                                         len(pwms_full) * top_occurrence)))

                        del_from_pwms_list = [i[0] for i in bad_pwms_indexes]
                        bad_pwms = [pwms_full[x] for x in del_from_pwms_list]

                        f.writelines(str(cluster) +
                                     "\t" + str(uper_triangle_similarity_matrix.mean()) +
                                     "\t" + str(median(uper_triangle_similarity_matrix)) +
                                     "\t" + str(min(uper_triangle_similarity_matrix)) + " " +
                                     "\t" + str(uper_triangle_similarity_matrix.std()) +
                                     "\t" + str(len(pwms_full)) +
                                     "\t" + str(len(bad_pwms)) +
                                     "\t" + str(int(((len(pwms_full) - len(bad_pwms)) / len(pwms_full)) * 100)))
                        if len(bad_pwms) > 0:
                            f.writelines("%\t" + str([i[0] for i in bad_pwms_indexes]) + "\n")
                            # total_bad_pwms += len(bad_pwms)
                            for x in bad_pwms:
                                # pwms.remove(x)
                                shutil.move(x, dst_for_bad_pwms)
                        else:
                            f.writelines("\n")
                            #print("Nothing")

if __name__ == '__main__':

    src_folder = '../data/out/ensemble/time_series/t0/0strand/'
    out = '../data/out/ensemble/time_series/t0/0strand/search_out/'
    dest_4_bd_pwms = '../data/out/ensemble/time_series/t0/0strand/bad_pwms/'
    db_folder = '../data/in/in_pwms/'
    db_file_type = '.mlp'
    db_count = 1
    db_prob = 0
    input_count = 1
    input_file_type = '.mlp'
    input_prob = 0
    min_pwms_in_cluster= 5
    top_n=2
    seed= 1
    ic_for_rep = 0
    mean_threshold=0.5
    z_score_threshold=0.0
    damp=0.7
    max_iter=1000
    convergence_iter=30
    top_occurrences=0.5
    occurrences_threshold=0.5
    db_type = 'folder'

    PredictedCompare(path_to_predicted_files=src_folder,
                     dst_for_bad_pwms=dest_4_bd_pwms,
                     output_folder=out,
                     db_folder=db_folder,
                     db_file_type = db_file_type,
                     db_count=db_count,
                     db_prob=db_prob,
                     tf_name='',
                     qa = 1,
                     seed=seed ,
                     ic_for_rep=ic_for_rep,
                     min_pwms_in_cluster=min_pwms_in_cluster, mean_threshold=mean_threshold, z_score_threshold=z_score_threshold,
                     top_occurrence=top_occurrences,
                     occurrences_threshold=occurrences_threshold,
                     top_n=top_n,
                     db_type=db_type,
                     input_count=input_count,
                     input_prob=input_prob,
                     input_file_type=input_file_type,
                     damp=damp,
                     max_iter=max_iter,
                     convergence_iter=convergence_iter,
                     )


import os
import sys
from PIL import Image,ImageFilter,ImageDraw
import numpy as np
import time
import random
import math
import json
import re

def _get_screen_size():
    size_str = os.popen('adb shell wm size').read()
    m = re.search('(\d+)x(\d+)', size_str)
    if m:
        width = m.group(1)
        height = m.group(2)
        return "{height}x{width}".format(height=height, width=width)

def open_accordant_config():
    screen_size = _get_screen_size()
    config_file = "{path}/config/{screen_size}/config.json".format(
        path=sys.path[0],
        screen_size=screen_size
    )
    if os.path.exists(config_file):
        with open(config_file, 'r') as f:
            # print("Load config file from {}".format(config_file))
            return json.load(f)
    else:
        with open('{}/config/default.json'.format(sys.path[0]), 'r') as f:
            # print("Load default config")
            return json.load(f)

config = open_accordant_config()

# Magic Number，不设置可能无法正常执行，请根据具体截图从上到下按需设置
under_game_score_y = config['under_game_score_y']
press_coefficient = config['press_coefficient']       # 长按的时间系数，请自己根据实际情况调节
piece_base_height_1_2 = config['piece_base_height_1_2']   # 二分之一的棋子底座高度，可能要调节
piece_body_width = config['piece_body_width']             # 棋子的宽度，比截图中量到的稍微大一点比较安全，可能要调节

def get_screenshot():
    os.system('adb shell screencap -p /sdcard/autojump.png')
    os.system('adb pull /sdcard/autojump.png .')

def game_over(im):
    pixels = im.load()
    point = pixels[0,0]
    print(point)
    if (30<point[0]<55) and (30<point[1]<55) and (30<point[2]<55):
        return True
    return False

def find_piece_x(im):
    w, h = im.size

    piece_x_sum = 0
    piece_x_c = 0
    piece_y_max = 0
    scan_x_border = int(w / 8)  # 扫描棋子时的左右边界
    scan_start_y = 0  # 扫描的起始y坐标
    im_pixel=im.load()
    # 以50px步长，尝试探测scan_start_y
    for i in range(int(h / 3), int( h*2 /3 ), 50):
        last_pixel = im_pixel[0,i]
        for j in range(1, w):
            pixel=im_pixel[j,i]
            # 不是纯色的线，则记录scan_start_y的值，准备跳出循环
            if pixel[0] != last_pixel[0] or pixel[1] != last_pixel[1] or pixel[2] != last_pixel[2]:
                scan_start_y = i - 50
                break
        if scan_start_y:
            break
    # print('scan_start_y: ', scan_start_y)

    # 从scan_start_y开始往下扫描，棋子应位于屏幕上半部分，这里暂定不超过2/3
    for i in range(scan_start_y, int(h * 2 / 3)):
        for j in range(scan_x_border, w - scan_x_border):  # 横坐标方面也减少了一部分扫描开销
            pixel = im_pixel[j,i]
            # 根据棋子的最低行的颜色判断，找最后一行那些点的平均值，这个颜色这样应该 OK，暂时不提出来
            if (50 < pixel[0] < 60) and (53 < pixel[1] < 63) and (95 < pixel[2] < 110):
                piece_x_sum += j
                piece_x_c += 1
                piece_y_max = max(i, piece_y_max)

    if not all((piece_x_sum, piece_x_c)):
        return 0
    piece_x = piece_x_sum / piece_x_c
    return int(piece_x)

def find_board_x(im):
    w, h = im.size
    conF = im.filter(ImageFilter.CONTOUR)
    L = conF.convert('L')
    # L.show()

    for j in range(int(h / 3),int(h * 2 / 3)):
        for i in range(w):
            r = L.getpixel((i,j))
            if r<200:
                # print(i,j)
                return i
    return 0
def find_piece_and_board(im):
    w, h = im.size

    piece_x_sum = 0
    piece_x_c = 0
    piece_y_max = 0
    board_x = 0
    board_y = 0
    scan_x_border = int(w / 8)  # 扫描棋子时的左右边界
    scan_start_y = 0  # 扫描的起始y坐标
    im_pixel=im.load()
    # 以50px步长，尝试探测scan_start_y
    for i in range(int(h / 3), int( h*2 /3 ), 50):
        last_pixel = im_pixel[0,i]
        for j in range(1, w):
            pixel=im_pixel[j,i]
            # 不是纯色的线，则记录scan_start_y的值，准备跳出循环
            if pixel[0] != last_pixel[0] or pixel[1] != last_pixel[1] or pixel[2] != last_pixel[2]:
                scan_start_y = i - 50
                break
        if scan_start_y:
            break
    # print('scan_start_y: ', scan_start_y)

    # 从scan_start_y开始往下扫描，棋子应位于屏幕上半部分，这里暂定不超过2/3
    for i in range(scan_start_y, int(h * 2 / 3)):
        for j in range(scan_x_border, w - scan_x_border):  # 横坐标方面也减少了一部分扫描开销
            pixel = im_pixel[j,i]
            # 根据棋子的最低行的颜色判断，找最后一行那些点的平均值，这个颜色这样应该 OK，暂时不提出来
            if (50 < pixel[0] < 60) and (53 < pixel[1] < 63) and (95 < pixel[2] < 110):
                piece_x_sum += j
                piece_x_c += 1
                piece_y_max = max(i, piece_y_max)

    if not all((piece_x_sum, piece_x_c)):
        return 0, 0, 0, 0
    piece_x = piece_x_sum / piece_x_c
    piece_y = piece_y_max - piece_base_height_1_2  # 上移棋子底盘高度的一半

    for i in range(int(h / 3), int(h * 2 / 3)):
        last_pixel = im_pixel[0, i]
        if board_x or board_y:
            break
        board_x_sum = 0
        board_x_c = 0

        for j in range(w):
            pixel = im_pixel[j,i]
            # 修掉脑袋比下一个小格子还高的情况的 bug
            if abs(j - piece_x) < piece_body_width:
                continue

            # 修掉圆顶的时候一条线导致的小 bug，这个颜色判断应该 OK，暂时不提出来
            if abs(pixel[0] - last_pixel[0]) + abs(pixel[1] - last_pixel[1]) + abs(pixel[2] - last_pixel[2]) > 10:
                board_x_sum += j
                board_x_c += 1
        if board_x_sum:
            board_x = board_x_sum / board_x_c
    # 按实际的角度来算，找到接近下一个 board 中心的坐标 这里的角度应该是30°,值应该是tan 30°, math.sqrt(3) / 3
    board_y = piece_y - abs(board_x - piece_x) * math.sqrt(3) / 3

    if not all((board_x, board_y)):
        return 0, 0, 0, 0

    return piece_x, piece_y, board_x, board_y
def getMaxQ(state,Q):
    return max(Q[state,:])

def getMaxQ_index(Q):
    index = np.argwhere(Q == max(Q[:]))
    # print(index[0][0])
    return index[0][0]

def q_learning(Q,GAMMA,init_action):
    # for i in range(episode):
    while True:

        get_screenshot()
        im = Image.open('./autojump.png')
        piece_x, piece_y, board_x, board_y = find_piece_and_board(im)
        dist = int(math.sqrt((board_x - piece_x) ** 2 + (board_y - piece_y) ** 2))
        press_time = 0
        # if dist%10 > 4:
        #     distance = int(dist/10) + 1
        # else:
        distance = int(dist/10)
        vaild_action = []
        for i in range(6, 41):
            if Q[distance][i] != -50:
                vaild_action.append(i)
        # print(vaild_action)
        if getMaxQ(distance,Q) == 0:
            if vaild_action != []:
                press_time = random.choice(vaild_action) * 25
            else:
                press_time = random.choice(init_action) * 25
        else:
            # print(np.argwhere(Q == getMaxQ(distance,Q)))
            # press_time = np.argwhere(Q == getMaxQ(distance,Q))[0][1] * 50
            press_time = getMaxQ_index(Q[distance]) * 25
            #print(press_time)
        # print('press_time:', press_time/50)
        # print('distance:',distance)
        os.system('adb shell input swipe 20 20 21 21 {time}'.format(time=press_time))
        time.sleep(4)
        #
        get_screenshot()
        im = Image.open('./autojump.png')
        if game_over(im):
            Q[distance][int(press_time / 25)] = -50 + GAMMA * getMaxQ(distance,Q)
            # print('score:',Q[distance][int(press_time / 50)])
            break
        else:
            Q[distance][int(press_time/25)] = 10 + GAMMA * getMaxQ(distance,Q)
            # print('score:',Q[distance][int(press_time / 50)])

def retry_button(w,h):
    left = w / 2
    top = 1003 * (h / 1280.0) + 10
    cmd = 'adb shell input swipe {x1} {y1} {x2} {y2} {duration}'.format(
        x1=left,
        y1=top,
        x2=left,
        y2=top,
        duration=100
    )
    os.system(cmd)

def train(episode,w,h):
    init_action = [i for i in range(1,41)]
    if os.path.exists("Qarray.txt"):
        Q = np.loadtxt("Qarray.txt")
        print("load Qarray.txt")
    else:
        # print(e)
        Q = np.zeros((int(w/10)+1,41))

    print(Q)
    for i in range(episode):
        print(i,'episode')
        q_learning(Q, 0.8,init_action)
        retry_button(w,h)
        time.sleep(1)
        np.savetxt("Qarray.txt",Q)

train(500,720,1280)
    # return 0
# get_screenshot()
# im = Image.open('./autojump.png')
# x = find_piece_x(im)
# i,j = find_boaOKzrd_x(im)
# draw = ImageDraw.Draw(im)
# draw.line((piece_x, piece_y) + (board_x, board_y), fill=2, width=3)
# draw.line((i, 0, i, im.size[1]), fill=(255, 0, 0))
# draw.line((0, j, im.size[0], j), fill=(255, 0, 0))
# draw.line((x, 0, x, im.size[1]), fill=(0, 0, 255))
# # draw.line((0, board_y, im.size[0], board_y), fill=(0, 0, 255))
# draw.ellipse((i - 10, j - 10, i + 10, j + 10), fill=(0, 0, 255))
# draw.ellipse((x - 10, j - 10, x + 10, j + 10), fill=(0, 0, 255))
# im.show()

# # print(int(w/10),h)

#
# print(Q)
# print(game_over(im))
# dist = abs(find_piece_x(im)-find_board_x(im))
# print(int(dist/10),dist%10)

# os.system('adb shell input swipe 20 20 21 21 500')