from serpent.game_agent import GameAgent

from serpent.enums import InputControlTypes

from serpent.frame_grabber import FrameGrabber
from serpent.input_controller import KeyboardKey

from serpent.machine_learning.reinforcement_learning.agents.rainbow_dqn_agent import RainbowDQNAgent
import pytesseract

import numpy as np

from serpent.config import config

from serpent.logger import Loggers

import serpent.cv
import signal
import time
import random

from mss import mss

from PIL import Image
import os

import cv2
import skimage

import pyautogui
import pywinauto
import keyboard

Wd, Hd = 1920, 1080
ACTIVATION_RANGE = 300
YOLO_DIRECTORY = "models"
CONFIDENCE = 0.36
THRESHOLD = 0.22
ACTIVATION_RANGE = 350
labelsPath = os.path.sep.join([YOLO_DIRECTORY, "coco-dataset.labels"])
LABELS = open(labelsPath).read().strip().split("\n")

np.random.seed(42)
COLORS = np.random.randint(0, 255, size=(len(LABELS), 3), dtype="uint8")
weightsPath = os.path.sep.join([YOLO_DIRECTORY, "yolov3-tiny.weights"])
configPath = os.path.sep.join([YOLO_DIRECTORY, "yolov3-tiny.cfg"])

net = cv2.dnn.readNetFromDarknet(configPath, weightsPath)
net.setPreferableBackend(cv2.dnn.DNN_BACKEND_DEFAULT)
net.setPreferableTarget(cv2.dnn.DNN_TARGET_OPENCL)
ln = net.getLayerNames()
ln = [ln[i[0] - 1] for i in net.getUnconnectedOutLayers()]
origbox = (int(Wd/2 - ACTIVATION_RANGE/2),
    int(Hd/2 - ACTIVATION_RANGE/2),
    int(Wd/2 + ACTIVATION_RANGE/2),
    int(Hd/2 + ACTIVATION_RANGE/2))
def set_pos(x, y):
    current_x, current_y = pyautogui.position()
    new_x = current_x + x
    new_y = current_y + y
    pywinauto.mouse.move(coords=(new_x, new_y))
def set_pos_aimbot(x, y):
    pywinauto.mouse.move(coords=(x, y))
class SerpentCODGameAgent(GameAgent):

    def __init__(self, **kwargs):
        super().__init__(**kwargs)

        self.frame_handlers["PLAY"] = self.handle_play
        self.frame_handler_setups["PLAY"] = self.setup_play
        self.frame_handler_pause_callbacks["PLAY"] = self.handle_play_pause

    def setup_play(self):
        self.environment = self.game.environments["GAME"](
            game_api=self.game.api,
            input_controller=self.input_controller,
            episodes_per_startregions_track=100000000000
        )

        self.game_inputs = [
            {
                "name": "CONTROLS",
                "control_type": InputControlTypes.DISCRETE,
                "inputs": self.game.api.combine_game_inputs(["MOVEMENT", "COMBAT", "CURSOR"])
            }
        ]
        self.agent = RainbowDQNAgent(
            "COD",
            game_inputs=self.game_inputs,
            callbacks=dict(
                after_observe=self.after_agent_observe,
                before_update=self.before_agent_update,
                after_update=self.after_agent_update
                ),
            rainbow_kwargs=dict(
                replay_memory_capacity=250000,
                observe_steps=100,
                batch_size=10,
                save_steps=300,
                model="datasets/rainbow_dqn_COD.pth"
                ),
            logger=Loggers.COMET_ML,
            logger_kwargs=dict(
                api_key="api_key_from_comet_ml",
                project_name="serpent-ai-cod",
                reward_func=self.reward
            )
            )
        self.analytics_client.track(event_key="COD", data={"name": "COD"})
        self.agent.logger.experiment.log_other("game", "COD")
        self.environment.new_episode(maximum_steps=350)  # 5 minutes
        self.overs = 0
        self.input_non_lethal = False
    def handle_play(self, game_frame, game_frame_pipeline):
        self.paused_at = None
        with mss() as sct:
            monitor_var = sct.monitors[1]
            monitor = sct.grab(monitor_var)
            valid_game_state = self.environment.update_startregions_state(monitor)
        if not valid_game_state:
            return None

        reward, over_boolean = self.reward(self.environment.startregions_state, 1.0)
        terminal = over_boolean

        self.agent.observe(reward=reward, terminal=terminal)

        if not terminal:
            game_frame_buffer = FrameGrabber.get_frames([0, 2, 4, 6], frame_type="PIPELINE")
            agent_actions = self.agent.generate_actions(game_frame_buffer)
            print("Current Action: ")
            print(agent_actions)
            str_agent_actions = str(agent_actions)
            if "MOVE MOUSE X" in str_agent_actions:
                set_pos(200, 0)
            if "MOVE MOUSE Y" in str_agent_actions:
                set_pos(0, 200)
            if "MOVE MOUSE XY" in str_agent_actions:
                set_pos(100, 100)
            if "MOVE MOUSE X2" in str_agent_actions:
                set_pos(-200, 0)
            if "MOVE MOUSE Y2" in str_agent_actions:
                set_pos(0, -200)
            if "MOVE MOUSE XY2" in str_agent_actions:
                set_pos(-100, -100)
            if "MOVE MOUSE XY3" in str_agent_actions:
                set_pos(-100, 100)
            if "MOVE MOUSE XY4" in str_agent_actions:
                set_pos(100, -100)
            if "LETHAL" in str_agent_actions:
                self.input_non_lethal = True
            self.human()
            self.environment.perform_input(agent_actions)
        else:
            self.environment.clear_input()
            self.agent.reset()

            time.sleep(30)
            #To Do
            #Choose Loadout (Meduim Range)
            self.environment.end_episode()
            self.environment.new_episode(maximum_steps=350)
            print("New Episode")
    def handle_play_pause(self):
        self.input_controller.handle_keys([])
    def num_there(self, s):
        return any(i.isdigit() for i in s)

    def get_health(self, image):
        img = Image.frombytes('RGB', image.size, image.rgb)
        red_O = 0
        for red in img.getdata():
            if red == (117,54,34):
                red_O += 1
        return red_O
    def get_xp(self, image_xp):
        img = Image.frombytes('RGB', image_xp.size, image_xp.rgb)
        pixels = 0
        for pixel in img.getdata():
            if pixel == (255,194,21):
                pixels += 1
        return pixels
    def is_startregions_over(self, image):
        image = Image.frombytes("RGB", image.size, image.bgra, "raw", "BGRX")
        ocr_result = pytesseract.image_to_string(image, lang='eng')
        print("Text: ")
        print(ocr_result)
        if "KILLCAM" in ocr_result:
            return True
        else:
            return False
    def human(self):
        with mss() as sct:
            W, H = None, None
            frame = np.array(sct.grab(origbox))
            frame = cv2.cvtColor(frame, cv2.COLOR_RGBA2BGR)

            if W is None or H is None:
                (H, W) = frame.shape[: 2]

            frame = cv2.UMat(frame)
            blob = cv2.dnn.blobFromImage(frame, 1 / 260, (150, 150),
                swapRB=False, crop=False)
            net.setInput(blob)
            layerOutputs = net.forward(ln)
            boxes = []
            confidences = []
            classIDs = []
            for output in layerOutputs:
                for detection in output:
                    scores = detection[5:]
                    classID = 0
                    confidence = scores[classID]
                    if confidence > CONFIDENCE:
                        box = detection[0: 4] * np.array([W, H, W, H])
                        (centerX, centerY, width, height) = box.astype("int")
                        x = int(centerX - (width / 2))
                        y = int(centerY - (height / 2))
                        boxes.append([x, y, int(width), int(height)])
                        confidences.append(float(confidence))
                        classIDs.append(classID)

            idxs = cv2.dnn.NMSBoxes(boxes, confidences, CONFIDENCE, THRESHOLD)
            if len(idxs) > 0:
                bestMatch = confidences[np.argmax(confidences)]

                # loop over the indexes we are keeping
                for i in idxs.flatten():
                    # extract the bounding box coordinates
                    (x, y) = (boxes[i][0], boxes[i][1])
                    (w, h) = (boxes[i][2], boxes[i][3])

                    # draw target dot on the frame
                    cv2.circle(frame, (int(x + w / 2), int(y + h / 5)), 5, (0, 0, 255), -1)

                    # draw a bounding box rectangle and label on the frame
                    # color = [int(c) for c in COLORS[classIDs[i]]]
                    cv2.rectangle(frame, (x, y),
                                    (x + w, y + h), (0, 0, 255), 2)

                    text = "TARGET {}%".format(int(confidences[i] * 100))
                    cv2.putText(frame, text, (x, y - 5),
                                cv2.FONT_HERSHEY_SIMPLEX, 0.5, (0, 255, 0), 2)

                    if bestMatch == confidences[i]:
                        mouseX = origbox[0] + (x + w/1.5)
                        mouseY = origbox[1] + (y + h/5)
                        mouseX = int(round(mouseX))
                        mouseY = int(round(mouseY))
                        set_pos_aimbot(mouseX, mouseY)
                        pywinauto.mouse.click(button='left', coords=(mouseX, mouseY))
                        pywinauto.mouse.click(button='left', coords=(mouseX, mouseY))
                        pywinauto.mouse.click(button='left', coords=(mouseX, mouseY))
                        pywinauto.mouse.click(button='left', coords=(mouseX, mouseY))
                        pywinauto.mouse.click(button='left', coords=(mouseX, mouseY))
                        pywinauto.mouse.click(button='left', coords=(mouseX, mouseY))
    def reward(self, game_state, object_reward_func):
        with mss() as sct:
            image = sct.grab(sct.monitors[1])
            value = self.get_health(image)
            print("Health: ")
            print(value * -1)
            monitor = {"top": 452, "left": 1000, "width": 144, "height": 51, "mon": 1}
            image_xp = sct.grab(monitor)
            xp = self.get_xp(image_xp)
            monitor_custom_game = {"top": 47, "left": 50, "width": 230, "height": 66, "mon": 1}
            image_over = sct.grab(monitor_custom_game)
            over_check = self.is_startregions_over(image_over)
            reward = 0.0
            over = False
            if over_check:
                reward = 0.0
                self.overs += 1
                if self.overs >= 6:
                    print("Game Over")
                    over = True
                    self.overs = 0
                else:
                    over = False
            else:
                reward = 0.0
                if value >= 1:
                    reward += -1.5
                elif value < 1:
                    reward += 0.0
                if value >= 1 and self.input_non_lethal:
                    reward += 2.5
                    self.input_non_lethal = False
                if xp >= 7:
                    reward += 3.0
            print("Reward: ")
            print(reward)
            return reward, over
    def after_agent_observe(self):
        self.environment.episode_step()

    def before_agent_update(self):
        self.input_controller.tap_key(KeyboardKey.KEY_ESCAPE)
        time.sleep(1)

    def after_agent_update(self):
        self.input_controller.tap_key(KeyboardKey.KEY_ESCAPE)
        time.sleep(3)


from __future__ import absolute_import
from __future__ import division
from __future__ import with_statement

import os

import keras
import numpy as np
from keras.layers import Flatten, GlobalMaxPooling1D
from keras.layers import Input, merge
from keras.layers.convolutional import Convolution1D
from keras.layers.core import Dense, Dropout
from keras.layers.embeddings import Embedding
from keras.models import Model
from nlplingo.sandbox.model import EventExtractionModel
from nlplingo.sandbox.model import MyRange

global keras_pair_model

class PairModel(EventExtractionModel):
    def __init__(self, params, event_domain, embeddings, causal_embeddings):
        """
        :type event_domain: nlplingo.event.event_domain.EventDomain
        :type embeddings: nlplingo.embeddings.word_embeddings.WordEmbedding
        :type causal_embeddings: nlplingo.embeddings.word_embeddings.WordEmbedding
        """
        super(PairModel, self).__init__(params, event_domain, embeddings, causal_embeddings,
                                           batch_size=params.get_int('pair.batch_size'),
                                           num_feature_maps=params.get_int('pair.num_feature_maps'))
        self.num_output = 2
        self.positive_weight = params.get_float('pair.positive_weight')
        self.epoch = params.get_int('pair.epoch')
        #self.existing_classes = read_file_to_set(params.get_string('existing_types'))
        #self.new_classes = read_file_to_set(params.get_string('new_types'))           # novel tasks types

    def fit(self, train_data_list, train_label, test_data_list, test_label):
        global keras_pair_model

        if self.verbosity == 1:
            print('- train_data_list=', train_data_list)
            print('- train_label=', train_label)

        none_label_index = 0
        sample_weight = np.ones(train_label.shape[0])
        #label_argmax = np.argmax(train_label, axis=1)
        label_argmax = train_label
        for i, label_index in enumerate(label_argmax):
            if label_index != none_label_index:
                sample_weight[i] = self.positive_weight

        #super(TriggerModel, self).fit(train_label, train_data_list, test_label, test_data_list,
        #            sample_weight=sample_weight, max_epoch=self.epoch)
        history = keras_pair_model.fit(train_data_list, train_label,
                                  sample_weight=sample_weight, batch_size=self.batch_size, nb_epoch=self.epoch)
                                  #validation_data=(test_data_list, test_label))
        return history

    def load_keras_model(self, filename=None):
        global keras_pair_model
        keras_pair_model = keras.models.load_model(filename, self.keras_custom_objects)

    def save_keras_model(self, filename):
        global keras_pair_model
        keras_pair_model.save(filename)

    def predict(self, test_data_list):
        global keras_pair_model

        try:
            pred_result = keras_pair_model.predict(test_data_list)
        except:
            self.load_keras_model(filename=os.path.join(self.model_dir, 'pair.hdf'))
            print('*** Loaded keras_pair_model ***')
            pred_result = keras_pair_model.predict(test_data_list)
        return pred_result

class MaxPoolEmbeddedPairModel(PairModel):
    def __init__(self, params, event_domain, embeddings, causal_embeddings):
        """
        :type params: nlplingo.common.parameters.Parameters
        :type event_domain: nlplingo.event.event_domain.EventDomain
        :type embeddings: nlplingo.embeddings.word_embeddings.WordEmbedding
        :type causal_embeddings: nlplingo.embeddings.word_embeddings.WordEmbedding
        """
        super(MaxPoolEmbeddedPairModel, self).__init__(params, event_domain, embeddings, causal_embeddings)
        self.train_embedding = False
        self.neighbor_dist = params.get_int('cnn.neighbor_dist')
        #self.use_trigger = params.get_boolean('use_trigger')
        #self.use_role = params.get_boolean('use_role')
        #self.role_use_head = params.get_boolean('role.use_head')
        self.number_of_hidden_nodes = params.get_int('nn.hidden_nodes')
        self.use_cnn = params.get_boolean('use_cnn')
        self.use_dropout = params.get_boolean('use_dropout')
        self.create_model()

    def create_model(self):
        global keras_pair_model

        pos_input1 = Input(shape=(self.sent_length,), dtype=u'int32', name=u'position_array1')
        pos_input2 = Input(shape=(self.sent_length,), dtype=u'int32', name=u'position_array2')

        # For each word the pos_array_input defines two coordinates:
        # distance(word,anchor) and distance(word,target).  Embed each distances
        # into an 'embedding_vec_length' dimensional vector space
        # pos_input1 = Input(shape=(2, self.sent_length), dtype=u'int32', name=u'position_array1')
        # anchor_pos_array1 = MySelect(0)(pos_input1)
        # target_pos_array1 = MySelect(1)(pos_input1)
        # anchor_embedding1 = Embedding(2 * self.sent_length, self.position_embedding_vec_len)(anchor_pos_array1)
        # target_embedding1 = Embedding(2 * self.sent_length, self.position_embedding_vec_len)(target_pos_array1)
        #
        # pos_input2 = Input(shape=(2, self.sent_length), dtype=u'int32', name=u'position_array2')
        # anchor_pos_array2 = MySelect(0)(pos_input2)
        # target_pos_array2 = MySelect(1)(pos_input2)
        # anchor_embedding2 = Embedding(2 * self.sent_length, self.position_embedding_vec_len)(anchor_pos_array2)
        # target_embedding2 = Embedding(2 * self.sent_length, self.position_embedding_vec_len)(target_pos_array2)

        trigger_window = 2 * self.neighbor_dist + 1
        # if self.role_use_head:
        #     role_window = 2 * self.neighbor_dist + 1
        # else:
        #     role_window = 2 * self.neighbor_dist
        #
        # if self.use_trigger and trigger_window == 0:
        #     raise ValueError('If using eventtrigger, trigger_window must be >= 0')
        # if self.use_role and role_window == 0:
        #     raise ValueError('If using role, role_window must be >= 0 (you must either use role_head or neighbor_dist > 0)')

        # dep_input1 = Input(shape=(1,), dtype=u'int32', name=u'dep_vector1')
        # dep_input2 = Input(shape=(1,), dtype=u'int32', name=u'dep_vector2')

        # Input is vector of embedding indice representing the sentence
        context_input1 = Input(shape=(self.sent_length + trigger_window,), dtype=u'int32', name=u'word_vector1')
        context_input2 = Input(shape=(self.sent_length + trigger_window,), dtype=u'int32', name=u'word_vector2')

        # context_cinput1 = Input(shape=(self.sent_length + trigger_window + role_window,), dtype=u'int32',
        #                        name=u'word_cvector1')
        # context_cinput2 = Input(shape=(self.sent_length + trigger_window + role_window,), dtype=u'int32',
        #                        name=u'word_cvector2')

        embeddings = Embedding(self.word_embeddings.shape[0], self.word_embeddings.shape[1],
                             weights=[self.word_embeddings], trainable=self.train_embedding)
        # causal_embeddings = Embedding(self.causal_word_embeddings.shape[0], self.causal_word_embeddings.shape[1],
        #                        weights=[self.causal_word_embeddings], trainable=self.train_embedding)

        all_words1 = embeddings(context_input1)
        all_words2 = embeddings(context_input2)

        # all_cwords1 = causal_embeddings(context_cinput1)
        # all_cwords2 = causal_embeddings(context_cinput2)
        #
        # dep_embeddings = Embedding(self.word_embeddings.shape[0], self.word_embeddings.shape[1],
        #                        weights=[self.word_embeddings], trainable=self.train_embedding)
        #
        # all_dep1 = dep_embeddings(dep_input1)
        # all_dep2 = dep_embeddings(dep_input2)

        #all_words = Embedding(self.word_embeddings.shape[0], self.word_embeddings.shape[1],
        #                      weights=[self.word_embeddings], trainable=self.train_embedding)(context_input)

        context_words1 = MyRange(0, -trigger_window)(all_words1)
        context_words2 = MyRange(0, -trigger_window)(all_words2)

        # merged1 = merge([context_words1, target_embedding1], mode=u'concat')
        # merged2 = merge([context_words2, target_embedding2], mode=u'concat')

        #Convolution1D(self.num_feature_maps, self.filter_length, border_mode=u'valid')

        # Input anchor and target words, plus +/- one context words
        # lex_vector = Input(shape=(3,self.word_vec_length), name='lex')

        # Merge sentance and lexcial features
        #merged_all = merge([maxpool1, maxpool2], mode=u'concat')
#        merged_all = keras.layers.multiply([maxpool1, maxpool2])
#        print('type(merged_all)={}'.format(type(merged_all)))

#        bi_dir = Bidirectional(LSTM(100, input_shape=all_words1.shape[1:], return_sequences=True))
#        maxpool1 = bi_dir(all_words1)
#        maxpool2 = bi_dir(all_words2)

        #sim_att = custom_attn.AttentionMMA(merge_mode="concat")([maxpool1, maxpool2])
        #merged_all = BatchNormalization()(sim_att)

        lex_vector1 = MyRange(-trigger_window, None)(all_words1)
        lex_vector2 = MyRange(-trigger_window, None)(all_words2)

        # if self.use_trigger and self.use_role:
        #     lex_vector1 = MyRange(-(trigger_window + role_window), None)(all_words1)
        #     lex_vector2 = MyRange(-(trigger_window + role_window), None)(all_words2)
        #     lex_cvector1 = MyRange(-(trigger_window + role_window), None)(all_cwords1)
        #     lex_cvector2 = MyRange(-(trigger_window + role_window), None)(all_cwords2)
        #     if role_window > 0:
        #         lex_cvector1 = MyRange(0, -role_window)(lex_cvector1)
        #         lex_cvector2 = MyRange(0, -role_window)(lex_cvector2)
        # else:
        #     if self.use_role:
        #         lex_vector1 = MyRange(-role_window, None)(all_words1)
        #         lex_vector2 = MyRange(-role_window, None)(all_words2)
        #         lex_cvector1 = MyRange(-role_window, None)(all_cwords1)
        #         lex_cvector2 = MyRange(-role_window, None)(all_cwords2)
        #     elif self.use_trigger:
        #         lex_vector1 = MyRange(-(trigger_window + role_window), None)(all_words1)
        #         lex_vector2 = MyRange(-(trigger_window + role_window), None)(all_words2)
        #         lex_cvector1 = MyRange(-(trigger_window + role_window), None)(all_cwords1)
        #         lex_cvector2 = MyRange(-(trigger_window + role_window), None)(all_cwords2)
        #         if role_window > 0:
        #             lex_vector1 = MyRange(0, -role_window)(lex_vector1)
        #             lex_vector2 = MyRange(0, -role_window)(lex_vector2)
        #             lex_cvector1 = MyRange(0, -role_window)(lex_cvector1)
        #             lex_cvector2 = MyRange(0, -role_window)(lex_cvector2)

        lex_flattened1 = Flatten()(lex_vector1)
        lex_flattened2 = Flatten()(lex_vector2)

        # lex_cflattened1 = Flatten()(lex_cvector1)
        # lex_cflattened2 = Flatten()(lex_cvector2)
        #
        # dep_flattened1 = Flatten()(all_dep1)
        # dep_flattened2 = Flatten()(all_dep2)

#        merged_all1 = merge([maxpool1, lex_flattened1], mode=u'concat')
#        merged_all2 = merge([maxpool2, lex_flattened2], mode=u'concat')

#        lex_layer = keras.layers.multiply([lex_flattened1, lex_flattened2])
#        lex_out = Dense(100, activation=u'sigmoid')(lex_layer)
#        clex_layer = keras.layers.multiply([lex_cflattened1, lex_cflattened2])
#        clex_out = Dense(100, activation=u'sigmoid')(clex_layer)
#        out = merge([lex_out, clex_out], mode=u'concat')

        #merged_final = keras.layers.multiply([maxpool1, maxpool2])
        #merged_final = keras.layers.multiply([merged_all1, merged_all2])
        #merged_final = merge([lex_flattened1, lex_flattened2], mode=u'concat')

        #merged_final = keras.layers.multiply([lex_flattened1, lex_flattened2])  # use this

        #premerge_1 = merge([lex_flattened1, lex_cflattened1, dep_flattened1], mode=u'concat')
        #premerge_2 = merge([lex_flattened2, lex_cflattened2, dep_flattened2], mode=u'concat')
        # premerge_1 = merge([lex_flattened1, dep_flattened1], mode=u'concat')
        # premerge_2 = merge([lex_flattened2, dep_flattened2], mode=u'concat')

        # Note: border_mode='same' to keep output the same width as the input
        if self.use_cnn:
            conv = Convolution1D(self.num_feature_maps, self.filter_length, border_mode=u'same')
            conv1 = conv(context_words1)
            conv2 = conv(context_words2)
            maxpool1 = GlobalMaxPooling1D()(conv1)
            maxpool2 = GlobalMaxPooling1D()(conv2)
            merged1 = merge([maxpool1, lex_flattened1], mode=u'concat')
            merged2 = merge([maxpool2, lex_flattened2], mode=u'concat')
            merged_final = keras.layers.multiply([merged1, merged2])
        else:
            merged_final = keras.layers.multiply([lex_flattened1, lex_flattened2])  # use this
        #merged_final = keras.layers.multiply([premerge_1, premerge_2])  # use this
        #merged_final = merge([premerge_1, premerge_2], mode=u'concat')

        # Dense MLP layer with dropout
        #dropout = Dropout(self.dropout)(merged_final)

        if self.use_dropout:
            dcn_input = Dropout(self.dropout)(merged_final)
        else:
            dcn_input = merged_final

        if self.number_of_hidden_nodes == 0:
            out_final = Dense(1, activation=u'sigmoid')(dcn_input)
        else:
            out = Dense(self.number_of_hidden_nodes, activation=u'sigmoid')(dcn_input)
            #dropout2 = Dropout(self.dropout)(out)
            out_final = Dense(1, activation=u'sigmoid')(out)

        #keras_pair_model = Model(input=[context_input1, context_input2, pos_input1, pos_input2, context_cinput1, context_cinput2, dep_input1, dep_input2], output=[out_final])
        keras_pair_model = Model(input=[context_input1, context_input2, pos_input1, pos_input2], output=[out_final])

        keras_pair_model.compile(optimizer='adadelta',
                                loss=u'binary_crossentropy',
                                metrics=['accuracy'])


#!/usr/bin/env python

# Copyright (C) 2017 by Paul Demorest and Maciej Serylak
# Licensed under the Academic Free License version 3.0
# This program comes with ABSOLUTELY NO WARRANTY.
# You are free to modify and redistribute this code as long
# as you do not remove the above attribution and reasonably
# inform recipients that you have modified the original work.

import numpy
import astropy.io.fits as pyfits

class PSRFITS:
    def __init__(self, filename=None):
        self.fits = None
        if filename is not None:
            self.open(filename)

    def open(self, filename):
        """Open the specified PSRFITS file. A pyfits object is
        created and stored as self.fits. For convenience, the
        primary HDU table header is stored as self.hdr, and
        the subint HDU table header as self.subintHdr."""
        self.fits = pyfits.open(filename, memmap=True, mode='readonly', save_backup=False)
        self.hdr = self.fits[0].header
        self.subint_hdr = self.fits['SUBINT'].header

    def get_freqs(self, row=0):
        """Return the frequency array from the specified subint."""
        return self.fits['SUBINT'].data['DAT_FREQ'][row]

    def get_data(self, start_row=0, end_row=None, downsamp=1, freq_downsamp=1, apply_scales=True, get_ft=False, squeeze=False):
        """Read the data from the specified rows and return it as a
        single array with [time, poln, chan] dimensions.
        Options:
            start_row: first subint read (0-based index)
            end_row: final subint to read. None implies endRow = startRow.
                     Negative values imply offset from the end, i.e.
                     getData(0,-1) would read the entire file. (Don't forget
                     that PSRFITS files are often huge so this might be a bad idea).
            down_samp: downsample the data in time as they are being read in.
                       The downsample factor should evenly divide the number
                       of spectra per row. downSamp = 0 means integrate each
                       row completely.
            freq_down_samp: downsample the data in frequency as they are being read in.
                            The downsample factor should evenly divide the number
                            of channels.
            apply_scales: set to False to avoid applying the scale/offset
                          data stored in the file.
            get_ft: if True return time and frequency arrays as well.
            squeeze: if True, "squeeze" the data array (remove len-1
                     dimensions).
         Notes: Only 8, 16, and 32 bit data are currently understood."""

        # Check if the mode of observations is correct.
        if self.hdr['OBS_MODE'].strip() != 'SEARCH':
            raise RuntimeError('get_data() only works on SEARCH-mode PSRFITS')

        # Get header parameters.
        nsamp_blk = self.subint_hdr['NSBLK']
        npol = self.subint_hdr['NPOL']
        nchan = self.subint_hdr['NCHAN']
        nbit = self.subint_hdr['NBITS']
        tbin = self.subint_hdr['TBIN']
        pol_type = self.subint_hdr['POL_TYPE']
        nrows_file = self.subint_hdr['NAXIS2']
        # print('nsamp_blk: {0}'.format(nsamp_blk))
        # print('npol: {0}'.format(npol))
        # print('nchan: {0}'.format(nchan))
        # print('nbit: {0}'.format(nbit))
        # print('tbin: {0}'.format(tbin))
        # print('pol_type: {0}'.format(pol_type))
        # print('nrows_file: {0}'.format(nrows_file))

        # Check if down sampling in time and frequency is enabled.
        if downsamp == 0:
            downsamp = 1
        if downsamp > nsamp_blk:
            downsamp = nsamp_blk
        if freq_downsamp == 0:
            freq_downsamp = 1
        if freq_downsamp > nchan:
            freq_downsamp = nchan
        if end_row is None:
            end_row = start_row
        if end_row < 0:
            end_row = nrows_file + end_row
        if nsamp_blk % downsamp > 0:
            raise RuntimeError('downsamp does not evenly divide NSBLK ({0}).'.format(nsamp_blk))
        if nchan % freq_downsamp > 0:
            raise RuntimeError('freq_downsamp does not evenly divide NCHAN ({0}).'.format(nchan))

        nrows_total = end_row - start_row + 1
        down_nsamp_blk = nsamp_blk / downsamp
        down_nchan = nchan / freq_downsamp
        down_tbin = tbin * downsamp
        # print('nrows_total: {0}'.format(nrows_total))
        # print('down_nsamp_blk: {0}'.format(down_nsamp_blk))
        # print('down_nchan: {0}'.format(down_nchan))
        # print('down_tbin: {0}'.format(down_tbin))

        # Data types of the signed and unsigned.
        if nbit == 8:
            signed_type = numpy.int8
            unsigned_type = numpy.uint8
        elif nbit == 16:
            signed_type = numpy.int16
            unsigned_type = numpy.uint16
        elif nbit == 32:
            signed_type = numpy.float32
            unsigned_type = numpy.float32
        else:
            raise RuntimeError('Not handled number of NBITS ({0}).'.format(nbit))

        # Allocate arrays.
        row_spectrum_result = numpy.zeros(nchan, dtype=numpy.float32)
        final_spectrum = numpy.zeros((nrows_total * down_nsamp_blk, npol, down_nchan), dtype=numpy.float32)
        if get_ft:
            final_freqs = numpy.zeros(down_nchan)
            final_timing = numpy.zeros(nrows_total * down_nsamp_blk)
        pol_sign = 1
        if 'AABB' in pol_type:
            pol_sign = 2

        # Iterate over rows (blocks).
        for irow in range(nrows_total):
            if apply_scales:
                offsets = self.fits['SUBINT'].data['DAT_OFFS'][irow + start_row]
                scales = self.fits['SUBINT'].data['DAT_SCL'][irow + start_row]
                weights = self.fits['SUBINT'].data['DAT_WTS'][irow + start_row]
                offsets = offsets.reshape((npol, nchan))
                scales = scales.reshape((npol, nchan))
                # weights = numpy.concatenate((weights, weights, weights, weights))
                # weights = weights.reshape((npol, nchan))
            if get_ft:
                row_timing = self.fits['SUBINT'].data['OFFS_SUB'][irow + start_row] - (self.fits['SUBINT'].data['TSUBINT'][irow + start_row] / 2.0)
                freqs_row = self.fits['SUBINT'].data['DAT_FREQ'][irow + start_row]
            row_data = self.fits['SUBINT'].data['DATA'][irow + start_row]
            # print('row_data: {0}'.format(row_data.shape))
            # print('offets: {0}'.format(offsets.shape))
            # print('scales: {0}'.format(scales.shape))
            # print('weights: {0}'.format(weights.shape))

        # Decode row_data for 16 bit data type file.
        if (nbit == 16):
            row_data = numpy.fromstring(row_data.tostring(), dtype=numpy.int16)
            row_data = row_data.reshape((nsamp_blk, npol, nchan))

        # Iterate over samples and polarisations.
        for isamp in range(down_nsamp_blk):
            for ipol in range(npol):
                if ipol < pol_sign:
                    data_type = unsigned_type
                else:
                    data_type = signed_type
                row_spectrum_result = row_data[isamp * downsamp:(isamp + 1) * downsamp, ipol, :].astype(data_type).mean(0)
                if get_ft:
                    final_timing[irow * down_nsamp_blk + isamp] = row_timing + (isamp + 0.5) * down_tbin
                if apply_scales:
                    row_spectrum_result *= scales[ipol, :]
                    row_spectrum_result += offsets[ipol, :]
                    row_spectrum_result *= weights[ipol, :]
                if freq_downsamp == 1:
                    final_spectrum[irow * down_nsamp_blk + isamp, ipol, :] = row_spectrum_result
                    if get_ft:
                        final_freqs[:] = freqs_row[:]
                else:
                    final_spectrum[irow * down_nsamp_blk + isamp, ipol, :] = row_spectrum_result.reshape((-1, freq_downsamp)).mean(1)
                    if get_ft:
                        final_freqs[:] = freqs_row.reshape((-1, freq_downsamp)).mean(1)

        if squeeze:
            final_spectrum = final_spectrum.squeeze()
        if get_ft:
            return (final_spectrum, final_timing, final_freqs)
        else:
            return final_spectrum


# -*- coding:utf-8 -*-

""" Tensorflow Implementation of Variational Autoencoder for sentences.
	Paper: Generating Sentences from a Continuous Space
	(https://arxiv.org/abs/1511.06349)
"""

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

import numpy as np
import tensorflow as tf

class VAE(object):
	def __init__(self, hparam, word_to_index, index_to_word):
		self.hparam = hparam
		self.word_to_index = word_to_index
		self.index_to_word = index_to_word

		self.seq_len = hparam.seq_len
		self.vocab_size = hparam.vocab_size
		self.embedding_dim = hparam.embedding_dim
		self.hidden_dim = hparam.hidden_dim
		self.latent_dim = hparam.latent_dim
		self.clip_norm = hparam.clip_norm
		self.num_layers = hparam.num_layers
		self.beam_width = hparam.beam_width
		self.rnn_type = hparam.rnn_type
		self.anneal_max = hparam.anneal_max
		self.anneal_bias = hparam.anneal_bias
		self.anneal_slope = hparam.anneal_slope

		self.SOS_ID = word_to_index["<sos>"]
		self.EOS_ID = word_to_index["<eos>"]
		self.UNK_ID = word_to_index["<unk>"]
		self.PAD_ID = word_to_index["<pad>"]

		self._build_global_helper()
		self._build_graph()
		self._build_summary()

		self.merged_summary = tf.summary.merge_all()
		self.init = tf.global_variables_initializer()

	def _build_global_helper(self):
		self.inputs = tf.placeholder(tf.int32, [None, self.seq_len])
		self.targets = tf.placeholder(tf.int32, [None, self.seq_len])

		self.enc_inputs = self.inputs
		self.dec_inputs = self.inputs
		self.dec_targets = self.targets
		self.enc_lengths = tf.count_nonzero(self.enc_inputs, 1, dtype=tf.int32)
		self.dec_lengths = self.enc_lengths
		self.batch_size = tf.shape(self.inputs)[0]

		self.lr = tf.Variable(0.0, trainable=False)
		self.new_lr = tf.placeholder(tf.float32, [])
		self.lr_update = tf.assign(self.lr, self.new_lr)

		self.global_step = tf.Variable(0, trainable=False)
		self.saver = tf.train.Saver()

	def _build_graph(self):
		self._build_forward_graph()
		self._build_backward_graph()

	def _build_forward_graph(self):
		with tf.variable_scope("encoder", reuse=None):
			embedding = tf.get_variable("embedding", [self.vocab_size, self.embedding_dim], dtype=tf.float32)
			self.embedded_enc_inputs = tf.nn.embedding_lookup(embedding, self.enc_inputs)
			self.embedded_dec_inputs = tf.nn.embedding_lookup(embedding, self.dec_inputs)

			enc_cell = self.create_rnn_cell(reuse=None)
			enc_init = enc_cell.zero_state(self.batch_size, tf.float32)
			enc_outputs, enc_states = tf.nn.dynamic_rnn(cell = enc_cell,
														initial_state = enc_init,
														inputs = self.embedded_enc_inputs,
														sequence_length = self.enc_lengths,
														dtype = tf.float32)
			if self.rnn_type == "gru":
				enc_vector = enc_states[-1]
			elif self.rnn_type == "block" or self.rnn_type == "basic":
				enc_vector = tf.concat([enc_states[-1].c, enc_states[-1].h], 1)

			self.enc_mean = tf.layers.dense(enc_vector, self.latent_dim, name="enc_mean")
			self.enc_logvar = tf.layers.dense(enc_vector, self.latent_dim, name="enc_logvar")
			epsilon = tf.truncated_normal(tf.shape(self.enc_mean))
			self.latent_sample = self.enc_mean + tf.exp(0.5 * self.enc_logvar) * epsilon

		self.training_rnn_outputs, self.training_logits = self.decoder_training(self.latent_sample)
		self.reconstructions = tf.argmax(self.training_logits, 2)
		self.predicted_ids = self.decoder_inference(self.latent_sample)

	def _build_backward_graph(self):
		self.nll_loss = self.nll_loss_fn()
		self.kl_w = self.kl_w_fn()
		self.kl_loss = self.kl_loss_fn()
		self.total_loss = self.nll_loss + self.kl_loss * self.kl_w

		params = tf.trainable_variables()
		gradients = tf.gradients(self.total_loss, params)
		clipped_gradients, _ = tf.clip_by_global_norm(gradients, self.clip_norm)
		self.train_op = tf.train.AdamOptimizer(self.lr).apply_gradients(
			zip(clipped_gradients, params), global_step=self.global_step)

	def create_rnn_cell(self, reuse):
		def create_single_cell(reuse):
			if self.rnn_type == "gru":
				cell = tf.contrib.rnn.GRUCell(self.hidden_dim,
										  	  kernel_initializer=tf.orthogonal_initializer(),
										  	  reuse=reuse)
			elif self.rnn_type == "block":
				cell = tf.contrib.rnn.LSTMBlockCell(self.hidden_dim, forget_bias=0.0, reuse=reuse)
			elif self.rnn_type == "basic":
				cell = tf.contrib.rnn.BasicLSTMCell(self.hidden_dim, forget_bias=0.0, state_is_tuple=True, reuse=reuse)
			else:
				ValueError("Unsupported rnn_type.")
			return cell
		cell = tf.contrib.rnn.MultiRNNCell([create_single_cell(reuse) for _ in range(self.num_layers)], state_is_tuple=True)
		return cell

	def _build_summary(self):
		tf.summary.scalar("total_loss", self.total_loss)
		tf.summary.scalar("nll_loss", self.nll_loss)
		tf.summary.scalar("kl_loss", self.kl_loss)
		tf.summary.scalar("kl_w", self.kl_w)
		# tf.summary.scalar("ratio", self.nll_loss / self.kl_loss)

	def decoder_training(self, latent_sample):
		with tf.variable_scope("encoder", reuse=True):
			embedding = tf.get_variable("embedding", [self.vocab_size, self.embedding_dim], dtype=tf.float32)
		with tf.variable_scope("decoder", reuse=None):
			if self.rnn_type == "gru":
				init_state = tuple([tf.layers.dense(latent_sample, self.hidden_dim, tf.nn.elu, name="trans0")] * self.num_layers)
			else:
				c = tf.layers.dense(latent_sample, self.hidden_dim, tf.nn.elu, name="trans1")
				h = tf.layers.dense(latent_sample, self.hidden_dim, tf.nn.elu, name="trans2")
				init_state = tuple([tf.contrib.rnn.LSTMStateTuple(c=c, h=h)] * self.num_layers)

			lin_proj = tf.layers.Dense(self.vocab_size, _scope='decoder/dense', _reuse=None)
			dec_cell = self.create_rnn_cell(reuse=None)
			helper = tf.contrib.seq2seq.TrainingHelper(inputs = self.embedded_dec_inputs,
													   sequence_length = self.dec_lengths)
			decoder = tf.contrib.seq2seq.BasicDecoder(cell = dec_cell,
													  helper = helper,
													  initial_state = init_state)
			decoder_output, _, _ = tf.contrib.seq2seq.dynamic_decode(decoder = decoder)
		return decoder_output.rnn_output, lin_proj.apply(decoder_output.rnn_output)

	def decoder_inference(self, latent_sample):
		with tf.variable_scope("encoder", reuse=True):
			embedding = tf.get_variable("embedding", [self.vocab_size, self.embedding_dim], dtype=tf.float32)
		with tf.variable_scope("decoder", reuse=True):
			if self.rnn_type == "gru":
				init_state = tuple([tf.layers.dense(latent_sample, self.hidden_dim, tf.nn.elu, name="trans0")] * self.num_layers)
			else:
				c = tf.layers.dense(latent_sample, self.hidden_dim, tf.nn.elu, name="trans1")
				h = tf.layers.dense(latent_sample, self.hidden_dim, tf.nn.elu, name="trans2")
				init_state = tuple([tf.contrib.rnn.LSTMStateTuple(c=c, h=h)] * self.num_layers)

			lin_proj = tf.layers.Dense(self.vocab_size, _scope='decoder/dense', _reuse=True)
			dec_cell = self.create_rnn_cell(reuse=True)
			decoder = tf.contrib.seq2seq.BeamSearchDecoder(cell=dec_cell,
													   embedding=embedding,
													   start_tokens=tf.tile(tf.constant([self.SOS_ID], dtype=tf.int32), [self.batch_size]),
													   end_token = tf.constant(self.EOS_ID, dtype=tf.int32),
													   initial_state = tf.contrib.seq2seq.tile_batch(init_state, self.beam_width),
													   beam_width = self.beam_width,
													   output_layer = lin_proj)
			decoder_output, _, _ = tf.contrib.seq2seq.dynamic_decode(decoder = decoder,
											maximum_iterations = 2 * tf.reduce_max(self.enc_lengths))
		return decoder_output.predicted_ids[:, :, 0]

	def nll_loss_fn(self):
		mask = tf.sequence_mask(self.dec_lengths, tf.reduce_max(self.dec_lengths), dtype=tf.float32)
		return tf.reduce_sum(tf.contrib.seq2seq.sequence_loss(
					logits = self.training_logits,
					targets = self.dec_targets,
					weights = mask,
					average_across_timesteps = False,
					average_across_batch = True))

	def kl_w_fn(self):
		return self.anneal_max * tf.sigmoid((self.anneal_slope) * (tf.cast(self.global_step, tf.float32) - tf.constant(self.anneal_bias / 2)))

	def kl_loss_fn(self):
		return 0.5 * tf.reduce_sum(tf.exp(self.enc_logvar) + tf.square(self.enc_mean) - 1 - self.enc_logvar) / tf.to_float(self.batch_size)

	def train_session(self, sess, inputs, targets):
		feed_dict = {self.inputs: inputs,
					 self.targets: targets}
		fetches = [self.train_op,
				   self.nll_loss,
				   self.kl_loss,
				   self.kl_w,
				   self.total_loss,
				   self.global_step,
				   self.merged_summary]
		_, nll_loss, kl_loss, kl_w, total_loss, step, summary = sess.run(fetches, feed_dict)
		return {"nll_loss": nll_loss,
				"kl_loss": kl_loss,
				"kl_w": kl_w,
				"total_loss": total_loss,
				"step": step,
				"summary": summary}

	def test_session(self, sess, inputs, targets):
		feed_dict = {self.inputs: inputs,
					 self.targets: targets}
		nll_loss, total_loss, kl_loss = sess.run([self.nll_loss, self.total_loss, self.kl_loss], feed_dict)
		return {"nll_loss": nll_loss,
				"total_loss": total_loss,
				"kl_loss": kl_loss}

	def random_generate_session(self, sess):
		feed_dict = {self.latent_sample: np.random.randn(1, self.latent_dim),
					 self.enc_lengths: [self.seq_len],
					 self.batch_size: 1}
		predicted_ids = sess.run(self.predicted_ids, feed_dict)[0]
		ret = 'Random Generation: %s' % ' '.join([self.index_to_word[index] for index in predicted_ids])
		return ret

	def reconstruct_session(self, sess, sentence):
		idx_list = [self.get_word_index(w) for w in sentence.lower().split()][:self.seq_len]
		idx_list = idx_list + [self.PAD_ID] * (self.seq_len - len(idx_list))
		ret1 = 'Original: %s' % ' '.join([self.index_to_word[index] for index in idx_list])

		predicted_ids = sess.run(self.reconstructions, {self.inputs: np.atleast_2d(idx_list)})[0]
		ret2 = 'Reconstr: %s' % ' '.join([self.index_to_word[index] for index in predicted_ids])
		return ret1, ret2

	def get_word_index(self, word):
		index = self.word_to_index[word]
		return index if index < self.vocab_size else self.UNK_ID

	def get_parameter_size(self):
		all_vars = tf.global_variables()
		total_count = 0
		for item in all_vars:
			if "Adam" in item.name:
				continue
			shape = item.get_shape().as_list()
			if len(shape) == 0:
				total_count += 1
			else:
				size =  1
				for val in shape:
					size *= val
				total_count += size
		return total_count


"""
Module containing functions for mutating ProteinNetRecords and feeding that data into further computations (e.g. Tensorflow).
These functions are fairly specific so may often be better used as inspiration to build users own solutions.
"""
import random
import logging

import numpy as np

from .data import LabeledFunction
from .maths import softmin_sample, softmax_sample

__all__ = ["ProteinNetMutator", "sequence_mutator", "per_position_mutator",
           "sample_deleterious", "sample_neutral"]

class ProteinNetMutator(LabeledFunction):
    """
    Map function generating mutated records.

    Apply a mutator function to a ProteinNet record and return the mutated sequence. This is a LabeledFunction that can be used to generate a TensorFlow Dataset. This setup is fairly specific to your downstream model design, so it will often be more useful to use it as a base to create an alternate implementation.

    Returns are in the form:
        ([wt_seq], mut_seq, [phi, psi, chi1]), label, [weights]

    Attributes
    ----------
    wildtype      : bool
        Outputs wildtype as well as mutant sequence.
    phi           : bool
        Outputs Phi backbone angles.
    psi           : bool
        Outputs Psi backbone angles.
    chi           : bool
        Outputs rotamer angles.
    mutator       : function
        Mutator function taking a ProteinNetRecord and returning the sampled variants and their deleteriousness. The return format depends on `per_position`. If per_position=False must return a tuple with the mutated sequence index array and whether it is deleterious (1/0). If per_position=True must return a tuple with mutant_seq, deleterious_inds, neutral_inds arrays.
    kwargs        : dict
        Keyword arguments passed to the mutator function.
    encoding      : dict
        Encoding mapping alphabetically encoded integer indeces to a new scheme.
    weights       : list
        List of float weights for WT, Deleterious and Neutral variants when mutating per position.
    func          : function
        Function applied when the class is called. This is a mutator applied to the whole sequence or per position derived from the initialisation parameters.
    output_shapes, output_types : tuple
        Tuple of output shapes and types (see data.LabeledFunction for details)
    """
    def __init__(self, mutator, per_position=False, include=('wt',),
                 weights=(0, 1, 1), encoding=None, **kwargs):
        """
        Initialise the mutator.

        Parameters
        ----------
        mutator       : function
            Mutator function taking a ProteinNetRecord and returning the sampled variants and their deleteriousness. The return format depends on `per_position`. If per_position=False must return a tuple with the mutated sequence index array and whether it is deleterious (1/0). If per_position=True must return a tuple with mutant_seq, deleterious_inds, neutral_inds arrays.
        per_position  : bool
            Return deleteriousness for each position rather than the entire sequence.
        include:      : sequence_like including some of {"wt", "phi", "psi", "chi1"}
            Variables to return for computation, alongside mutant sequence.
        weights       : list
            weightings for [wt, deleterious, neutral] variants when processing per residue variants
        encoding      : dict
            Optional dictionary mapping alphabetically encoded AA integer indeces to a new scheme (e.g. that used in UniRep).
        **kwargs      : dict
            arguments passed on to mutator
        """
        self.wildtype = 'wt' in include
        self.phi = 'phi' in include
        self.psi = 'psi' in include
        self.chi = 'chi1' in include

        self.mutator = mutator
        self.kwargs = kwargs

        self.encoding = encoding
        self.weights = weights

        if per_position:
            self.func = self._per_position_func
        else:
            self.func = self._whole_seq_func

        # Calculate output shape/type (Mutant sequence is always output)
        self.output_shapes = [[None]]
        self.output_types = ['int32']

        if self.wildtype:
            self.output_shapes.insert(0, [None])
            self.output_types.insert(0, 'int32')

        if self.phi:
            self.output_shapes.append([None])
            self.output_types.append('float32')

        if self.psi:
            self.output_shapes.append([None])
            self.output_types.append('float32')

        if self.chi:
            self.output_shapes.append([None])
            self.output_types.append('float32')

        if per_position:
            self.output_shapes = (tuple(self.output_shapes), [None, 1], [None])
            self.output_types = (tuple(self.output_types), 'int32', 'int32')
        else:
            self.output_shapes = (tuple(self.output_shapes), [])
            self.output_types = (tuple(self.output_types), 'int32')

    def _per_position_func(self, record):
        """
        Function producing per position mutants and labels from a ProteinNetRecord

        Parameters
        ----------
        record : ProteinNetRecord
            Record to mutate.
        """
        mut_seq, deleterious, neutral = self.mutator(record, **self.kwargs)

        mut_weights = np.full(deleterious.shape, self.weights[0])
        mut_weights[deleterious] = self.weights[1]
        mut_weights[neutral] = self.weights[2]

        if self.encoding is not None:
            mut_seq = np.array([self.encoding[x] for x in mut_seq])

        if self.wildtype:
            if self.encoding is not None:
                wt_seq = np.array([self.encoding[x] for x in record.primary_ind])
            else:
                wt_seq = record.primary_ind

            output = [wt_seq + 1, mut_seq + 1] # +1 to leave 0 free for padding
        else:
            output = [mut_seq + 1]

        if self.phi:
            output.append(record.rama[1])

        if self.psi:
            output.append(record.rama[2])

        if self.chi:
            output.append(record.chi)

        return tuple(output), deleterious.reshape((len(deleterious), 1)), mut_weights

    def _whole_seq_func(self, record):
        """
        Function applying mutator to a record and producing mutant data and a single label for the sequence from a ProteinNetRecord.

        Parameters
        ----------
        record : ProteinNetRecord
            Record to mutate.
        """
        mut_seq, label = self.mutator(record, **self.kwargs)

        if self.encoding is not None:
            mut_seq = np.array([self.encoding[x] for x in mut_seq])

        if self.wildtype:
            if self.encoding is not None:
                wt_seq = np.array([self.encoding[x] for x in record.primary_ind])
            else:
                wt_seq = record.primary_ind

            output = [wt_seq + 1, mut_seq + 1] # +1 to leave 0 free for padding
        else:
            output = [mut_seq]

        if self.phi:
            output.append(record.rama[1])

        if self.psi:
            output.append(record.rama[2])

        if self.chi:
            output.append(record.chi)

        return tuple(output), label

def sequence_mutator(record, p_deleterious=0.5, max_mutations=3,
                     max_deleterious=0.01, min_neutral=0.1):
    """
    Generate mutated sequences from a ProteinNetRecord with a few deleterious or neutral variants.

    Generate mutated sequences from a ProteinNetRecord with a few deleterious and/or neutral variants. First randomly choose to generate a deleterious or neutral sequence then sample some of the corresponding variant types based on the records MSA frequencies.

    Parameters
    ----------
    record          : ProteinNetRecord
        Record to mutate.
    p_deleterious   : float
        Probability of returning a deleterious set of variants.
    max_mutations   : int
        Maximum number of mutations to make.
    max_deleterious : float
        Maximum MSA frequency for a variant to be considered deleterious.
    min_neutral     : float
        Minimum MSA frequency for a variant to be considered neutral.

    Returns
    -------
    tuple
        Tuple of the format (seq, deleterious). The first entry is the mutated amino acid sequence, encoded with integer indeces and the second is 1 if the sequence is deleterious and 0 if neutral.
    """
    deleterious = int(random.random() < p_deleterious)
    seq = record.primary_ind.copy()

    num_muts = random.randint(1, max_mutations)

    if deleterious:
        positions, subs = sample_deleterious(num=num_muts, wt_seq=seq,
                                             pssm=record.evolutionary,
                                             max_freq=max_deleterious)
    else:
        positions, subs = sample_neutral(num=num_muts, wt_seq=seq,
                                         pssm=record.evolutionary,
                                         min_freq=min_neutral)

    seq[positions] = subs

    return seq, deleterious

def per_position_mutator(record, max_deleterious=2, max_neutral=4,
                         max_deleterious_freq=0.01, min_neutral_freq=0.1):
    """
    Generate mutated sequences from ProteinNetRecords with labels identifying deleterious and neutral mutations.

    Generate mutated sequences from ProteinNetRecords with labels identifying where deleterious and neutral mutations have been made.
    Will always generate at least one variant.

    Parameters
    ----------
    record               : ProteinNetRecord
        Record to mutate.
    max_deleterious      : int
        Maximum number of deleterious variants to make.
    max_neutral          : int
        Maximum number of neutral variants to make.
    max_deleterious_freq : float
        Maximum MSA frequency for a variant to be considered deleterious.
    min_neutral_freq     : float
        Minimum MSA frequency for a variant to be considered neutral.

    Returns
    -------
    tuple
        Tuple of the format seq, deleterious, neutral. The first entry is the mutated sequence, the second a list of positions with deleterious variants and the third a list of positions with neutral variants.
    """
    seq = record.primary_ind.copy()

    # Masks identifying where deletious and neutral mutations are made
    deleterious = np.zeros(seq.shape, dtype=int)
    neutral = np.zeros(seq.shape, dtype=int)

    # Sample number of substitutions [deleterious, neutral]
    num_deleterious = random.randint(0, max_deleterious)
    num_neutral = random.randint(0 if num_deleterious else 1, max_neutral)

    pos = None # Need a mask for sample_neutral if not making any del subs
    del_subs = 0 # Number of subs actually made
    neut_subs = 0
    if num_deleterious:
        try:
            pos, subs = sample_deleterious(num=num_deleterious,
                                           wt_seq=record.primary_ind,
                                           pssm=record.evolutionary,
                                           max_freq=max_deleterious_freq)

            seq[pos] = subs
            deleterious[pos] = 1
            del_subs = len(pos)

        except ValueError:
            pass

    if num_neutral:
        try:
            pos, subs = sample_neutral(num=num_neutral,
                                       wt_seq=record.primary_ind,
                                       pssm=record.evolutionary,
                                       min_freq=min_neutral_freq,
                                       mask=pos)

            seq[pos] = subs
            neutral[pos] = 1
            neut_subs = len(subs)

        except ValueError:
            pass

    if not del_subs and not neut_subs:
        raise ValueError('No valid mutations could be made')

    return seq, deleterious, neutral

def sample_deleterious(num, pssm, wt_seq, max_freq=0.025, mask=None):
    """
    Sample deleterious mutations from a MSA frequency matrix.

    Randomly choose a selection of deleterious variants from a MSA frequency matrix.

    Parameters
    ----------
    num      : int
        Number of mutations to make.
    pssm     : float ndarray (20, N)
        MSA frequency matrix to determine neutral and deleterious variants.
    wt_seq   : int ndarray (N,)
        WT sequence of the protein (as int indeces corresponding to the MSA matrix rows).
    max_freq : float
        Maximum frequency considered deleterious.
    mask     : int array_like
        Array of positions not to mutate.

    Returns
    -------
    tuple
        Numpy array of position indeces chosen and an array of the alternate amino acid in each position (as MSA row indeces).
    """
    if num == 0:
        raise ValueError('num must be > 0')

    # Select valid positions
    pssm[wt_seq, np.arange(pssm.shape[1])] = 0
    positions = np.nonzero(np.sum(pssm < max_freq, 0))[0]
    if not mask is None:
        positions = positions[~np.isin(positions, mask)]

    if positions.size == 0:
        raise ValueError('No valid neutral substitutions')

    if num > positions.size:
        logging.warning(('Too few positions to sample %s deleterious mutations, '
                         'sampling as many as possible (%s)'), num, positions.size)
        num = positions.size

    positions = np.random.choice(positions, replace=False, size=num)

    # Sample and make substitutions
    pssm = pssm[:, positions]
    substitutions = np.apply_along_axis(lambda x: softmin_sample(x, max_value=max_freq), 0, pssm)
    return positions, substitutions

def sample_neutral(num, pssm, wt_seq, min_freq=0.025, mask=None):
    """
    Sample deleterious mutations froma pssm

    Parameters
    ----------
    num      : int
        Number of mutations to make.
    pssm     : float ndarray (20, N)
        MSA frequency matrix to determine neutral and deleterious variants.
    wt_seq   : int ndarray (N,)
        WT sequence of the protein (as int indeces corresponding to the MSA matrix rows).
    min_freq : float
        Minimum frequency considered neutral.
    mask     : int array_like
        Array of positions not to mutate.

    Returns
    -------
    tuple
        Numpy array of position indeces chosen and an array of the alternate amino acid in each position (as MSA row indeces).
    """
    if num == 0:
        return None, None

    # Select valid positions
    pssm[wt_seq, np.arange(pssm.shape[1])] = 0
    positions = np.nonzero(np.sum(pssm > min_freq, 0))[0]
    if not mask is None:
        positions = positions[~np.isin(positions, mask)]

    if positions.size == 0:
        raise ValueError('No valid neutral substitutions')

    if num > positions.size:
        logging.warning(('Too few positions to sample %s deleterious mutations, '
                         'sampling as many as possible (%s)'), num, positions.size)
        num = positions.size

    positions = np.random.choice(positions, replace=False, size=num)

    # Sample and make substitutions
    pssm = pssm[:, positions]
    substitutions = np.apply_along_axis(lambda x: softmax_sample(x, min_value=min_freq), 0, pssm)
    return positions, substitutions


from __future__ import absolute_import

from builtins import object

import numpy as np
import logging

from relaax.common import profiling
from relaax.server.common import session
from relaax.common.algorithms.lib import utils
from relaax.common.algorithms.lib import observation

from .. import dqn_config
from .. import dqn_model
from . import dqn_utils

logger = logging.getLogger(__name__)
profiler = profiling.get_profiler(__name__)

class Trainer(object):
    def __init__(self, parameter_server, metrics, exploit):
        self.ps = parameter_server
        self.metrics = metrics
        self._exploit = exploit

        self.session = session.Session(dqn_model.AgentModel())
        self.session.op_initialize()

        self.replay_buffer = dqn_utils.ReplayBuffer(dqn_config.config.replay_buffer_size,
                                                    dqn_config.config.alpha)
        self.observation = observation.Observation(dqn_config.config.input.history)

        self.last_action = None
        self.local_step = 0
        self.last_target_weights_update = 0
        self.agent_weights_id = 0

    @profiler.wrap
    def begin(self):
        self.get_action()

    @profiler.wrap
    def step(self, reward, state, terminal):
        self.local_step += 1

        if self.local_step % dqn_config.config.update_target_weights_interval == 0:
            self.session.op_update_target_weights()

        self.receive_experience()

        if self.local_step > dqn_config.config.start_sample_step:
            self.update()
            if reward is not None:
                self.ps.session.op_add_rewards_to_model_score_routine(reward_sum=reward, reward_weight=1)

        # metrics
        if state is not None:
            self.metrics.histogram('state', state)

        if reward is None:
            self.observation.add_state(state)
        else:
            self.push_experience(reward, state, terminal)

        if terminal:
            self.observation.add_state(None)

        assert self.last_action is None
        self.get_action()

    @profiler.wrap
    def update(self):
        experience = self.replay_buffer.sample(dqn_config.config.batch_size)
        self.send_experience(experience)

    @profiler.wrap
    def send_experience(self, experience):
        batch = dict(zip(experience[0], zip(*[d.values() for d in experience])))
        q_next_target = self.session.op_get_q_target_value(next_state=batch["next_state"])
        q_next = self.session.op_get_q_value(state=batch["next_state"])

        feeds = dict(state=batch["state"],
                     reward=batch["reward"],
                     action=batch["action"],
                     terminal=batch["terminal"],
                     q_next_target=q_next_target,
                     q_next=q_next)

        gradients = self.session.op_compute_gradients(**feeds)

        for i, g in enumerate(utils.Utils.flatten(gradients)):
            self.metrics.histogram('gradients_%d' % i, g)

        self.ps.session.op_submit_gradients(gradients=gradients, step_inc=1, agent_step=self.agent_weights_id)

    @profiler.wrap
    def receive_experience(self):
        weights, self.agent_weights_id = self.ps.session.op_get_weights_signed()
        self.session.op_assign_weights(weights=weights)

    def push_experience(self, reward, state, terminal):
        assert not self.observation.is_none()
        assert self.last_action is not None

        old_state = self.observation.get_state()
        if state is not None:
            self.observation.add_state(state)

        if dqn_config.config.output.q_values:
            action = np.squeeze(np.argmax(self.last_action)).astype(np.int32)
        else:
            action = self.last_action

        self.replay_buffer.append(dict(state=old_state,
                                       action=action,
                                       reward=reward,
                                       terminal=terminal,
                                       next_state=self.observation.get_state()))

        self.last_action = None

    def get_action(self):
        if self.observation.is_none():
            self.last_action = None
        else:
            q_value = self.session.op_get_q_value(state=[self.observation.get_state()])
            self.last_action = self.session.op_get_action(local_step=self.local_step, q_value=q_value)

            assert self.last_action is not None

            # metrics
            self.metrics.histogram('action', self.last_action)


import tensorflow as tf                #神经机器翻译（Neural machine translation, NMT）是最近提出的机器翻译方法。
import numpy as np                     #应该是pretrain一系列的东西
import time                            #f和g在这里面定义
import sys
import re
sys.path.append("..")
from utils.data import load_dataset, load_paired_dataset
from utils.vocab import build_vocab_from_file, load_vocab
from model import NMT
from utils import constants
from options import *

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

def create_model(sess, args, src_vocab_size, tgt_vocab_size, src_vocab_rev, tgt_vocab_rev, mode=constants.TRAIN,
                 reuse=None, load_pretrained_model=False, direction="", model_save_dir=None):
    sess.run(tf.tables_initializer())

    with tf.variable_scope(constants.NMT_VAR_SCOPE + direction, reuse=reuse):
        with tf.variable_scope("src"):
            src_emb = tf.get_variable("embedding", shape=[src_vocab_size, args.emb_dim])
        with tf.variable_scope("dst"):
            tgt_emb = tf.get_variable("embedding", shape=[tgt_vocab_size, args.emb_dim])

        model = NMT(mode, args.__dict__, src_vocab_size, tgt_vocab_size, src_emb, tgt_emb,
                    src_vocab_rev, tgt_vocab_rev, direction)
                     #一系列环境词库变量，fg等装入model中

    if load_pretrained_model:
        if model_save_dir is None:
            model_save_dir = args.nmt_model_save_dir
            if direction not in model_save_dir:
                if direction[::-1] in model_save_dir:
                    model_save_dir = re.sub(direction[::-1], direction, model_save_dir)
                else:
                    model_save_dir = os.path.join(model_save_dir, direction)
        print(model_save_dir)
        try:
            print("Loading nmt model from", model_save_dir)
            model.saver.restore(sess, model_save_dir)
        except Exception as e:
            print("Error! Loading nmt model from", model_save_dir)
            print("Again! Loading nmt model from", tf.train.latest_checkpoint(model_save_dir))
            model.saver.restore(sess, tf.train.latest_checkpoint(model_save_dir))
    else:
        if reuse is None:
            print("Creating model with new parameters.")
            sess.run(tf.global_variables_initializer()) #开始训练
        else:
            print("Reuse parameters.")
    return model

def train(model):
    best = {
        "loss": 100.0,
        "step": 0,
    }
    global_step = 0
    for i in range(args.n_epoch):
        print("Epoch:%d" % i)
        sess.run(train_iterator.initializer)
        n_batch = 0
        t0 = time.time()
        while True:
            try:
                src = sess.run(train_next_op)
                feed_dict = {
                    model.input_ids: src["ids"],
                    model.input_length: src["length"],
                    model.target_ids_in: src["trans_ids_in"],
                    model.target_ids_out: src["trans_ids_out"],
                    model.target_length: src["trans_length"]
                }

                assert src["length"].shape == src["trans_length"].shape
                assert src["ids"].shape[0] == src["trans_ids"].shape[0]
                ops = [model.loss,
                       model.logits,
                       model.train_op, ]
                res = sess.run(ops, feed_dict=feed_dict)   #完成前项传播反向传播更新，每一百  sess等于一个函数的返回值

                n_batch += 1
                global_step += 1
                if n_batch % 100 == 0:
                    eval_loss = eval(model)  #评估损失函数
                    if eval_loss < best["loss"]:
                        best["loss"] = eval_loss
                        best["step"] = global_step
                        print("Model save at: %s" % (args.nmt_model_save_dir+str(n_batch)))
                        model.saver.save(sess, args.nmt_model_save_dir, global_step=n_batch)

                    print("Epoch/n_batch:%d/%d\tTrain_loss:%.3f\tEval_loss:%.3f\tMin_eval_loss:%.3f\tBest_step:%d"
                          "\tTime:%d" % (i, n_batch, res[0], eval_loss, best["loss"], best["step"], time.time() - t0))

            except tf.errors.OutOfRangeError:  # next epoch
                print("Train---Total N batch:{}".format(n_batch))
                break

def eval(model):
    sess.run(dev_iterator.initializer)
    n_batch = 0
    t0 = time.time()
    loss = []
    while True:
        try:
            src = sess.run(dev_next_op)
            feed_dict = {
                model.input_ids: src["ids"],
                model.input_length: src["length"],
                model.target_ids_in: src["trans_ids_in"],
                model.target_ids_out: src["trans_ids_out"],
                model.target_length: src["trans_length"]
            }
            assert src["length"].shape == src["trans_length"].shape
            assert src["ids"].shape[0] == src["trans_ids"].shape[0]

            ops = [model.loss]
            res = sess.run(ops, feed_dict=feed_dict)
            loss.append(res[0])
            n_batch += 1
        except tf.errors.OutOfRangeError:  # next epoch
            print("Eval---Total N batch:{}\tLoss:{}\tCost time:{}".format(n_batch, np.mean(loss), time.time() - t0))
            break
    return np.mean(loss)

def inference(model, A, B, sess, args, src_test_iterator, src_test_next, src_vocab_rev, result_dir=None, step=None): #要改
    sess.run([src_test_iterator.initializer]) #需要改动

    n_batch = 0
    t0 = time.time()
    probs = []
    if result_dir is not None:
        if not os.path.exists(result_dir):
            os.makedirs(result_dir)
        src_test_file = args.test_data[A].split("/")[-1]
        step = str(step) + "_" if step is not None else ''
        result_save_path1 = result_dir + "/" + step + src_test_file + '-' + str(B) + '.tsf'
        result_save_path2 = result_dir + "/" + step + src_test_file + '.tsf'
        print("Result save path:" + result_save_path1 + ", " + result_save_path2)
        dst_f1 = open(result_save_path1, "w")
        dst_f2 = open(result_save_path2, "w")

    while True:#重要
        try:
            n_batch += 1
            src = sess.run(src_test_next)
            feed_dict = {
                model.input_ids: src["ids"],
                model.input_length: src["length"],
            }
            predictions = sess.run(model.predictions, feed_dict=feed_dict)
            probs.extend(predictions["log_probs"].reshape(-1))

            if result_dir is not None:
                def save_prediction():
                    batch_size = predictions["tokens"].shape[0]
                    decode_width = predictions["tokens"].shape[1]
                    src_tokens = src_vocab_rev.lookup(tf.cast(src["ids"], tf.int64))
                    src_tokens = sess.run(src_tokens)
                    for i in range(batch_size):
                        src_tokens_ = src_tokens[i][:src["length"][i]]
                        src_sent = " ".join(src_tokens_)
                        for j in range(decode_width):
                            log_probs = predictions["log_probs"][i][j]
                            pred_tokens = predictions["tokens"][i][j][:predictions["length"][i][j] - 1]  # ignore </s>
                            pred_sent = " ".join(pred_tokens)
                            dst_f1.write("%s\t%s\t%s\n" % (src_sent, pred_sent, log_probs))
                            if j == 0:
                                dst_f2.write("%s\n" % pred_sent)

                save_prediction()

        except tf.errors.OutOfRangeError:  # next epoch
            print("INFERENCE---Total N batch:{}\tGenerate probs:{}\tCost time:{}".format(
                n_batch, np.mean(probs), time.time() - t0))
            break
    if result_dir is not None:
        dst_f1.close()
        dst_f2.close()
    return np.mean(probs), result_save_path2

def get_nmt_direction(direction="0-1"): #fg转换
    if direction == "0-1":
        A = 0
        B = 1
    else:
        A = 1
        B = 0
    return A, B
def get_nmt_direction1(direction="0-2"): #新加的
    if direction == "0-2":
        A = 0
        C = 2
    else:
        A = 1
        C = 2
    return A, C
def get_nmt_direction2(direction="1-2"): #新加的
    if direction == "1-2":
        C = 2
        B = 1
    else:
        C = 2
        B = 1
    return C, B
if __name__ == "__main__":
    args = load_nmt_arguments()

    # === Get translation direction
    A, B = get_nmt_direction(args.nmt_direction)
    print("A=%s, B=%s" % (A, B))

    A, C = get_nmt_direction(args.nmt_direction)#新加的
    print("A=%s, C=%s" % (A, C))

    B, C = get_nmt_direction(args.nmt_direction)#新加的
    print("B=%s, C=%s" % (B, C))

    # ===  Build vocab and load data
    if not os.path.isfile(args.global_vocab_file):
        build_vocab_from_file(args.train_data, args.global_vocab_file)
    vocab, vocab_size = load_vocab(args.global_vocab_file)
    src_vocab = tgt_vocab = vocab
    src_vocab_size = tgt_vocab_size = vocab_size
    print('Vocabulary size:src:%s' % vocab_size)

    vocab_rev = tf.contrib.lookup.index_to_string_table_from_file(
        args.global_vocab_file,  # target vocabulary file(each lines has a word)
        vocab_size=vocab_size - constants.NUM_OOV_BUCKETS,
        default_value=constants.UNKNOWN_TOKEN)
    src_vocab_rev = tgt_vocab_rev = vocab_rev

    with tf.device("/cpu:0"):  # Input pipeline should always be placed on the CPU. 在cpu上输入？ 指定Session在第1块CPU上运行
        print("Use x'->y to update model f(x->y)") #需要改
        train_iterator = load_paired_dataset(args.tsf_train_data[B], args.train_data[B],
                                             src_vocab, tgt_vocab, batch_size=args.batch_size)
        dev_iterator = load_paired_dataset(args.tsf_dev_data[B], args.dev_data[B],
                                           src_vocab, tgt_vocab, batch_size=args.batch_size)

        src_test_iterator = load_dataset(args.test_data[A], src_vocab, mode=constants.INFER)

        train_next_op = train_iterator.get_next()
        dev_next_op = dev_iterator.get_next()
        src_test_next_op = src_test_iterator.get_next()

    # === Create session
    tf_config = tf.ConfigProto()
    tf_config.gpu_options.allow_growth = True
    tf_config.gpu_options.per_process_gpu_memory_fraction = 0.4  #在现有的基础上吗？
    sess = tf.Session(config=tf_config)

    # === Train
    if args.mode == "train":#两个对称的model
        # Prepare for model saver
        print("Prepare for model saver") #模型的保存器
        args.nmt_model_save_dir = "%s/%s-%s/" % (args.nmt_model_save_dir, A, B) #要改，应该是实际名称加一个0-1
        print("Model save dir:", args.nmt_model_save_dir)
        if not os.path.exists(args.nmt_model_save_dir):
            os.makedirs(args.nmt_model_save_dir)

        dump_args_to_yaml(args, args.nmt_model_save_dir) #把地址转换成为静态的yaml的中间态

        # Initial and build model 一个训练模型一个推理模型？
        train_model = create_model(sess, args, src_vocab_size, tgt_vocab_size, src_vocab_rev, tgt_vocab_rev,
                                   mode=constants.TRAIN, direction=args.nmt_direction)
        infer_model = create_model(sess, args, src_vocab_size, tgt_vocab_size, src_vocab_rev, tgt_vocab_rev,
                                   mode=constants.INFER, direction=args.nmt_direction, reuse=True)
        train(train_model)

    # === Inference
    elif args.mode == "inference":
        print("Prepare for model saver")
        final_model_save_path = "%s%s-%s/" % (args.nmt_model_save_dir, A, B) #要改
        print("Model save path:", final_model_save_path)
        eval_model = create_model(sess, args, src_vocab_size, tgt_vocab_size, src_vocab_rev, tgt_vocab_rev,
                                  mode=constants.EVAL, direction=args.nmt_direction, load_pretrained_model=True,
                                  model_save_dir=final_model_save_path)
        infer_model = create_model(sess, args, src_vocab_size, tgt_vocab_size, src_vocab_rev, tgt_vocab_rev,
                                   mode=constants.INFER, direction=args.nmt_direction, reuse=True)
        inference(infer_model, A=A, B=B, sess=sess, args=args, src_test_iterator=src_test_iterator,
                  src_test_next=src_test_iterator.get_next(),
                  result_dir=args.tsf_result_dir, src_vocab_rev=src_vocab_rev)

    # === Final inference (after dualRL training)
    elif args.mode == "final_inference":  # todo: check data loader
        print("Prepare for model saver")
        final_model_save_path = "%s%s-%s/" % (args.final_model_save_dir, A, B) #要改

        args.decode_type = constants.RANDOM

        print("Model save path:", final_model_save_path)
        eval_model = create_model(sess, args, src_vocab_size, tgt_vocab_size, src_vocab_rev, tgt_vocab_rev,
                                  mode=constants.EVAL, direction=args.nmt_direction, load_pretrained_model=True,
                                  model_save_dir=final_model_save_path)
        infer_model = create_model(sess, args, src_vocab_size, tgt_vocab_size, src_vocab_rev, tgt_vocab_rev,
                                   mode=constants.INFER, direction=args.nmt_direction, reuse=True)
        print("INFERENCE TYPE:%s" % args.decode_type)
        inference(infer_model, A=A, B=B, sess=sess, args=args, src_test_iterator=src_test_iterator,
                  src_test_next=src_test_iterator.get_next(),
                  result_dir=args.final_tsf_result_dir, src_vocab_rev=src_vocab_rev)


def predict_on_image_names(image_names, config, model_path_id="/home/data/output/resnet_v1_101_coco_fcis_end2end_ohem-nebraska/train-nebraska/e2e",epoch=8):
    import argparse
    import os
    import sys
    import logging
    import pprint
    import cv2
    from utils.image import resize, transform
    import numpy as np
    # get config
    os.environ['PYTHONUNBUFFERED'] = '1'
    os.environ['MXNET_CUDNN_AUTOTUNE_DEFAULT'] = '0'
    os.environ['MXNET_ENABLE_GPU_P2P'] = '0'
    cur_path = os.path.abspath(".")
    sys.path.insert(0, os.path.join(cur_path, '../external/mxnet', config.MXNET_VERSION))
    import mxnet as mx
    print("use mxnet at", mx.__file__)
    from core.tester import im_detect, Predictor
    from symbols import *
    from utils.load_model import load_param
    from utils.show_masks import show_masks
    from utils.tictoc import tic, toc
    from nms.nms import py_nms_wrapper
    from mask.mask_transform import gpu_mask_voting, cpu_mask_voting
    # get symbol
    ctx_id = [int(i) for i in config.gpus.split(',')]
    sym_instance = eval(config.symbol)()
    sym = sym_instance.get_symbol(config, is_train=False)

    # set up class names
    num_classes = 2
    classes = ['cp']

    # load demo data
    data = []
    for im_name in image_names:
        assert os.path.exists(im_name), ('%s does not exist'.format(im_name))
        im = cv2.imread(im_name, cv2.IMREAD_COLOR | cv2.IMREAD_IGNORE_ORIENTATION)
        target_size = config.SCALES[0][0]
        max_size = config.SCALES[0][1]
        im, im_scale = resize(im, target_size, max_size, 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})

    # get predictor
    data_names = ['data', 'im_info']
    label_names = []
    data = [[mx.nd.array(data[i][name]) for name in 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(data_names, data[i])] for i in xrange(len(data))]
    provide_label = [None for i in xrange(len(data))]

     # loading the last epoch that was trained, 8
    arg_params, aux_params = load_param(model_path_id, epoch, process=True)
    predictor = Predictor(sym, data_names, label_names,
                          context=[mx.gpu(ctx_id[0])], max_data_shapes=max_data_shape,
                          provide_data=provide_data, provide_label=provide_label,
                          arg_params=arg_params, aux_params=aux_params)

    all_classes = []
    all_configs = []
    all_masks = []
    all_dets = []
    all_ims = []

    # warm up
    for i in xrange(2):
        data_batch = mx.io.DataBatch(data=[data[0]], label=[], pad=0, index=0,
                                     provide_data=[[(k, v.shape) for k, v in zip(data_names, data[0])]],
                                     provide_label=[None])
        scales = [data_batch.data[i][1].asnumpy()[0, 2] for i in xrange(len(data_batch.data))]
        _, _, _, _ = im_detect(predictor, data_batch, data_names, scales, config)

    # test
    for idx, im_name in enumerate(image_names):
        data_batch = mx.io.DataBatch(data=[data[idx]], label=[], pad=0, index=idx,
                                     provide_data=[[(k, v.shape) for k, v in zip(data_names, data[idx])]],
                                     provide_label=[None])
        scales = [data_batch.data[i][1].asnumpy()[0, 2] for i in xrange(len(data_batch.data))]

        tic()
        scores, boxes, masks, data_dict = im_detect(predictor, data_batch, data_names, scales, config)
        im_shapes = [data_batch.data[i][0].shape[2:4] for i in xrange(len(data_batch.data))]

        if not config.TEST.USE_MASK_MERGE:
            all_boxes = [[] for _ in xrange(num_classes)]
            all_masks = [[] for _ in xrange(num_classes)]
            nms = py_nms_wrapper(config.TEST.NMS)
            for j in range(1, num_classes):
                indexes = np.where(scores[0][:, j] > 0.7)[0]
                cls_scores = scores[0][indexes, j, np.newaxis]
                cls_masks = masks[0][indexes, 1, :, :]
                try:
                    if config.CLASS_AGNOSTIC:
                        cls_boxes = boxes[0][indexes, :]
                    else:
                        raise Exception()
                except:
                    cls_boxes = boxes[0][indexes, j * 4:(j + 1) * 4]

                cls_dets = np.hstack((cls_boxes, cls_scores))
                keep = nms(cls_dets)
                all_boxes[j] = cls_dets[keep, :]
                all_masks[j] = cls_masks[keep, :]
            dets = [all_boxes[j] for j in range(1, num_classes)]
            masks = [all_masks[j] for j in range(1, num_classes)]
        else:
            masks = masks[0][:, 1:, :, :]
            im_height = np.round(im_shapes[0][0] / scales[0]).astype('int')
            im_width = np.round(im_shapes[0][1] / scales[0]).astype('int')
            print (im_height, im_width)
            boxes = clip_boxes(boxes[0], (im_height, im_width))
            result_masks, result_dets = gpu_mask_voting(masks, boxes, scores[0], num_classes,
                                                        100, im_width, im_height,
                                                        config.TEST.NMS, config.TEST.MASK_MERGE_THRESH,
                                                        config.BINARY_THRESH, ctx_id[0])

            dets = [result_dets[j] for j in range(1, num_classes)]
            masks = [result_masks[j][:, 0, :, :] for j in range(1, num_classes)]
        print ('testing {} {:.4f}s'.format(im_name, toc()))
        # visualize
        for i in xrange(len(dets)):
            keep = np.where(dets[i][:,-1]>0.7)
            dets[i] = dets[i][keep]
            masks[i] = masks[i][keep]

        all_classes.append(classes)
        all_configs.append(config)
        all_masks.append(masks)
        all_dets.append(dets)
        im = cv2.imread(im_name)
        all_ims.append(im)
    return all_ims, all_dets, all_masks, all_configs, all_classes