"""
 Copyright (c) 2020 Intel Corporation
 Licensed under the Apache License, Version 2.0 (the "License");
 you may not use this file except in compliance with the License.
 You may obtain a copy of the License at
      http://www.apache.org/licenses/LICENSE-2.0
 Unless required by applicable law or agreed to in writing, software
 distributed under the License is distributed on an "AS IS" BASIS,
 WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 See the License for the specific language governing permissions and
 limitations under the License.
"""
import networkx as nx
import torch

from nncf.dynamic_graph.graph import NNCFGraph
from nncf.dynamic_graph.operator_metatypes import NoopMetatype, HardTanhMetatype, TanhMetatype, RELUMetatype, \
    PRELUMetatype, ELUMetatype, GELUMetatype, SigmoidMetatype, SoftmaxMetatype, AvgPool2dMetatype, MaxPool2dMetatype, \
    DropoutMetatype, Conv1dMetatype, Conv2dMetatype, Conv3dMetatype, BatchNormMetatype, CatMetatype, AddMetatype, \
    SubMetatype, DivMetatype, MulMetatype, LinearMetatype, MatMulMetatype, MinMetatype, MaxMetatype, MeanMetatype
from nncf.nncf_logger import logger as nncf_logger
from nncf.nncf_network import NNCFNetwork
from nncf.pruning.export_utils import PruningOperationsMetatypeRegistry, identity_mask_propagation, get_input_masks, \
    fill_input_masks
from nncf.pruning.utils import get_sources_of_node

PRUNING_OPERATOR_METATYPES = PruningOperationsMetatypeRegistry("operator_metatypes")

# pylint: disable=protected-access
class DefaultMetaOp:
    subtypes = []
    additional_types = []

    def mask_propagation(self, model: NNCFNetwork, nx_node: dict, graph: NNCFGraph, nx_graph: nx.DiGraph):
        """
        Propagate mask through a node using masks of all inputs and pruning mask of current node (if any).
        Should set the following attributes:
        input_masks - list of masks of input nodes (None if there is no mask in some input)
        output_mask - resulting mask of nx_node operation
        accept_pruned_input - can this operation work with pruned input or not
        :param model: model to prune
        :param nx_node: node from networkx graph to propagate mask through it
        :param graph: graph of model to prune
        :param nx_graph: networkx graph
        """
        raise NotImplementedError

    def input_prune(self, model: NNCFNetwork, nx_node, graph: NNCFGraph, nx_graph: nx.DiGraph):
        """
        Prune nx_node by input_masks (if masks is not none and operation support it).
        """

    def output_prune(self, model: NNCFNetwork, nx_node, graph: NNCFGraph, nx_graph: nx.DiGraph):
        """
        Prune nx_node by output_mask (if mask is not none and operation support it).
        """

    @classmethod
    def get_all_op_aliases(cls):
        """
        :return: list of all aliases of types in metatype
        """
        op_types = []
        for subtype in cls.subtypes:
            op_types.extend(subtype.get_all_aliases())
        op_types = list(set(op_types)) + cls.additional_types
        return op_types

@PRUNING_OPERATOR_METATYPES.register('model_input')
class Input(DefaultMetaOp):
    subtypes = [NoopMetatype]

    def mask_propagation(self, model, nx_node, graph, nx_graph):
        nx_node['input_masks'] = []
        nx_node['output_mask'] = None
        nx_node['accept_pruned_input'] = False

@PRUNING_OPERATOR_METATYPES.register('identity_mask_propagation')
class IdentityMaskForwardOps(DefaultMetaOp):
    subtypes = [HardTanhMetatype, TanhMetatype, RELUMetatype, PRELUMetatype, ELUMetatype, GELUMetatype, SigmoidMetatype,
                SoftmaxMetatype, AvgPool2dMetatype, MaxPool2dMetatype, DropoutMetatype]
    additional_types = ["h_sigmoid", "h_swish"]

    def mask_propagation(self, model: NNCFNetwork, nx_node, graph: NNCFGraph, nx_graph: nx.DiGraph):
        identity_mask_propagation(nx_node, nx_graph)
        nx_node['accept_pruned_input'] = True

@PRUNING_OPERATOR_METATYPES.register('convolution')
class Convolution(DefaultMetaOp):
    subtypes = [Conv1dMetatype, Conv2dMetatype, Conv3dMetatype]

    def mask_propagation(self, model: NNCFNetwork, nx_node, graph: NNCFGraph, nx_graph: nx.DiGraph):
        output_mask = None
        accept_pruned_input = True
        is_depthwise = False
        input_masks = get_input_masks(nx_node, nx_graph)

        nncf_node = graph._nx_node_to_nncf_node(nx_node)
        node_module = model.get_module_by_scope(nncf_node.op_exec_context.scope_in_model)

        if node_module.pre_ops:
            output_mask = node_module.pre_ops['0'].op.binary_filter_pruning_mask

        # In case of group convs we can't prune by output filters
        if node_module.groups != 1:
            if node_module.weight.size(1) == 1:
                # Depthwise case
                is_depthwise = True
                output_mask = input_masks[0]
            else:
                accept_pruned_input = False
                output_mask = None

        nx_node['input_masks'] = input_masks
        nx_node['output_mask'] = output_mask
        nx_node['accept_pruned_input'] = accept_pruned_input
        nx_node['is_depthwise'] = is_depthwise

    def input_prune(self, model: NNCFNetwork, nx_node, graph: NNCFGraph, nx_graph: nx.DiGraph):
        input_mask = nx_node['input_masks'][0]
        if input_mask is None:
            return
        bool_mask = torch.tensor(input_mask, dtype=torch.bool)
        new_num_channels = int(torch.sum(input_mask))

        nncf_node = graph._nx_node_to_nncf_node(nx_node)
        node_module = model.get_module_by_scope(nncf_node.op_exec_context.scope_in_model)
        is_depthwise = nx_node['is_depthwise']
        old_num_clannels = int(node_module.weight.size(1))

        if is_depthwise:
            # In depthwise case prune output channels by input mask, here only fix for new number of input channels
            node_module.groups = new_num_channels
            node_module.in_channels = new_num_channels
        else:
            out_channels = node_module.weight.size(0)
            broadcasted_mask = bool_mask.repeat(out_channels).view(out_channels, bool_mask.size(0))
            new_weight_shape = list(node_module.weight.shape)
            new_weight_shape[1] = new_num_channels

            node_module.in_channels = new_num_channels
            node_module.weight = torch.nn.Parameter(node_module.weight[broadcasted_mask].view(new_weight_shape))

        nncf_logger.info('Pruned Convolution {} by input mask. Old input filters number: {}, new filters number:'
                         ' {}.'.format(nx_node['key'], old_num_clannels, new_num_channels))

    def output_prune(self, model: NNCFNetwork, nx_node, graph: NNCFGraph, nx_graph: nx.DiGraph):
        mask = nx_node['output_mask']
        if mask is None:
            return

        bool_mask = torch.tensor(mask, dtype=torch.bool)

        nncf_node = graph._nx_node_to_nncf_node(nx_node)
        node_module = model.get_module_by_scope(nncf_node.op_exec_context.scope_in_model)
        old_num_clannels = int(node_module.weight.size(0))

        node_module.out_channels = int(torch.sum(mask))
        node_module.weight = torch.nn.Parameter(node_module.weight[bool_mask])

        if node_module.bias is not None and not nx_node['is_depthwise']:
            node_module.bias = torch.nn.Parameter(node_module.bias[bool_mask])

        nncf_logger.info('Pruned Convolution {} by pruning mask. Old output filters number: {}, new filters number:'
                         ' {}.'.format(nx_node['key'], old_num_clannels, node_module.out_channels))

@PRUNING_OPERATOR_METATYPES.register('batch_norm')
class BatchNorm(DefaultMetaOp):
    subtypes = [BatchNormMetatype]

    def mask_propagation(self, model: NNCFNetwork, nx_node, graph: NNCFGraph, nx_graph: nx.DiGraph):
        identity_mask_propagation(nx_node, nx_graph)
        nx_node['accept_pruned_input'] = True

    def input_prune(self, model: NNCFNetwork, nx_node, graph: NNCFGraph, nx_graph: nx.DiGraph):
        input_mask = nx_node['input_masks'][0]
        if input_mask is None:
            return

        nncf_node = graph._nx_node_to_nncf_node(nx_node)
        node_module = model.get_module_by_scope(nncf_node.op_exec_context.scope_in_model)

        bool_mask = torch.tensor(input_mask, dtype=torch.bool)
        old_num_clannels = int(node_module.weight.size(0))
        new_num_channels = int(torch.sum(input_mask))

        node_module.num_features = new_num_channels
        node_module.weight = torch.nn.Parameter(node_module.weight[bool_mask])
        node_module.bias = torch.nn.Parameter(node_module.bias[bool_mask])
        node_module.running_mean = torch.nn.Parameter(node_module.running_mean[bool_mask], requires_grad=False)
        node_module.running_var = torch.nn.Parameter(node_module.running_var[bool_mask], requires_grad=False)

        nncf_logger.info('Pruned BatchNorm {} by input mask. Old num features: {}, new num features:'
                         ' {}.'.format(nx_node['key'], old_num_clannels, new_num_channels))

@PRUNING_OPERATOR_METATYPES.register('concat')
class Concat(DefaultMetaOp):
    subtypes = [CatMetatype]

    def all_inputs_from_convs(self, nx_node, nx_graph, graph):
        """
        Return whether all input sources of nx_node is convolutions or not
        :param nx_node: node to determine it's sources
        :param nx_graph:  networkx graph to work with
        :param graph:  NNCF graph to work with
        """
        inputs = [u for u, _ in nx_graph.in_edges(nx_node['key'])]
        input_masks = get_input_masks(nx_node, nx_graph)

        for i, inp in enumerate(inputs):
            # If input has mask ->  it went from convolution (source of this node is a convolution)
            if input_masks[i] is not None:
                continue
            nncf_input_node = graph._nx_node_to_nncf_node(nx_graph.nodes[inp])
            source_nodes = get_sources_of_node(nncf_input_node, graph, Convolution.get_all_op_aliases() +
                                               StopMaskForwardOps.get_all_op_aliases() +
                                               Input.get_all_op_aliases())
            sources_types = [node.op_exec_context.operator_name for node in source_nodes]
            if any([t in sources_types for t in StopMaskForwardOps.get_all_op_aliases()]):
                return False
        return True

    def check_concat(self, nx_node, nx_graph, graph):
        if self.all_inputs_from_convs(nx_node, nx_graph, graph):
            return True
        return False

    def mask_propagation(self, model: NNCFNetwork, nx_node, graph: NNCFGraph, nx_graph: nx.DiGraph):
        result_mask = None
        accept_pruned_input = False

        if self.check_concat(nx_node, nx_graph, graph):
            input_masks, filled_input_masks = fill_input_masks(nx_node, nx_graph)

            if all([mask is None for mask in input_masks]):
                result_mask = None
            else:
                result_mask = torch.cat(filled_input_masks)
                accept_pruned_input = True

        nx_node['input_masks'] = input_masks
        nx_node['output_mask'] = result_mask
        nx_node['accept_pruned_input'] = accept_pruned_input

@PRUNING_OPERATOR_METATYPES.register('elementwise')
class Elementwise(DefaultMetaOp):
    subtypes = [AddMetatype, SubMetatype, DivMetatype, MulMetatype]

    def mask_propagation(self, model: NNCFNetwork, nx_node, graph: NNCFGraph, nx_graph: nx.DiGraph):
        input_masks = get_input_masks(nx_node, nx_graph)

        nx_node['input_masks'] = input_masks
        nx_node['output_mask'] = None
        nx_node['accept_pruned_input'] = False

@PRUNING_OPERATOR_METATYPES.register('stop_propagation_ops')
class StopMaskForwardOps(DefaultMetaOp):
    subtypes = [MeanMetatype, MaxMetatype, MinMetatype, LinearMetatype, MatMulMetatype]

    def mask_propagation(self, model: NNCFNetwork, nx_node, graph: NNCFGraph, nx_graph: nx.DiGraph):
        input_masks = get_input_masks(nx_node, nx_graph)

        nx_node['input_masks'] = input_masks
        nx_node['output_mask'] = None
        nx_node['accept_pruned_input'] = False

class ModelPruner:
    CAN_PRUNE_ATTR = 'can_prune'

    def __init__(self, model: NNCFNetwork, graph: NNCFGraph, nx_graph: nx.DiGraph):
        self.model = model
        self.graph = graph
        self.nx_graph = nx_graph

    def node_propagate_can_prune_attr(self, nx_node_key):
        nx_node = self.nx_graph.nodes[nx_node_key]
        is_conv = self.graph.node_type_fn(nx_node) in Convolution().get_all_op_aliases()
        return not is_conv or (is_conv and nx_node['is_depthwise'])

    def node_accept_different_inputs(self, nx_node):
        """
        Return whether nx_node accept pruned and not pruned inputs as inputs at the same time.
        """
        node_type = self.graph.node_type_fn(nx_node)
        return node_type in Concat.get_all_op_aliases()

    @staticmethod
    def get_class_by_type_name(type_name):
        """
        Return class of metaop that corresponds to type_name type.
        """
        cls = PRUNING_OPERATOR_METATYPES.get_operator_metatype_by_op_name(type_name)
        if cls is None:
            raise RuntimeError("Class {} is not found".format(type_name))
            # cls = StopMaskForwardOps
        return cls

    def mask_propagation(self):
        """
        Mask propagation in graph:
        to propagate masks run method mask_propagation (of metaop of current node) on all nodes in topological order.
        """
        sorted_nodes = [self.nx_graph.nodes[node_name] for node_name in nx.topological_sort(self.nx_graph)]
        for node in sorted_nodes:
            node_type = self.graph.node_type_fn(node)
            cls = self.get_class_by_type_name(node_type)()
            cls.mask_propagation(self.model, node, self.graph, self.nx_graph)
        nncf_logger.info('Finished mask propagation in graph')

    def propagate_can_prune_attr_up(self):
        """
        Propagating can_prune attribute in reversed topological order.
        This attribute depends on accept_pruned_input and can_prune attributes of output nodes.
        Node can_prune is True if all outputs accept_pruned_input is True and all outputs
        (except convs because conv can be pruned by input and output independently).
        """
        for node_name in self.nx_graph.nodes:
            self.nx_graph.nodes[node_name][ModelPruner.CAN_PRUNE_ATTR] = True

        reversed_sorted_nodes = reversed([self.nx_graph.nodes[name] for name in nx.topological_sort(self.nx_graph)])
        for node in reversed_sorted_nodes:
            # Check all output nodes accept_pruned_input attribute
            out_edges = self.nx_graph.out_edges(node['key'])
            outputs_accept_pruned_input = all(self.nx_graph.nodes[key]['accept_pruned_input'] for _, key in out_edges)

            # Check all output nodes can_prune attribute
            outputs_will_be_pruned = all([self.nx_graph.nodes[key][ModelPruner.CAN_PRUNE_ATTR] for _, key in out_edges
                                          if self.node_propagate_can_prune_attr(key)])
            node[ModelPruner.CAN_PRUNE_ATTR] = outputs_accept_pruned_input and outputs_will_be_pruned

        nncf_logger.info('Propagated can_prune attribute up')

    def propagate_can_prune_attr_down(self):
        """
        Propagating can_prune attribute down to fix all branching cases with one pruned and one not pruned
        branches.
        """
        sorted_nodes = [self.nx_graph.nodes[name] for name in nx.topological_sort(self.nx_graph)]
        for node in sorted_nodes:
            # Propagate attribute only in not conv case
            if self.node_propagate_can_prune_attr(node['key']):
                in_edges = self.nx_graph.in_edges(node['key'])
                can_prune = all(self.nx_graph.nodes[key][ModelPruner.CAN_PRUNE_ATTR] for key, _ in in_edges)
                can_prune_any = any(self.nx_graph.nodes[key][ModelPruner.CAN_PRUNE_ATTR] for key, _ in in_edges)

                if (not self.node_accept_different_inputs(node) and not can_prune) or \
                        (self.node_accept_different_inputs(node) and not can_prune_any):
                    node[ModelPruner.CAN_PRUNE_ATTR] = can_prune

        nncf_logger.info('Propagated can_prune attribute down')

    def apply_mask(self):
        """
        Applying propagated masks for all nodes in topological order:
        if all inputs of node can_prune -> running input_prune method for this node
        if node[ModelPruner.CAN_PRUNE_ATTR] -> running output_prune method for this node
        """
        sorted_nodes = [self.nx_graph.nodes[name] for name in nx.topological_sort(self.nx_graph)]
        with torch.no_grad():
            for node in sorted_nodes:
                node_type = self.graph.node_type_fn(node)
                node_cls = self.get_class_by_type_name(node_type)()

                in_edges = self.nx_graph.in_edges(node['key'])
                can_prune_input = all(self.nx_graph.nodes[key][ModelPruner.CAN_PRUNE_ATTR] for key, _ in in_edges)
                if can_prune_input:
                    node_cls.input_prune(self.model, node, self.graph, self.nx_graph)

                if node[ModelPruner.CAN_PRUNE_ATTR]:
                    node_cls.output_prune(self.model, node, self.graph, self.nx_graph)
        nncf_logger.info('Finished mask applying step')

    def prune_model(self):
        """
        Model pruner work in three stages:
        1. Mask propagation: propagate pruning masks through the graph.
        2. Propagate can_prune attribute (up and then down) through the graph.This attribute shows can we really
        prune some node or not.
        3. Applying masks accordingly with can_prune attribute (only when can prune).
        :return:
        """
        nncf_logger.info('Start pruning model')
        self.mask_propagation()
        self.propagate_can_prune_attr_up()
        self.propagate_can_prune_attr_down()
        self.apply_mask()
        nncf_logger.info('Finished pruning model')


#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
Created on Mon Jul 26 16:22:22 2021

@author: Tom Hadfield

Class for docking elaborations using GOLD and doing operations with docked structures

"""

#import elabHotspotsBase
#import elabHotspotsCCDC

#########Standard Libraries##########
import json
import numpy as np
from numpy import linalg
from scipy import stats
import pandas as pd
import multiprocessing as mp #For parallelising
import time
import matplotlib.pyplot as plt
import sys
from random import sample
from functools import partial
import glob

#########RDKit Modules############
import rdkit
from rdkit import Chem
from rdkit.Chem import rdmolops
from rdkit.Chem import QED
from rdkit.Chem import rdFMCS
from rdkit.Chem import Draw
from rdkit.Chem import Descriptors
from rdkit.Chem import rdMolDescriptors
from rdkit.Chem import rdMMPA
from rdkit.Chem import AllChem
from rdkit.Chem import rdchem
from rdkit import DataStructs
from rdkit.Chem.Draw import IPythonConsole
from rdkit import RDLogger
RDLogger.DisableLog('rdApp.*') #Supress annoying RDKit output
from rdkit import RDConfig
from rdkit.Chem import ChemicalFeatures
from rdkit.Chem.QED import qed

import pickle
import os

from ccdc.docking import Docker
from ccdc.io import MoleculeReader, EntryReader
import tempfile

import sys
from argparse import ArgumentParser
from platform import platform
from pathlib import Path
from os import mkdir, chdir
from shutil import rmtree, copy as cp
from time import time
from dataclasses import dataclass, field
from multiprocessing import Pool

#Create feature factory
fdefName = os.path.join(RDConfig.RDDataDir,'BaseFeatures.fdef')
factory = ChemicalFeatures.BuildFeatureFactory(fdefName)

class docking:

    def __init__(self):
        pass

    def ligandEfficiency(self, mols, scores):
        #Calculate the ligand efficiency as docking_score/number_heavy_atoms

        smiles = []
        le = []

        for i, m in enumerate(mols):
            smiles.append(Chem.MolToSmiles(m))
            le.append(scores[i]/m.GetNumHeavyAtoms())

        leDF = pd.DataFrame({'smiles':smiles, 'ligEfficiency': le}).sort_values('ligEfficiency', ascending = False)
        return leDF

    def vectorDistance(self, p1, p2):
        return np.linalg.norm(p1 - p2)

    def assessAllDocks(self, mols, hotspot, single = True):
        smiles = []
        dist = []
        for m in mols:
            smiles.append(Chem.MolToSmiles(m))
            dist.append(self.assessSingleDock(m, hotspot, single))

        distanceDF = pd.DataFrame({'smiles':smiles, 'distance':dist}).sort_values('distance', ascending = True)
        return distanceDF

    def assessSingleDock(self, mol, hotspot, single = True):
        #Compute the distance between the pharmacophores in a mol and the corresponding hotspot

        if single == True:
            #i.e. we're comparing to a single hotspot
            #Check how close the pharmacophores in the docked molecule are to the hotspot of interest
            distanceToPharm = []
            feats = factory.GetFeaturesForMol(mol)
            for feat in feats:
                if feat.GetFamily() == hotspot['type']:
                    #Compute Distance to the  hotspot and take the pharmacophore with the smallest distance
                    pharmPosition = np.array(mol.GetConformer().GetAtomPosition(feat.GetAtomIds()[0]))
                    distanceToPharm.append(self.vectorDistance(pharmPosition, hotspot['position']))
            return min(distanceToPharm)
        else:
            #i.e. we're assessing closeness to multiple hotspots

            distances = []
            feats = factory.GetFeaturesForMol(mol)

            for k in hotspot.keys():
                pharmPoint = hotspot[k]
                distanceToPharm = []

                for feat in feats:
                    if feat.GetFamily() == pharmPoint['type']:
                        #Compute Distance to the hotspot and take the pharmacophore with the smallest distance
                        pharmPosition = np.array(mol.GetConformer().GetAtomPosition(feat.GetAtomIds()[0]))
                        distanceToPharm.append(self.vectorDistance(pharmPosition, pharmPoint['position']))
                if len(distanceToPharm) == 0:
                    #i.e. no matching pharmacophores
                    distances.append(100) #Append really big distance
                elif len(distanceToPharm) > 0:
                    distances.append(min(distanceToPharm))

            return max(distances) #Return the maximum of the distances to the hotspot centres

            '''
            distances = []
            for i in [1,2]:
                hotspotOfInterest = hotspot[i]
                feats = factory.GetFeaturesForMol(mol)
                distanceToPharm = []
                for feat in feats:
                   if feat.GetFamily() == hotspotOfInterest['type']:
                        #Compute Distance to the three hotspot and take the pharmacophore with the smallest distance
                        pharmPosition = np.array(mol.GetConformer().GetAtomPosition(feat.GetAtomIds()[0]))
                        distanceToPharm.append(self.vector_distance(pharmPosition, hotspotOfInterest['position']))

            if len(distanceToPharm) == 0:
                #i.e. there were no matching pharmacophores
                distances.append(100) #Append really big distance
            elif len(distanceToPharm) > 0:
                distances.append(min(distanceToPharm))
            return max(distances) #Return the maximum of the distances to the hotspot centres
            '''

    def computeDistanceToClosestHotspot(mol, hotspots, single = True):
        if single == True:
            allDistances = []
            for k in hotspots.keys():
                allDistances.append(assessSingleDock(mol, hotspots[k], True))

            return min(allDistances)
        else:
            return assessSingleDock(mol, hotspots, False)

    def flexDocking(self, ligandsToDockSDF, constraintFile, cavityLigandFile, protein, n_docks = 10, outputFile = 'docked_ligands.sdf', returnFitnessScore = False, conf_file = None, trackFlexibility = None):

        #for trackFlexibility - if not None, then should be a list: [resID, chiNumber] - #e.g. ['TYR119', 'chi2']

        if conf_file is None:
            return 'Please specify a conf file that supports flexible docking'

        else:
            #i.e. we use a conf_file to set up the docking (if we're going to be doing something like flexible docking

            #First we need to change the conf_file so that it has the right ligand file so we dock the correct ligands

            with open(conf_file, 'r') as f:
                cf = f.readlines()

            cf_new = []

            for l in cf:
                if 'ligand_data_file' in l:
                    cf_new.append(f'ligand_data_file {ligandsToDockSDF} {n_docks}\n')
                else:
                    cf_new.append(l)

            with open(conf_file, 'w') as f:
                for l in cf_new:
                    f.write(l)

            #Now the conf_file should have the correct ligand information but everything else should be the same

            #Set up docking with conf file
            settings = Docker.Settings.from_file(conf_file)
            docker = Docker(settings = settings)

        #Set scaffold constraint:
        scaffold =  MoleculeReader(constraintFile)[0]
        settings.add_constraint(settings.ScaffoldMatchConstraint(scaffold))

        settings.fitness_function = 'plp'
        settings.autoscale = 10.
        settings.early_termination = False
        batch_tempd = tempfile.mkdtemp()
        settings.output_directory = batch_tempd
        settings.output_file = outputFile

        '''
        #Define protein, binding site, scaffold and ligands to dock
        protein = settings.proteins[0]
        crystal_ligand = MoleculeReader(cavityLigandFile)[0]
        crystal_ligand.identifier = 'crystalLigand' #Change the identifier of the cavity ligand
        settings.binding_site = settings.BindingSiteFromLigand(protein, crystal_ligand, 10.0) #Define the binding site

        ligandsToDock = MoleculeReader(ligandsToDockSDF)
        settings.add_ligand_file(ligandsToDockSDF, 10) #Generate 10 poses per

        #Set scaffold constraint:
        scaffold =  MoleculeReader(constraintFile)[0]
        settings.add_constraint(settings.ScaffoldMatchConstraint(scaffold))
        '''

        #Dock the ligands
        results = docker.dock()

        #Get docked poses:
        docksLocation = results.ligands.file_name[0]
        topRanked = glob.glob(f'{docksLocation[0:docksLocation.index(outputFile)]}ranked*_m*_1.sdf')
        mols = [Chem.SDMolSupplier(fn)[0] for fn in topRanked]

        if trackFlexibility is not None:

            greatest_torsion_changes = []

            #Get Mol Indices
            for i in range(len(topRanked)):
                docksForOneLigand = glob.glob(f'{docksLocation[0:docksLocation.index(outputFile)]}ranked*_m{i + 1}_*.sdf')
                torsion_changes = []

                for dock in docksForOneLigand:
                    rt = EntryReader(dock)[0].attributes['Gold.Protein.RotatedTorsions']

                    for l in rt.split('\n'):
                        if trackFlexibility[0] in l and trackFlexibility[1] in l:
                            end_torsion = float(l[l.index('final') + 5:l.index('input')])
                            start_torsion = float(l[l.index('input') + 5:l.index('|library')])
                    torsion_changes.append(np.abs(end_torsion - start_torsion))

                greatest_torsion_changes.append(max(torsion_changes))

        locs = glob.glob(f'{docksLocation[0:docksLocation.index(outputFile)]}ranked*_m*.sdf')

        #Either return just the docks or the docks and associated fitness scores

        if returnFitnessScore and trackFlexibility is None:
            scores = [float(EntryReader(fn)[0].attributes['Gold.PLP.Fitness']) for fn in topRanked] #Get PLP scores for each of the top-ranked docks
            return mols, scores
        elif returnFitnessScore and trackFlexibility is not None:
            scores = [float(EntryReader(fn)[0].attributes['Gold.PLP.Fitness']) for fn in topRanked] #Get PLP scores for each of the top-ranked docks
            return mols, scores, greatest_torsion_changes
        else:
            return mols

    def dockLigands(self, ligandsToDockSDF, constraintFile, cavityLigandFile, protein, n_docks = 10, outputFile = 'docked_ligands.sdf', returnFitnessScore = False):

        #Method for the (rigid) docking of ligands, without Multiprocessing
        #For faster results, it is recommended that you use dockLigandsMP()
        #To perform docking with a flexible side chain, we recommend that you use flexDocking()

        #Set up docking class
        docker = Docker()
        settings = docker.settings
        settings.add_protein_file(protein)
        protein = settings.proteins[0]

        crystal_ligand = MoleculeReader(cavityLigandFile)[0]
        crystal_ligand.identifier = 'crystalLigand' #Change the identifier of the cavity ligand
        settings.binding_site = settings.BindingSiteFromLigand(protein, crystal_ligand, 10.0) #Define the binding site

        ligandsToDock = MoleculeReader(ligandsToDockSDF)
        settings.add_ligand_file(ligandsToDockSDF, 10) #Generate 10 poses per ligand

        #Set scaffold constraint:
        scaffold =  MoleculeReader(constraintFile)[0]
        settings.add_constraint(settings.ScaffoldMatchConstraint(scaffold))

        settings.fitness_function = 'plp'
        settings.autoscale = 10.
        settings.early_termination = False
        batch_tempd = tempfile.mkdtemp()
        settings.output_directory = batch_tempd
        settings.output_file = outputFile

        #Dock the ligands
        results = docker.dock()

        #Get docked poses:
        docksLocation = results.ligands.file_name[0]
        topRanked = glob.glob(f'{docksLocation[0:docksLocation.index(outputFile)]}ranked*_m*_1.sdf')
        mols = [Chem.SDMolSupplier(fn)[0] for fn in topRanked]

        #Either return just the docks or the docks and associated fitness scores

        if returnFitnessScore:
            scores = [float(EntryReader(fn)[0].attributes['Gold.PLP.Fitness']) for fn in topRanked] #Get PLP scores for each of the top-ranked docks
            return mols, scores
        else:
            return mols

    def dockLigandsMP(self, ligandsToDockSDF, constraintFile, cavityLigandFile, protein, n_processes = 1, ndocks = 10, outputFile = 'docked_ligands.sdf', returnFitnessScore = False):

        """
        Dock the molecules from the supplied input file in parallel.
        Adapted from the gold_multi_map.py file supplied by ccdc
        """

        t0 = time()  # Script start time

        if not n_processes > 0:
            print(f"Error! Number of processes must be an integer greater than zero.")
            sys.exit(1)

        #Get number of molecules to be docked
        with EntryReader(ligandsToDockSDF) as reader:
            n_molecules = len(reader)

        print(f"There are {n_molecules} molecules to dock on {n_processes} processes...")

        output_dir = tempfile.mkdtemp() #Create temporary output file

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

        # Determine the sets of parameters defining the chunks...
        chunks = []  # List of records that define the chunks

        # Work out the size of each chunk, and hence the start and finish indices in the input file...
        basic_size = n_molecules // n_processes  # Basic size of a chunk, which must obviously be integral
        remainder  = n_molecules %  n_processes  # Number of molecules that would not be included in basic-sized chunks

        finish = 0  # Finish index

        for n in range(1, n_processes + 1):  # Recall that the number of chunks is the same as the number of processes
            start = finish + 1  # Start index
            chunk_size = basic_size + 1 if n <= remainder else basic_size  # Add one to the basic chunk sizes until the remainder are taken care of
            finish = start + chunk_size - 1
            chunks.append(chunkForMP(n=n, start=start, finish=finish, output_dir=output_dir, crystal_ligand = cavityLigandFile, constraint = constraintFile, protein = protein, ligand_file = ligandsToDockSDF, ndocks=ndocks))
            print(f"chunk {n}: size: {chunk_size}; start: {start}, finish: {finish}.")

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

        # Dock the chunks in parallel...
        with Pool(n_processes) as pool:
            _ = pool.map(do_chunk, chunks)  # No output; docks are saved in the output_dir directory

        #print(f"Finished docking in {time() - t0:.1f} seconds.")

        #Now return the top ranked mols
        topRanked = sorted(glob.glob(f'{output_dir}/*/ranked*_m*_1.sdf')) #sort the list for consistency
        mols = [Chem.SDMolSupplier(fname)[0] for fname in topRanked]

        # All done.
        print(f"Finished in {time() - t0:.1f} seconds.")

        if returnFitnessScore:
            scores = [float(EntryReader(fn)[0].attributes['Gold.PLP.Fitness']) for fn in topRanked] #Get PLP scores for each of the top-ranked docks
            return mols, scores
        else:
            return mols

    def prepareForDocking(self, Hotspot, core, fname):

        #Hotspot is a HotspotSingle or HotspotMulti

        Hotspot.dockingFname = fname #Save the name of the file we'll save the hotspots to.
        Hotspot.countsIdx, Hotspot.countsSDFIdx = self.getSDFs(Hotspot.profileElabs, core, fname) #Saves sdfs to fname

        return Hotspot

    def getMultipleSDFs(self, data, core, outNameCore):

        #Iterate over the different indices (i.e. the different pharmacophoric profiles) and pass to getSDFs

        for idx in data['idx'].drop_duplicates():
            d = data.loc[data['idx'] == idx]
            self.getSDFs(d, core, f'{outNameCore}_p{idx}.sdf')

    def getSDFs(self, data, core, outName):
        #Core is the embedded mol we use for the constrained embedding

        mols = []
        dfIdx = []
        sdfIdx = []
        sdfIdxNum = 0
        for idx, g in enumerate(list(data['gen'])):
            dfIdx.append(idx)
            m = Chem.MolFromSmiles(g)

            try:
                AllChem.ConstrainedEmbed(m, core)
                fail = 0
            except:
                fail = 1
                print(f'Unable to do constrained embedding on the following molecule: {g}')

            #Now addHs and reembed:
            mHs = Chem.AddHs(m)
            try:
                AllChem.ConstrainedEmbed(mHs, m)
                fail = 0
            except:
                fail = 1
                print(f'Unable to do constrained embedding on the following molecule: {g}')

            if fail == 0:
                mols.append(mHs)
            else:
                mols.append('Could not embed molecule - not docking')

        #Now write mols to file
        w = Chem.SDWriter(outName)
        for m in mols:
            if m != 'Could not embed molecule - not docking':
                w.write(m)
                sdfIdx.append(sdfIdxNum)
                sdfIdxNum += 1
            else:
                sdfIdx.append(-1) #Indicates that we weren't able to dock this molecule

        return dfIdx, sdfIdx

    def qed_flag(self, full, frag, threshold = 0):
        #returns 1 if the elaborated mol has a higher qed than the original fragment and 0 otherwise

        qed_frag = qed(frag)
        qed_full = qed(full)

        if qed_full - qed_frag > threshold:
            return 1
        else:
            return 0

####Additional Code for parallel docking####

@dataclass
class chunkForMP:

    #Contains the information to provide to GOLD for docking
    #Store a bunch of chunks in a list and use them to do the docking in parallel

    n: int       # Chunk number
    start: int   # Index of first molecule in chunk
    finish: int  # Index of last molecule in chunk
    ndocks: int  # Number of docking poses to generate
    #conf_file: Path   # GOLD configuration file
    output_dir: Path  # Output dir, in which the chunk sub-directory will be created
    crystal_ligand: Path #Location of the crystal ligand (used for binding site definition)
    constraint: Path #Location of mol2 file used to constrain the docking
    protein: Path #Location of protein pdb File
    ligand_file: Path #Location of ligands to be docked

    dir: Path = field(init=False) # Sub-directory for chunk, see __post_init__ below

    def __post_init__(self):
        self.dir = f'{self.output_dir}/chunk_{self.n:02d}'

def do_chunk(chunk):

    """
    Dock a chunk of the input file.

    :param chunk: a record holding the parameters defining the chunk

    As we can't return a GOLD results object from a pool process (it can't be pickled as it wraps C++ objects),
    we simply return a boolean recording whether GOLD exited with a 'success' status code.
    """

    #Set up docking class
    docker = Docker()
    settings = docker.settings
    settings.add_protein_file(chunk.protein)

    settings.fitness_function = 'plp'
    settings.autoscale = 10.
    settings.early_termination = False
#     batch_tempd = tempfile.mkdtemp()
#     settings.output_directory = batch_tempd
#     settings.output_file = outputFile

    # Create and enter the sub-directory for this chunk...
    mkdir(chunk.dir)
    chdir(chunk.dir)
    settings.output_directory = '.'  # Ensure GOLD writes output to the chunk sub-directory

    #Define protein, binding site, scaffold and ligands to dock
    protein = settings.proteins[0]
    crystal_ligand = MoleculeReader(chunk.crystal_ligand)[0]
    crystal_ligand.identifier = 'crystalLigand' #Change the identifier of the cavity ligand
    settings.binding_site = settings.BindingSiteFromLigand(protein, crystal_ligand, 10.0) #Define the binding site

    # Specify the chunk of molecules to dock...
#     ligand_file = settings.ligand_files[0]  # The ligand file info will be overwritten, so store for reference below
    settings.clear_ligand_files() #Clear just to be sure
    settings.add_ligand_file(chunk.ligand_file, ndocks=chunk.ndocks, start=chunk.start, finish=chunk.finish)

    #Set scaffold constraint:
    scaffold =  MoleculeReader(chunk.constraint)[0]
    settings.add_constraint(settings.ScaffoldMatchConstraint(scaffold))

    # Run docking...
    print(f"Starting chunk {chunk.n} (ligand indices {chunk.start} - {chunk.finish})...")
    docker = Docker(settings=settings)
    results = docker.dock()

    print(f"Finished chunk {chunk.n}.")
    # As we can't return the results (as they are not picklable) and the poses have already been written to disk, we just return the status code

    return results.ligands.file_name[0] #return docks location


# DQN
import random
import numpy as np
import cv2
import PIL
from keras.models import Model
from keras.layers import Conv2D, MaxPooling2D, Dense, Lambda, Flatten, Input, average, Add
from keras.backend import repeat_elements
from keras.optimizers import Adam
import keras
import tensorflow as tf
from collections import deque

class DQNAgent:

    def __init__(self, inp_action_size):
        self.action_size = inp_action_size
        self.name="DDDQN"

        self.memory = deque(maxlen=2000)

        self.gamma = 0.95
        self.epsilon = 0.6
        self.epsilon_decay = 0.98
        self.epsilon_min = 0.01
        self.learning_rate = 0.0005

        self.tau = 200
        self.tau_counter = 0
        self.DQNetwork = self._build_model(self.name+"-l")
        self.TargetNetwork = self._build_model(self.name+"-t")
        self.sync_networks(0)

        # self.model = self._build_model()
        self.store_list = []

    def _build_model(self,name,  input_shape=(460, 460, 1)):
        h, w, d=input_shape
        # Constructing a Duelling Deep Q-Network
        # model = Sequential()
        # model.add(Lambda(lambda x: x / 127.5 - 1., input_shape=input_shape, output_shape=input_shape))
        # model.add(Conv2D(32, (3, 3), activation='relu', name='conv1', input_shape=input_shape, padding='same'))
        # model.add(MaxPooling2D(pool_size=(2, 2)))
        # model.add(Conv2D(64, (4, 4), activation='relu', name='conv2', padding='same', strides=(2, 2)))
        # model.add(MaxPooling2D(pool_size=(2, 2)))
        # model.add(Flatten())
        # model.add(Dense(256, activation='relu'))
        # model.add(Dense(2, activation='linear'))
        # inputs = Input(shape=input_shape)
        # x = Lambda(lambda x: x / 127.5 - 1, input_shape=input_shape, output_shape=input_shape)(inputs)
        # x = Conv2D(32, (3, 3), activation='relu', name='conv1', padding='same', input_shape=input_shape)(x)
        # x = MaxPooling2D(pool_size=(2, 2))(x)
        # x = Conv2D(64, (4, 4), activation='relu', name='conv2', padding='same')(x)
        # x = MaxPooling2D(pool_size=(2, 2))(x)
        # x = Flatten()(x)
        #
        # value = Dense(1, activation='relu')(x)
        # action_advantage = Dense(self.action_size, activation='relu')(x)
        # added=keras.layers.add([action_advantage, action_advantage])
        # # averaged=keras.layers.average([action_advantage, action_advantage])
        # # multiplied=repeat_elements(value, self.action_size, 0)
        # average_action_advantage = keras.backend.mean(action_advantage, axis=[-1], keepdims=True)
        # keras.layers.average(action_advantage)
        # # k=repeat_elements(value, self.action_size, -1)
        # # Q_values = Add()([repeat_elements(value, self.action_size, -1), average_action_advantage])
        #
        # model = Model(inputs, average_action_advantage)
        # model.compile(loss='mse', optimizer=Adam(lr=self.learning_rate))
        #
        # model.summary()
        #
        # return model
        # We use tf.variable_scope here to know which network we're using (DQN or target_net)
        # it will be useful when we will update our w- parameters (by copy the DQN parameters)
        with tf.variable_scope(name):
            # We create the placeholders
            # *state_size means that we take each elements of state_size in tuple hence is like if we wrote
            # [None, 100, 120, 4]
            self.inputs_ = tf.placeholder(tf.float32, [None, h,w,d], name="inputs")

            #
            self.ISWeights_ = tf.placeholder(tf.float32, [None, 1], name='IS_weights')

            self.actions_ = tf.placeholder(tf.float32, [None, self.action_size], name="actions_")

            # Remember that target_Q is the R(s,a) + ymax Qhat(s', a')
            self.target_Q = tf.placeholder(tf.float32, [None], name="target")

            """
            First convnet:
            CNN
            ELU
            """
            # Input is 100x120x4
            self.conv1 = tf.layers.conv2d(inputs=self.inputs_,
                                          filters=32,
                                          kernel_size=[8, 8],
                                          strides=[4, 4],
                                          padding="VALID",
                                          kernel_initializer=tf.contrib.layers.xavier_initializer_conv2d(),
                                          name="conv1")

            self.conv1_out = tf.nn.elu(self.conv1, name="conv1_out")

            """
            Second convnet:
            CNN
            ELU
            """
            self.conv2 = tf.layers.conv2d(inputs=self.conv1_out,
                                          filters=64,
                                          kernel_size=[4, 4],
                                          strides=[2, 2],
                                          padding="VALID",
                                          kernel_initializer=tf.contrib.layers.xavier_initializer_conv2d(),
                                          name="conv2")

            self.conv2_out = tf.nn.elu(self.conv2, name="conv2_out")

            """
            Third convnet:
            CNN
            ELU
            """
            self.conv3 = tf.layers.conv2d(inputs=self.conv2_out,
                                          filters=128,
                                          kernel_size=[4, 4],
                                          strides=[2, 2],
                                          padding="VALID",
                                          kernel_initializer=tf.contrib.layers.xavier_initializer_conv2d(),
                                          name="conv3")

            self.conv3_out = tf.nn.elu(self.conv3, name="conv3_out")

            self.flatten = tf.layers.flatten(self.conv3_out)

            ## Here we separate into two streams
            # The one that calculate V(s)
            self.value_fc = tf.layers.dense(inputs=self.flatten,
                                            units=512,
                                            activation=tf.nn.elu,
                                            kernel_initializer=tf.contrib.layers.xavier_initializer(),
                                            name="value_fc")

            self.value = tf.layers.dense(inputs=self.value_fc,
                                         units=1,
                                         activation=None,
                                         kernel_initializer=tf.contrib.layers.xavier_initializer(),
                                         name="value")

            # The one that calculate A(s,a)
            self.advantage_fc = tf.layers.dense(inputs=self.flatten,
                                                units=512,
                                                activation=tf.nn.elu,
                                                kernel_initializer=tf.contrib.layers.xavier_initializer(),
                                                name="advantage_fc")

            self.advantage = tf.layers.dense(inputs=self.advantage_fc,
                                             units=self.action_size,
                                             activation=None,
                                             kernel_initializer=tf.contrib.layers.xavier_initializer(),
                                             name="advantages")

            # Agregating layer
            # Q(s,a) = V(s) + (A(s,a) - 1/|A| * sum A(s,a'))
            self.output = self.value + tf.subtract(self.advantage,
                                                   tf.reduce_mean(self.advantage, axis=1, keepdims=True))

            # Q is our predicted Q value.
            self.Q = tf.reduce_sum(tf.multiply(self.output, self.actions_), axis=1)

            # The loss is modified because of PER
            self.absolute_errors = tf.abs(self.target_Q - self.Q)  # for updating Sumtree

            self.loss = tf.reduce_mean(self.ISWeights_ * tf.squared_difference(self.target_Q, self.Q))

            self.optimizer = tf.train.RMSPropOptimizer(self.learning_rate).minimize(self.loss)
            return tf.get_default_graph()

    # The weights from the target network are updated to the DQN-weights
    def sync_networks(self, sync_num):
        name = "Models/model_sync_{}.h5".format(sync_num)
        self.save(name, self.DQNetwork)
        self.load(name, self.TargetNetwork)

    def pre_process(self, image):
        image = np.array(image)
        # print('image', image)
        gray = cv2.cvtColor(image[30:490, 0:460], cv2.COLOR_BGR2GRAY)
        gray = gray.reshape(1, 460, 460, 1)
        # gray1=gray.reshape(460, 460)
        # PIL.Image.fromarray(gray1).show()
        return gray

    def _remember(self, state, action, reward, next_state, done):
        self.memory.append((state, action, reward, next_state, done))

    def _store(self, state, action, reward, done):
        self.store_list.append([state, action, reward, done])

    # def store_to_remember(self):
    #     self.store_list[-1][3] = True
    #     self.store_list[-1][2] = -10
    #     for x in range(len(self.store_list) - 1):
    #         elem = self.store_list[x]
    #         next_elem = self.store_list[x + 1]
    #         self._remember(elem[0], elem[1], next_elem[2], next_elem[0], next_elem[3])
    #     self.store_list = []

    # def act_store(self, state, reward, done):
    #     state = self.pre_process(state)
    #     action = self.act(state)
    #     self._store(state, action, reward, done)
    #     return action

    def predict(self, state):
        return self.TargetNetwork.predict(state)

    def act(self, state, pre_process=False):
        if pre_process:
            state = self.pre_process(state)
        if np.random.rand() <= self.epsilon:
            return random.randrange(self.action_size)
        act_values = self.TargetNetwork.predict(state)
        return np.argmax(act_values[0])

    def replay(self, batch_size, minibatch):
        # minibatch = random.sample(self.memory, batch_size)
        state, action, reward, next_state, done = minibatch

        for state, action, reward, next_state, done in minibatch:

            self.tau_counter += 1
            # Updating target network - part of fixed Q-targets
            if self.tau_counter == self.tau:
                self.tau_counter = 0
                self.sync_networks(random.randint(0, 100000000))

            target = reward
            if not done:
                print('pred', self.TargetNetwork.predict(next_state))

                # Double learning
                target_act = np.argmax(self.DQNetwork.predict(next_state))
                target = reward + self.gamma * self.TargetNetwork.predict(next_state)[0][target_act]
            print('target', target, self.epsilon)
            target_f = self.DQNetwork.predict(state)
            target_f[0][action] = target

            train_history = self.DQNetwork.fit(state, target_f, epochs=1, verbose=0)

        if self.epsilon > self.epsilon_min:
            self.epsilon *= self.epsilon_decay

    def load(self, name, model):
        model.load_weights(name)

    def save(self, name, model):
        with tf.Session() as sess:
            sess.run(model)
            tf.train.Saver().save(sess, name)


#!/usr/bin/env python3

import argparse
import math
import os
import random
import shutil
import sys

import matplotlib as mpl
import numpy as np
import torch
import torch.multiprocessing as mp
import torch.nn as nn

from common.config import Config
from common.environment import initialize_environment
from common.network import build_encoder
from common.utils import check_devices, check_logging
from sac import build_model, train, test, test_render

mpl.use('Agg')

try:
    mp.set_start_method('spawn', force=True)
except RuntimeError:
    pass

ROOT_DIR = os.path.dirname(os.path.abspath(__file__))

def get_config():
    def gpu_type(device):
        try:
            return int(device)
        except ValueError:
            return device

    parser = argparse.ArgumentParser(description='Train or test Soft Actor-Critic controller.')
    parser.add_argument('--mode', type=str, choices=['train', 'test', 'test_render'], default='train',
                        help='mode (default: train)')
    parser.add_argument('--gpu', type=gpu_type, default=None, nargs='+', metavar='CUDA_DEVICE',
                        help="GPU device indexes. if --sampler-gpu option is specified, gpu in this option is not used in sampler."
                             "(int for CUDA device or 'c'/'cpu' for CPU) "
                             "(use 'cuda:0' if no following arguments; use CPU if not present)")
    parser.add_argument('--sampler-gpu', type=gpu_type, default=None, nargs='+', metavar='SAMPLER_CUDA_DEVICE',
                        help="GPU device indexes to use for running sampler. if this option is not specified, use --gpu option"
                             "(int for CUDA device or 'c'/'cpu' for CPU) "
                             "(use 'cuda:0' if no following arguments; use CPU if not present)")
    parser.add_argument('--env', type=str, default='Pendulum-v1',
                        help='environment to train on (default: Pendulum-v1)')
    parser.add_argument('--n-frames', type=int, default=1,
                        help='concatenate original N consecutive observations as a new observation (default: 1)')
    parser.add_argument('--n-repeat-actions', type=int, default=1,
                        help='number of repeat action (default: 1, which is no repeat)')
    parser.add_argument('--n-past-actions', type=int, default=1,
                        help='number of past action that will be used as part of state (default: 1 which is no past action state)')
    parser.add_argument('--render', action='store_true',
                        help='render the environment')
    parser.add_argument('--vision-observation', action='store_true',
                        help='use rendered images as observation')
    parser.add_argument('--image-size', type=int, default=[96], metavar='SIZE', nargs='*',
                        help='image size of vision observation (default: 96)')
    parser.add_argument('--grayscale', action='store_true',
                        help='use grayscale image as an observation')
    parser.add_argument('--hidden-dims', type=int, default=[], nargs='+', metavar='DIM',
                        help='hidden dimensions of FC controller')
    parser.add_argument('--activation', type=str, choices=['ReLU', 'LeakyReLU', 'Tanh', 'ELU', 'SiLU'], default='ReLU',
                        help='activation function in controller networks (default: ReLU)')
    parser.add_argument('--dry-run-init-env', action='store_true', default=False,
                        help='dry run when initialize environment for the first time only')
    parser.add_argument('--half-training', action='store_true', default=False,
                        help='train in native FP16')
    carla_environment_group = parser.add_argument_group('carla environment')
    carla_environment_group.add_argument('--camera-size', type=int, default=None, nargs='*',
                                         help='size of camera in environment (default: same as image size)')
    carla_environment_group.add_argument('--camera-fov', type=int, default=None,
                                         help='camera FOV (default: None)')
    carla_environment_group.add_argument('--fps-mode', type=str, choices=['low', 'high'], default='HIGH',
                                         help='set FPS of CARLA. low is 5FPS, high is 10FPS (default: high)')
    encoder_group = parser.add_argument_group('state encoder')
    encoder_group.add_argument('--encoder-arch', type=str, choices=['FC', 'RNN', 'CNN', 'VAE', 'BETAVAE', 'EFFICIENTNET', 'RESNET', 'TINY_CNN'], default='FC',
                               help='architecture of state encoder network (default: FC)')
    encoder_group.add_argument('--separate-encoder', action='store_true', default=False,
                               help='use separate encoder for Actor and Critic (default: False)')
    encoder_group.add_argument('--state-dim', type=int, default=None, metavar='DIM',
                               help='target state dimension of encoded state '
                                    '(use env.observation_space.shape if not present)')
    encoder_group.add_argument('--encoder-activation', type=str, choices=['ReLU', 'LeakyReLU', 'Tanh', 'ELU', 'SiLU'], metavar='ACTIVATION',
                               help='activation function in state encoder networks '
                                    '(use activation function in controller if not present)')
    fc_encoder_group = parser.add_argument_group('FC state encoder')
    fc_encoder_group.add_argument('--encoder-hidden-dims', type=int, default=[], nargs='+', metavar='DIM',
                                  help='hidden dimensions of FC state encoder')
    rnn_encoder_group = parser.add_argument_group('RNN state encoder')
    rnn_encoder_group.add_argument('--encoder-hidden-dims-before-rnn', type=int, default=[], nargs='+', metavar='DIM',
                                   help='hidden FC dimensions before GRU layers in RNN state encoder')
    rnn_encoder_group.add_argument('--encoder-hidden-dims-rnn', type=int, default=[], nargs='+', metavar='DIM',
                                   help='GRU hidden dimensions of RNN state encoder')
    rnn_encoder_group.add_argument('--encoder-hidden-dims-after-rnn', type=int, default=[], nargs='+', metavar='DIM',
                                   help='hidden FC dimensions after GRU layers in RNN state encoder')
    rnn_encoder_group.add_argument('--skip-connection', action='store_true', default=False,
                                   help='add skip connection beside GRU layers in RNN state encoder')
    rnn_encoder_group.add_argument('--trainable-hidden', action='store_true', default=False,
                                   help='set initial hidden of GRU layers trainable '
                                        '(use zeros as initial hidden if not present)')
    rnn_encoder_group.add_argument('--step-size', type=int, default=16,
                                   help='number of continuous steps for update (default: 16)')
    cnn_encoder_group = parser.add_argument_group('CNN state encoder')
    cnn_encoder_group.add_argument('--encoder-hidden-channels', type=int, default=[], nargs='+', metavar='CHN',
                                   help='channels of hidden conv layers in CNN state encoder')
    cnn_encoder_group.add_argument('--kernel-sizes', type=int, default=[], nargs='+', metavar='K',
                                   help='kernel sizes of conv layers in CNN state encoder (defaults: 3)')
    cnn_encoder_group.add_argument('--strides', type=int, default=[], nargs='+', metavar='S',
                                   help='strides of conv layers in CNN state encoder (defaults: 1)')
    cnn_encoder_group.add_argument('--paddings', type=int, default=[], nargs='+', metavar='P',
                                   help='paddings of conv layers in CNN state encoder (defaults: K // 2)')
    cnn_encoder_group.add_argument('--poolings', type=int, default=[], nargs='+', metavar='K',
                                   help='max pooling kernel size after activation function in CNN state encoder (defaults: 1)')
    cnn_encoder_group.add_argument('--batch-normalization', action='store_true', default=False,
                                   help='use batch normalization in CNN state encoder')
    vae_encoder_group = parser.add_argument_group('VAE state encoder')
    vae_encoder_group.add_argument('--weight-path', type=str, default=None, help='path to weight file')
    parser.add_argument('--max-episode-steps', type=int, default=10000,
                        help='max steps per episode (default: 10000)')
    parser.add_argument('--n-epochs', type=int, default=1000,
                        help='number of training epochs (default: 1000)')
    parser.add_argument('--n-episodes', type=int, default=100,
                        help='number of test episodes (default: 100)')
    parser.add_argument('--n-updates', type=int, default=256,
                        help='number of learning updates per epoch (default: 256)')
    parser.add_argument('--batch-size', type=int, default=256,
                        help='batch size (default: 256)')
    parser.add_argument('--n-samplers', type=int, default=4,
                        help='number of parallel samplers (default: 4)')
    parser.add_argument('--buffer-capacity', type=int, default=1000000, metavar='CAPACITY',
                        help='capacity of replay buffer (default: 1000000)')
    parser.add_argument('--update-sample-ratio', type=float, default=2.0, metavar='RATIO',
                        help='speed ratio of training and sampling '
                             '(sample speed <= training speed / ratio (ratio should be larger than 1.0)) '
                             '(default: 2.0)')
    parser.add_argument('--gamma', type=float, default=0.99,
                        help='discount factor for rewards (default: 0.99)')
    parser.add_argument('--soft-tau', type=float, default=0.01, metavar='TAU',
                        help='soft update factor for target networks (default: 0.01)')
    parser.add_argument('--n-bootstrap-step', type=int, default=1, metavar='BOOTSTRAP',
                        help='n-step return bootstrap (default: 1)')
    parser.add_argument('--normalize-rewards', action='store_true',
                        help='normalize rewards for training')
    parser.add_argument('--reward-scale', type=float, default=1.0, metavar='SCALE',
                        help='reward scale factor for normalized rewards (default: 1.0)')
    parser.add_argument('--use-popart', default=False, action='store_true',
                        help='normalize target using POP-ART')
    parser.add_argument('--beta', type=float, default=1e-4, metavar='BETA',
                        help='beta parameter of POP-ART (default: 1e-4)')
    parser.add_argument('--deterministic', action='store_true', help='deterministic in evaluation')
    lr_group = parser.add_argument_group('learning rate')
    lr_group.add_argument('--lr', type=float, default=1E-4,
                          help='learning rate (can be override by the following specific learning rate) (default: 0.0001)')
    lr_group.add_argument('--critic-lr', type=float, default=None,
                          help='learning rate for critic networks (use LR above if not present)')
    lr_group.add_argument('--actor-lr', type=float, default=None,
                          help='learning rate for actor networks (use LR above if not present)')
    alpha_group = parser.add_argument_group('temperature parameter')
    alpha_group.add_argument('--alpha-lr', type=float, default=None,
                             help='learning rate for temperature parameter (use ACTOR_LR above if not present)')
    alpha_group.add_argument('--initial-alpha', type=float, default=1.0, metavar='ALPHA',
                             help='initial value of temperature parameter (default: 1.0)')
    alpha_group.add_argument('--adaptive-entropy', action='store_true',
                             help='auto update temperature parameter while training')
    alpha_group.add_argument('--target-entropy', type=float, default=None,
                             help='target entropy to adapt to (default: -dim(action))')
    parser.add_argument('--actor-update-frequency', type=int, default=1, metavar='ACTOR_UPDATE_FREQUENCY',
                        help='frequency of actor update per critic update (default: 1)')
    parser.add_argument('--weight-decay', type=float, default=0.0,
                        help='weight decay (default: 0.0)')
    parser.add_argument('--clip-gradient', action='store_true',
                        help='clip gradient on optimizer step')
    parser.add_argument('--random-seed', type=int, default=0, metavar='SEED',
                        help='random seed (default: 0)')
    parser.add_argument('--log-episode-video', action='store_true',
                        help='save rendered episode videos to TensorBoard logs')
    parser.add_argument('--log-dir', type=str, default=os.path.join(ROOT_DIR, 'logs'),
                        help='folder to save TensorBoard logs')
    parser.add_argument('--checkpoint-dir', type=str, default=os.path.join(ROOT_DIR, 'checkpoints'),
                        help='folder to save checkpoint')
    parser.add_argument('--checkpoint-save-frequency', type=int, default=10,
                        help='how often to save checkpoint(default: 10)')
    parser.add_argument('--load-checkpoint', action='store_true',
                        help='load latest checkpoint in checkpoint dir')
    parser.add_argument('--epoch-number', type=int, default=None,
                        help='number of epoch to load')
    parser.add_argument('--record-video', action='store_true',
                        help='record video directly from environment')
    parser.add_argument('--video-dir', type=str, default=None,
                        help='path to directory to store recorded video')
    args = parser.parse_args()
    if len(sys.argv) == 1:
        parser.print_help()
        exit()

    config = Config(vars(args))

    return config

def initialize(config):
    random.seed(config.random_seed)
    np.random.seed(config.random_seed)
    torch.manual_seed(config.random_seed)

    initialize_hyperparameters(config)
    initialize_environment(config)
    build_encoder(config)
    check_devices(config)
    check_logging(config)
    copy_environment_file(config)

def initialize_hyperparameters(config):
    config.activation = {
        'ReLU': nn.ReLU(inplace=True),
        'LeakyReLU': nn.LeakyReLU(negative_slope=0.3, inplace=True),
        'Tanh': nn.Tanh(),
        'ELU': nn.ELU(inplace=True),
        'SiLU': nn.SiLU(inplace=True)
    }.get(config.activation)

    config.FC_encoder = (config.encoder_arch == 'FC')
    config.RNN_encoder = (config.encoder_arch == 'RNN')
    config.CNN_encoder = (config.encoder_arch == 'CNN')

    config.VAE_encoder = (config.encoder_arch == 'VAE')
    config.RESNET_encoder = (config.encoder_arch == 'RESNET')
    config.BETAVAE_encoder = (config.encoder_arch == 'BETAVAE')
    config.EFFICIENTNET_encoder = (config.encoder_arch == 'EFFICIENTNET')
    config.TINY_CNN_encoder = (config.encoder_arch == 'TINY_CNN')

    if config.encoder_arch in ['VAE', 'BETAVAE']:
        config.encoder_type = 'VAE'
    elif config.encoder_arch in ['CNN', 'RESNET', 'EFFICIENTNET', 'TINY_CNN']:
        config.encoder_type = 'CNN'
    else:
        config.encoder_type = None

    if config.encoder_activation is not None:
        config.encoder_activation = {
            'ReLU': nn.ReLU(inplace=True),
            'LeakyReLU': nn.LeakyReLU(negative_slope=0.3, inplace=True),
            'Tanh': nn.Tanh(),
            'ELU': nn.ELU(inplace=True),
            'SiLU': nn.SiLU(inplace=True)
        }.get(config.encoder_activation)

    if len(config.image_size) == 1:
        config.image_size = (config.image_size, config.image_size)

    if config.camera_size is None:
        config.camera_size = config.image_size
    else:
        if len(config.camera_size) == 1:
            config.camera_size = (config.camera_size, config.camera_size)

    if config.CNN_encoder:
        kernel_sizes = config.kernel_sizes
        strides = config.strides
        paddings = config.paddings
        poolings = config.poolings
        while len(kernel_sizes) < len(config.encoder_hidden_channels):
            kernel_sizes.append(3)
        while len(strides) < len(kernel_sizes):
            strides.append(1)
        while len(paddings) < len(kernel_sizes):
            paddings.append(kernel_sizes[len(paddings)] // 2)
        while len(poolings) < len(kernel_sizes):
            poolings.append(1)

    config.n_samples_per_update = config.batch_size
    if config.RNN_encoder:
        config.n_samples_per_update *= config.step_size

    config.critic_lr = (config.critic_lr or config.lr)
    config.actor_lr = (config.actor_lr or config.lr)
    config.alpha_lr = (config.alpha_lr or config.actor_lr)

    if config.n_bootstrap_step > 1:
        config.gamma = math.pow(config.gamma, config.n_bootstrap_step)

def copy_environment_file(config):
    shutil.copy('common/carla_environment/environment.py', config.log_dir)
    shutil.copy('common/carla_environment/manual_route_planner.py', config.log_dir)
    shutil.copy('common/network.py', config.log_dir)

def main():
    config = get_config()
    initialize(config)

    model = build_model(config)
    if config.mode == 'train' and config.initial_epoch < config.n_epochs:
        train(model, config)
    elif config.mode == 'test':
        test(model, config)
    elif config.mode == 'test_render':
        test_render(model, config)

if __name__ == '__main__':
    main()


import os
import cv2
import numpy as np
from tensorflow.keras.models import load_model

from utils import getPathToLatestModel, getLabelIndices, loadComicImage

pathToModel = getPathToLatestModel()
model = load_model(pathToModel)
labelIndices = getLabelIndices()

helloImage = loadComicImage(os.path.join('comic_images', 'hello.png'))
yesImage = loadComicImage(os.path.join('comic_images', 'yes.png'))
noImage = loadComicImage(os.path.join('comic_images', 'no.png'))
rows, cols, channels = helloImage.shape

alpha = 0.7
xOffset = (480 - cols) // 2
yOffset = (640 - rows) // 2

videoCapture = cv2.VideoCapture(0)
while True:
    _, frame = videoCapture.read()
    frame = cv2.flip(frame, 1)

    # center crop image
    croppedFrame = frame[:, 80:560]

    image = cv2.resize(croppedFrame, (224, 224))
    image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
    image = image.reshape((1, 224, 224, 3))
    image = (image / 127.0) - 1

    prediction = model.predict(image)[0]

    classIndex = np.argmax(prediction)
    classConfidence = prediction[classIndex] * 100
    classLabel = labelIndices.get(classIndex)

    addedImage = None
    if classLabel == 'greet' and classConfidence > 80:
        addedImage = cv2.addWeighted(frame[xOffset:xOffset + rows, yOffset:yOffset + cols, :], alpha, helloImage[0:rows, 0:cols], 1 - alpha, 0)
    elif classLabel == 'thumbs_down' and classConfidence > 80:
        addedImage = cv2.addWeighted(frame[xOffset:xOffset + rows, yOffset:yOffset + cols, :], alpha, noImage[0:rows, 0:cols], 1 - alpha, 0)
    elif classLabel == 'thumbs_up' and classConfidence > 80:
        addedImage = cv2.addWeighted(frame[xOffset:xOffset + rows, yOffset:yOffset + cols, :], alpha, yesImage[0:rows, 0:cols], 1 - alpha, 0)

    if addedImage is not None:
        frame[xOffset:xOffset + rows, yOffset:yOffset + cols] = addedImage

    cv2.putText(frame, f'{classLabel} ({classConfidence:.2f}%)', (10, 50), cv2.FONT_HERSHEY_SIMPLEX, 1, (255, 0, 0), 2)
    cv2.imshow('ResponDeep', frame)

    if cv2.waitKey(1) & 0xff == ord('q'):
        break

videoCapture.release()
cv2.destroyAllWindows()


from .node import Node, NodeFactory
from math import pi, fabs
import numpy as np
from time import time

class MotorFactory(NodeFactory):
    def get_node(self, network, node_id = 0x10, reduction = 12.45):
        motor = Motor(node_id, reduction)
        network[node_id] = motor
        return motor

class Motor(Node):
    CMD_GET_API_VERSION = 0x01
    CMD_ODRIVE_HEARTBEAT = 0x02
    CMD_ODRIVE_ESTOP = 0x03
    CMD_RESET_ODRIVE = 0x04
    CMD_CLEAR_ERRORS = 0x05
    CMD_SAVE_CONFIGURATION = 0x06
    CMD_SET_AXIS_NODE_ID = 0x07
    CMD_SET_AXIS_REQUESTED_STATE = 0x08
    CMD_GET_HARDWARE_STATUS = 0x09
    CMD_GET_MOTOR_STATUS = 0x0A
    CMD_GET_CONTROLLER_MODES = 0x0B
    CMD_SET_CONTROLLER_MODES = 0x0C
    CMD_GET_CONTROLLER_PID = 0x0D
    CMD_SET_CONTROLLER_PID = 0x0E
    CMD_GET_LIMITS = 0x0F
    CMD_SET_LIMITS = 0x10
    CMD_SET_INPUT_POS = 0x11
    CMD_SET_INPUT_VEL = 0x12
    CMD_SET_INPUT_TORQUE = 0x13
    CMD_ERASE_ODRIVE = 0x14
    CMD_UNLOCK = 0x15

    def __init__(self, node_id, reduction):
        super(Motor, self).__init__(node_id)
        self.reduction = reduction
        self._factor = reduction / 2.0 / pi
        self._status = {
            'state': 0,
            'error': 0,
            'motor_err': 0,
            'encoder_err': 0,
            'controller_err': 0,
            'control_mode': 3,
            'input_mode': 1,
            'timestamp': 0
        }

    def associate_network(self, network):
        super(Motor, self).associate_network(network)
        self.network.subscribe(self.can_id(self.CMD_ODRIVE_HEARTBEAT), self.heartbeat_callback)

    def heartbeat_callback(self, can_id, data, timestamp):
        state,err,merr,eerr,cerr,cmode,imode,mtype = list(np.frombuffer(data, dtype='<B'))
        self._status.update({
            'state': state,
            'error': err,
            'motor_err': merr,
            'encoder_err': eerr,
            'controller_err': cerr,
            'control_mode': cmode,
            'input_mode': imode,
            'timestamp': timestamp
        })

    def status(self):
        print(time())
        return self._status

    @Node.get_func_decorator(CMD_GET_API_VERSION, '<u2,<u2,<u4')
    def get_api_ver(self, main_ver, sub_ver, uuid):
        return {
            'device_uuid': uuid,
            'main_version': main_ver,
            'sub_version': sub_ver,
        }

    @Node.send_func_decorator(CMD_ODRIVE_ESTOP)
    def estop(self):
        pass

    @Node.send_func_decorator(CMD_RESET_ODRIVE)
    def reboot(self):
        pass

    @Node.send_func_decorator(CMD_CLEAR_ERRORS)
    def clear_errors(self):
        pass

    @Node.send_func_decorator(CMD_SAVE_CONFIGURATION)
    def save_configuration(self):
        pass

    @Node.send_func_decorator(CMD_SET_AXIS_NODE_ID, '<u2,<u2,<u4')
    def set_can_id(self, id = 0x10, rate = 1000, bitrate = 1000000):
        bitrate_id = 2
        if bitrate == 250000:
            bitrate_id = 0
        elif bitrate == 500000:
            bitrate_id = 1
        return (id, bitrate_id, rate)

    @Node.send_func_decorator(CMD_SET_AXIS_REQUESTED_STATE, '<u4')
    def enable(self):
        return (8, )

    @Node.send_func_decorator(CMD_SET_AXIS_REQUESTED_STATE, '<u4')
    def disable(self):
        return (1, )

    @Node.send_func_decorator(CMD_SET_AXIS_REQUESTED_STATE, '<u4')
    def calibrate(self):
        return (3, )

    @Node.get_func_decorator(CMD_GET_HARDWARE_STATUS, '<f4,<f4')
    def get_vbus(self, vbus, temperature):
        return vbus

    @Node.get_func_decorator(CMD_GET_HARDWARE_STATUS, '<f4,<f4')
    def get_temperature(self, vbus, temperature):
        return temperature

    @Node.get_func_decorator(CMD_GET_MOTOR_STATUS, '<f4,<f2,<f2')
    def get_status(self, pos, vel, torque):
        return (pos / self._factor, vel / self._factor, torque * self.reduction)

    @Node.get_func_decorator(CMD_GET_MOTOR_STATUS, '<f4,<f2,<f2')
    def get_pos(self, pos, vel, torque):
        return pos / self._factor

    @Node.get_func_decorator(CMD_GET_MOTOR_STATUS, '<f4,<f2,<f2')
    def get_vel(self, pos, vel, torque):
        return vel / self._factor

    @Node.get_func_decorator(CMD_GET_MOTOR_STATUS, '<f4,<f2,<f2')
    def get_torque(self, pos, vel, torque):
        return torque * self.reduction

    # @Node.get_func_decorator(CMD_GET_MOTOR_STATUS, '<f4,<f2,<f2')
    # def get_phase_current(self, pos, vel, torque):
    #     k = 240
    #     if self.reduction < 50 and self.reduction > -50:
    #         k = 470
    #     return torque / 8.27 * k

    @Node.get_func_decorator(CMD_GET_CONTROLLER_MODES, '<B,<B')
    def get_controller_modes(self, control_mode, input_mode):
        return (control_mode, input_mode)

    @Node.get_func_decorator(CMD_GET_CONTROLLER_MODES, '<B,<B,<f4,<f2')
    def get_ramp_mode_ramp(self, control_mode, input_mode, ramp,_):
        assert input_mode == 2
        return ramp / fabs(self._factor)

    @Node.get_func_decorator(CMD_GET_CONTROLLER_MODES, '<B,<B,<f4,<f2')
    def get_filter_mode_bandwidth(self, control_mode, input_mode, bandwidth,_):
        assert input_mode == 3
        return bandwidth

    @Node.get_func_decorator(CMD_GET_CONTROLLER_MODES, '<B,<B,<f2,<f2,<f2')
    def get_traj_mode_params(self, control_mode, input_mode, max_vel, max_acc, max_dec):
        assert input_mode == 5
        return (fabs(max_vel / self._factor), fabs(max_acc / self._factor), fabs(max_dec / self._factor))

    @Node.send_func_decorator(CMD_SET_CONTROLLER_MODES, '<B,<B')
    def position_mode(self):
        return (3, 1)

    @Node.send_func_decorator(CMD_SET_CONTROLLER_MODES, '<B,<B')
    def velocity_mode(self):
        return (2, 1)

    @Node.send_func_decorator(CMD_SET_CONTROLLER_MODES, '<B,<B')
    def torque_mode(self):
        return (1, 1)

    @Node.send_func_decorator(CMD_SET_CONTROLLER_MODES, '<B,<B,<f4')
    def position_filter_mode(self, bandwidth = 20):
        return (3, 3, bandwidth)

    @Node.send_func_decorator(CMD_SET_CONTROLLER_MODES, '<B,<B,<f2,<f2,<f2')
    def position_traj_mode(self, max_vel, max_acc, max_dec):
        return (3, 5, fabs(max_vel * self._factor), fabs(max_acc * self._factor), fabs(max_dec * self._factor))

    @Node.send_func_decorator(CMD_SET_CONTROLLER_MODES, '<B,<B,<f4')
    def velocity_ramp_mode(self, ramp):
        return (2, 2, ramp * fabs(self._factor))

    @Node.get_func_decorator(CMD_GET_CONTROLLER_PID, '<f4,<f2,<f2')
    def get_controller_pid(self, pp, vp, vi):
        return (pp, vp, vi)

    @Node.send_func_decorator(CMD_SET_CONTROLLER_PID, '<f4,<f2,<f2')
    def set_controller_pid(self, pos_p, vel_p, vel_i):
        return (pos_p, vel_p, vel_i)

    @Node.get_func_decorator(CMD_GET_LIMITS, '<f4,<f4')
    def get_limits(self, vel_limit, current_limit):
        return (vel_limit, current_limit)

    @Node.get_func_decorator(CMD_GET_LIMITS, '<f4,<f4')
    def get_vel_limit(self, vel_limit, current_limit):
        return vel_limit / fabs(self._factor)

    @Node.get_func_decorator(CMD_GET_LIMITS, '<f4,<f4')
    def get_current_limit(self, vel_limit, current_limit):
        return current_limit

    @Node.send_func_decorator(CMD_SET_LIMITS, '<f4,<f4')
    def set_limits(self, vel_limit = 0, current_limit = 0):
        return (vel_limit, current_limit)

    @Node.send_func_decorator(CMD_SET_LIMITS, '<f4,<f4')
    def set_vel_limit(self, vel_limit):
        return (vel_limit * fabs(self._factor), 0)

    @Node.send_func_decorator(CMD_SET_INPUT_POS, '<f4,<f2,<f2')
    def set_pos(self, pos, vel_ff = 0, torque_ff = 0):
        return (pos * self._factor, vel_ff * self._factor, torque_ff / self._factor)

    @Node.send_func_decorator(CMD_SET_INPUT_VEL, '<f4,<f4')
    def set_vel(self, vel, torque_ff=0):
        return (vel * self._factor, torque_ff / self._factor)

    @Node.send_func_decorator(CMD_SET_INPUT_TORQUE, '<f4')
    def set_torque(self, torque):
        return (torque / self._factor, )

    @Node.send_func_decorator(CMD_ERASE_ODRIVE)
    def reset_motor(self):
        pass

    @Node.send_func_decorator(CMD_UNLOCK)
    def unlock(self):
        pass


import os
import sys
import time
import torch
import argparse
import numpy as np
import torch.nn as nn
import torch.optim as optim
from datetime import datetime
from torch.nn import DataParallel
from util.logging import init_log
from torchvision import transforms
from torch.optim import lr_scheduler
from dataset.facebank import FaceBank
from util.visualize import Visualizer
from torch.utils.data import DataLoader
from backbone.mobilefacenet import MobileFaceNet
from backbone.resnet import ResNet_50, ResNet_101, ResNet_152
from backbone.irse import IR_50, IR_SE_50, IR_101, IR_SE_101, IR_152, IR_SE_152
from backbone.vgg import VGG11, VGG11_BN, VGG13, VGG13_BN, VGG16, VGG16_BN, VGG19, VGG19_BN

class FaceRecognizer(nn.Module):

    def __init__(self, feature, classnum):
        super(FaceRecognizer, self).__init__()
        self.feature = feature
        self.classifier = nn.Sequential(*[
            nn.Linear(512, 256),
            nn.Tanh(),
            nn.Linear(256, classnum)
        ])

    def load_feature(self, file_path):
        self.feature.load_state_dict(torch.load(file_path)['net_state_dict'])

    def train(self, mode=True):
        self.feature.eval()
        self.training = mode
        for module in self.classifier.children():
            module.train(mode)
        return self

    def eval(self):
        return self.train(False)

    def forward(self, x):
        out = self.feature(x)
        logit = self.classifier(out)
        return logit

def main(args):
    seed = 117
    np.random.seed(seed)
    torch.manual_seed(seed)
    torch.cuda.manual_seed(seed)
    # gpu initialize
    multi_gpus = False
    if len(args.gpus.split(',')) > 1:
        multi_gpus = True
    os.environ['CUDA_VISIBLE_DEVICES'] = args.gpus
    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')

    # log initialize
    save_dir = os.path.join(args.save_dir, args.model_pre +
                            args.backbone.upper() + '_' +
                            args.margin_type.upper() + '_' +
                            datetime.now().strftime('%Y%m%d_%H%M%S'))
    if os.path.exists(save_dir):
        raise NameError('Model directionary exists!')
    os.makedirs(save_dir)
    logging = init_log(save_dir)
    _print = logging.info

    # dataset loader
    # prepare dataset
    with open(args.train_file_list) as f:
        img_label_list = f.read().splitlines()
    np.random.shuffle(img_label_list)
    train_len = len(img_label_list) - len(img_label_list) // 10
    trans = transforms.Compose([
        transforms.ToTensor(),  # range [0, 255.0] => [0.0, 1.0]
        transforms.Normalize(mean=(0.5, 0.5, 0.5), std=(
            0.5, 0.5, 0.5))  # range [0.0, 1.0] => [-1.0, 1.0]
    ])
    # training dataset
    train_dataset = FaceBank(
        args.train_root, img_label_list[:train_len], transform=trans)
    train_loader = DataLoader(train_dataset, batch_size=args.batch_size,
                              shuffle=True, num_workers=8, drop_last=False)
    # testing dataset
    test_dataset = FaceBank(args.train_root, img_label_list[train_len:], transform=trans)
    test_loader = DataLoader(test_dataset, batch_size=128,
                             shuffle=False, num_workers=4, drop_last=False)

    # define backbone
    if args.backbone.upper() == 'RESNET_50':
        feature = ResNet_50([112, 112])
    elif args.backbone.upper() == 'RESNET101':
        feature = ResNet_101([112, 112])
    elif args.backbone.upper() == 'RESNET152':
        feature = ResNet_152([112, 112])
    elif args.backbone.upper() == 'IR_50':
        feature = IR_50([112, 112])
    elif args.backbone.upper() == 'IR_SE_50':
        feature = IR_SE_50([112, 112])
    elif args.backbone.upper() == 'IR_101':
        feature = IR_101([112, 112])
    elif args.backbone.upper() == 'IR_SE_101':
        feature = IR_SE_101([112, 112])
    elif args.backbone.upper() == 'IR_152':
        feature = IR_152([112, 112])
    elif args.backbone.upper() == 'IR_SE_152':
        feature = IR_SE_152([112, 112])
    elif args.backbone.upper() == 'VGG11':
        feature = VGG11()
    elif args.backbone.upper() == 'VGG11_BN':
        feature = VGG11_BN()
    elif args.backbone.upper() == 'VGG13':
        feature = VGG13()
    elif args.backbone.upper() == 'VGG13_BN':
        feature = VGG13_BN()
    elif args.backbone.upper() == 'VGG16':
        feature = VGG16()
    elif args.backbone.upper() == 'VGG16_BN':
        feature = VGG16_BN()
    elif args.backbone.upper() == 'VGG19':
        feature = VGG19()
    elif args.backbone.upper() == 'VGG19_BN':
        feature = VGG19_BN()
    elif args.backbone.upper() == 'MOBILEFACENET':
        feature = MobileFaceNet(feature_dim=512)
    else:
        raise NameError(args.backbone, ' is not availabe!')

    # define classifer
    recognizer = FaceRecognizer(feature, train_dataset.class_nums)
    recognizer.load_feature(args.backbone_path)

    if args.resume:
        print('Resuming the model parameters from: ', args.net_path)
        recognizer.load_state_dict(torch.load(args.net_path)['recognizer_state_dict'])

    # define optimizers for different layer
    criterion = torch.nn.CrossEntropyLoss().to(device)
    optimizer_ft = optim.SGD(recognizer.classifier.parameters(), lr=0.1, momentum=0.9, weight_decay=5e-4, nesterov=True)
    milestone_ratios = [0.33, 0.60, 0.88]
    exp_lr_scheduler = lr_scheduler.MultiStepLR(
        optimizer_ft, milestones=[int(r * args.total_epoch) for r in milestone_ratios], gamma=0.1)

    if multi_gpus:
        recognizer = DataParallel(recognizer).to(device)
    else:
        recognizer = recognizer.to(device)

    best_acc, best_loss, best_iters, total_iters = 0., 0., 0, 0
    vis = Visualizer(env=args.model_pre + args.backbone + '_' + args.margin_type.upper())
    for epoch in range(1, args.total_epoch + 1):

        # train model
        _print('Train Epoch: {}/{} ...'.format(epoch, args.total_epoch))
        recognizer.train()

        since = time.time()
        for data in train_loader:
            img, label = data[0].to(device), data[1].to(device)
            optimizer_ft.zero_grad()

            logits = recognizer(img)
            loss = criterion(logits, label)
            loss.backward()
            optimizer_ft.step()

            total_iters += 1
            if total_iters % 10 == 0:
                # current training accuracy
                _, predict = torch.max(logits.data, 1)
                total = label.size(0)
                correct = (np.array(predict.cpu()) ==
                           np.array(label.data.cpu())).sum()
                time_cur = (time.time() - since) / 100
                since = time.time()
                vis.plot_curves({'softmax loss': loss.item()}, iters=total_iters, title='train loss',
                                xlabel='iters', ylabel='train loss')
                vis.plot_curves({'train accuracy': correct / total}, iters=total_iters, title='train accuracy', xlabel='iters',
                                ylabel='train accuracy')

                _print("Iters: {:0>6d}/[{:0>2d}], loss: {:.4f}, train_accuracy: {:.4f}, time: {:.2f} s/iter, learning rate: {}".format(
                    total_iters, epoch, loss.item(), correct/total, time_cur, optimizer_ft.param_groups[0]['lr']))

            # save model
            if total_iters % args.save_freq == 0:
                msg = 'Saving checkpoint: {}'.format(total_iters)
                _print(msg)
                if multi_gpus:
                    recognizer_state_dict = recognizer.module.state_dict()
                else:
                    recognizer_state_dict = recognizer.state_dict()
                if not os.path.exists(save_dir):
                    os.mkdir(save_dir)
                torch.save({
                    'iters': total_iters,
                    'recognizer_state_dict': recognizer_state_dict
                }, os.path.join(save_dir, 'Iter_%06d_recognizer.ckpt' % total_iters))

            # test accuracy
            if total_iters % args.test_freq == 0:

                # test model on testing dataset
                recognizer.eval()
                test_acc, test_loss = 0., 0.
                with torch.no_grad():
                    for data in test_loader:
                        img, label = data[0].to(device), data[1].to(device)
                        logits = recognizer(img)
                        loss = criterion(logits, label)
                        test_loss += loss.detach().cpu().numpy() * img.shape[0]
                        _, predict = torch.max(logits.data, 1)
                        test_acc += (np.array(predict.cpu()) == np.array(label.data.cpu())).sum()
                    test_acc /= len(test_dataset)
                    test_loss /= len(test_dataset)
                _print('Testing Ave Accuracy: {:.4f}, Test Ave Loss: {:.4f}'.format(test_acc * 100, test_loss))
                if best_acc <= test_acc * 100:
                    best_acc = test_acc * 100
                    best_loss = test_loss
                    best_iters = total_iters

                _print('Current Best Accuracy: {:.4f}, Current Best Loss: {:.4f} in iters: {}'.format(
                    best_acc, best_loss, best_iters))

                vis.plot_curves({'testing': best_acc}, iters=total_iters, title='test accuracy', xlabel='iters', ylabel='test accuracy')
                recognizer.train()
        exp_lr_scheduler.step()
    _print('Finally Best Accuracy: {:.4f}, Best Loss: {:.4f} in iters: {},'.format(best_acc, best_loss, best_iters))
    vis.save()
    print('finishing training')

def parse_args(argv):
    parser = argparse.ArgumentParser(
        description='PyTorch for deep face recognition')
    parser.add_argument('--train_root', type=str,
                        default=r'Dataset\Human_Face_Dataset\facebank-112x112', help='train image root')
    parser.add_argument('--train_file_list', type=str,
                        default=r'Dataset\Human_Face_Dataset\facebank-112x112.list', help='train list')

    parser.add_argument('--backbone', type=str, default='IR_SE_50',
                        help='IR_SE_50, IR_SE_101, IR_SE_152...')
    parser.add_argument('--margin_type', type=str, default='ArcFace',
                        help='The margin type')
    parser.add_argument('--batch_size', type=int,
                        default=200, help='batch size')
    parser.add_argument('--total_epoch', type=int,
                        default=18, help='total epochs')

    parser.add_argument('--save_freq', type=int,
                        default=300, help='save frequency')
    parser.add_argument('--test_freq', type=int,
                        default=300, help='test frequency')
    parser.add_argument('--backbone_path', type=str,
                        default='', help='backbone model')
    parser.add_argument('--resume', type=int,
                        default=False, help='resume model')
    parser.add_argument('--net_path', type=str,
                        default='', help='resume model')
    parser.add_argument('--save_dir', type=str,
                        default='finetune_model', help='model save dir')
    parser.add_argument('--model_pre', type=str,
                        default='FRS_', help='model prefix')
    parser.add_argument('--gpus', type=str,
                        default='0', help='model prefix')
    return parser.parse_args(argv)

if __name__ == "__main__":
    main(parse_args(sys.argv[1:]))


# ----------------------------------------------
# Face search using VGG16
# ----------------------------------------------

import caffe
from datetime import datetime
import numpy as np
import sys, getopt
import cv2
import os
import glob
import pickle

from keras.applications.vgg16 import VGG16
from keras.preprocessing import image
from keras.models import Model

os.environ['KERAS_BACKEND'] = 'tensorflow'

from keras.models import load_model
from keras.preprocessing import image

def interpret_output(output, img_width, img_height):
	classes = ["face"]
	w_img = img_width
	h_img = img_height
	threshold = 0.2
	iou_threshold = 0.5
	num_class = 1
	num_box = 2
	grid_size = 11
	probs = np.zeros((grid_size,grid_size,2,20))
	class_probs = np.reshape(output[0:grid_size*grid_size*num_class],(grid_size,grid_size,num_class))
	scales = np.reshape(output[grid_size*grid_size*num_class:grid_size*grid_size*num_class+grid_size*grid_size*2],(grid_size,grid_size,2))
	boxes = np.reshape(output[grid_size*grid_size*num_class+grid_size*grid_size*2:],(grid_size,grid_size,2,4))
	offset = np.transpose(np.reshape(np.array([np.arange(grid_size)]*grid_size*2),(2,grid_size,grid_size)),(1,2,0))

	boxes[:,:,:,0] += offset
	boxes[:,:,:,1] += np.transpose(offset,(1,0,2))
	boxes[:,:,:,0:2] = boxes[:,:,:,0:2] / grid_size
	boxes[:,:,:,2] = np.multiply(boxes[:,:,:,2],boxes[:,:,:,2])
	boxes[:,:,:,3] = np.multiply(boxes[:,:,:,3],boxes[:,:,:,3])

	boxes[:,:,:,0] *= w_img
	boxes[:,:,:,1] *= h_img
	boxes[:,:,:,2] *= w_img
	boxes[:,:,:,3] *= h_img

	for i in range(2):
		for j in range(num_class):
			probs[:,:,i,j] = np.multiply(class_probs[:,:,j],scales[:,:,i])
	filter_mat_probs = np.array(probs>=threshold,dtype='bool')
	filter_mat_boxes = np.nonzero(filter_mat_probs)
	boxes_filtered = boxes[filter_mat_boxes[0],filter_mat_boxes[1],filter_mat_boxes[2]]
	probs_filtered = probs[filter_mat_probs]
	classes_num_filtered = np.argmax(probs,axis=3)[filter_mat_boxes[0],filter_mat_boxes[1],filter_mat_boxes[2]]

	argsort = np.array(np.argsort(probs_filtered))[::-1]
	boxes_filtered = boxes_filtered[argsort]
	probs_filtered = probs_filtered[argsort]
	classes_num_filtered = classes_num_filtered[argsort]

	for i in range(len(boxes_filtered)):
		if probs_filtered[i] == 0 : continue
		for j in range(i+1,len(boxes_filtered)):
			if iou(boxes_filtered[i],boxes_filtered[j]) > iou_threshold :
				probs_filtered[j] = 0.0

	filter_iou = np.array(probs_filtered>0.0,dtype='bool')
	boxes_filtered = boxes_filtered[filter_iou]
	probs_filtered = probs_filtered[filter_iou]
	classes_num_filtered = classes_num_filtered[filter_iou]

	result = []
	for i in range(len(boxes_filtered)):
		result.append([classes[classes_num_filtered[i]],boxes_filtered[i][0],boxes_filtered[i][1],boxes_filtered[i][2],boxes_filtered[i][3],probs_filtered[i]])

	return result

def iou(box1,box2):
	tb = min(box1[0]+0.5*box1[2],box2[0]+0.5*box2[2])-max(box1[0]-0.5*box1[2],box2[0]-0.5*box2[2])
	lr = min(box1[1]+0.5*box1[3],box2[1]+0.5*box2[3])-max(box1[1]-0.5*box1[3],box2[1]-0.5*box2[3])
	if tb < 0 or lr < 0 : intersection = 0
	else : intersection =  tb*lr
	return intersection / (box1[2]*box1[3] + box2[2]*box2[3] - intersection)

def show_results(MODE,img,results, img_width, img_height, keras_model, features, img_paths):
	img_cp = img.copy()
	for i in range(len(results)):
		x = int(results[i][1])
		y = int(results[i][2])
		w = int(results[i][3])//2
		h = int(results[i][4])//2

		if(w<h):
			w=h
		else:
			h=w

		xmin = x-w
		xmax = x+w
		ymin = y-h
		ymax = y+h
		if xmin<0:
			xmin = 0
		if ymin<0:
			ymin = 0
		if xmax>img_width:
			xmax = img_width
		if ymax>img_height:
			ymax = img_height

		cv2.rectangle(img_cp,(xmin,ymin),(xmax,ymax),(0,255,0),2)
		cv2.rectangle(img_cp,(xmin,ymin-20),(xmax,ymin),(125,125,125),-1)
		cv2.putText(img_cp,results[i][0] + ' : %.2f' % results[i][5],(xmin+5,ymin-7),cv2.FONT_HERSHEY_SIMPLEX,0.5,(0,0,0),1)

		target_image=img_cp
		margin=w/4

		x=xmin
		y=ymin
		w*=2
		h*=2

		x2=x-margin
		y2=y-margin
		w2=w+margin*2
		h2=h+margin*2

		if(x2<0):
			x2=0
		if(y2<0):
			y2=0
		if(x2+w2>=img.shape[1]):
			w2=img.shape[1]-1-x2
		if(y2+h2>=img.shape[0]):
			h2=img.shape[0]-1-y2

		face_image = img[y2:y2+h2, x2:x2+w2]

		if(face_image.shape[0]<=0 or face_image.shape[1]<=0):
			continue

		IMAGE_SIZE_KERAS=224

		img_keras = cv2.resize(face_image, (IMAGE_SIZE_KERAS,IMAGE_SIZE_KERAS))
		img_keras = img_keras[...,::-1]  #BGR 2 RGB
		img_keras = np.array(img_keras, dtype=np.float32)
		img_keras.shape = (1,) + img_keras.shape
		img_keras -= 128
		feature = keras_model.predict(img_keras)[0]
		feature = feature / np.linalg.norm(feature)  # Normalize

		cv2.rectangle(target_image, (x2,y2), (x2+w2,y2+h2), color=(0,0,255), thickness=3)
		offset=16

		query = feature
		dists = np.linalg.norm(features - query, axis=1)  # Do search
		ids = np.argsort(dists)[:30] # Top 30 results
		scores = [(dists[id], img_paths[id]) for id in ids]

		for j in range(5):
			prob=scores[j][0]
			file=scores[j][1]

			IMAGE_SIZE_THUMB=128

			img = cv2.imread(file)
			img = cv2.resize(img, (IMAGE_SIZE_THUMB,IMAGE_SIZE_THUMB))

			x=j*IMAGE_SIZE_THUMB
			y=i*(IMAGE_SIZE_THUMB+16)
			h, w, _ = img.shape
			target_image[y:y+h, x:x+w] = img

			cv2.putText(target_image, ""+str(prob), (x,y+IMAGE_SIZE_THUMB+16), cv2.FONT_HERSHEY_COMPLEX_SMALL, 0.8, (0,0,250));

		if(MODE == "capture"):
			for k in range(1000):
				if(os.path.exists("runtime/faces/"+str(k)+".jpg")):
					continue
				img_keras = cv2.resize(face_image, (IMAGE_SIZE_KERAS,IMAGE_SIZE_KERAS))
				cv2.imwrite("runtime/faces/"+str(k)+".jpg", img_keras)
				feature_path = 'runtime/feature/'+str(k)+'.pkl'
				pickle.dump(feature, open(feature_path, 'wb'))
				features.append(pickle.load(open(feature_path, 'rb')))
				img_paths.append('runtime/' + os.path.splitext(os.path.basename(feature_path))[0] + '.jpg')
				return True

	cv2.imshow('face search',img_cp)
	return False

def get_mean(binary_proto,width,height):
	mean_filename=binary_proto
	proto_data = open(mean_filename, "rb").read()
	a = caffe.io.caffe_pb2.BlobProto.FromString(proto_data)
	mean  = caffe.io.blobproto_to_array(a)[0]
	print "mean value of "+binary_proto+" is "+str(mean)+" shape "+str(mean.shape)

	shape=(mean.shape[0],height,width);
	mean=mean.copy()
	mean.resize(shape)

	print "resized mean value is "+str(mean)

	return mean

def main(argv):
	DEMO_IMG=""
	MODE="search"

	# Parse command
	if len(sys.argv) >= 2:
		MODE = sys.argv[1]
		if len(sys.argv) >= 3:
			DEMO_IMG = sys.argv[2]
	else:
		print("usage: python face_search.py [capture/search]")
		sys.exit(1)

	if MODE!="search" and MODE!="capture":
		print("command not found")
		sys.exit(1)

	# Create folder
	if MODE=="capture":
		if not os.path.exists("local/faces"):
			os.mkdir("local/faces")
		if not os.path.exists("local/feature"):
			os.mkdir("local/feature")
		if not os.path.exists("runtime/faces"):
			os.mkdir("runtime/faces")
		if not os.path.exists("runtime/feature"):
			os.mkdir("runtime/feature")

	# Read image features
	features = []
	img_paths = []
	for feature_path in glob.glob("local/feature/*"):
	    features.append(pickle.load(open(feature_path, 'rb')))
	    img_paths.append('local/faces/' + os.path.splitext(os.path.basename(feature_path))[0] + '.jpg')
	for feature_path in glob.glob("runtime/feature/*"):
	    features.append(pickle.load(open(feature_path, 'rb')))
	    img_paths.append('runtime/faces/' + os.path.splitext(os.path.basename(feature_path))[0] + '.jpg')

	# Read yolo
	model_filename = './pretrain/face.prototxt'
	weight_filename = './pretrain/face.caffemodel'
	net = caffe.Net(model_filename, weight_filename, caffe.TEST)

	# Read models
	base_model = VGG16(weights='imagenet')
	keras_model = Model(inputs=base_model.input, outputs=base_model.get_layer('fc1').output)
	keras_model.summary()

	while True:
		cap = cv2.VideoCapture(0)
		ret, frame = cap.read() #BGR
		img=frame
		img = img[...,::-1]  #BGR 2 RGB
		inputs = img.copy() / 255.0

		if(DEMO_IMG!=""):
			img = cv2.imread(DEMO_IMG)
			img = img[...,::-1]  #BGR 2 RGB
			inputs = img.copy() / 255.0

		transformer = caffe.io.Transformer({'data': net.blobs['data'].data.shape})
		transformer.set_transpose('data', (2,0,1))
		out = net.forward_all(data=np.asarray([transformer.preprocess('data', inputs)]))
		img_cv = cv2.cvtColor(img, cv2.COLOR_RGB2BGR)
		results = interpret_output(out['layer20-fc'][0], img.shape[1], img.shape[0])
		if(show_results(MODE,img_cv,results, img.shape[1], img.shape[0], keras_model, features, img_paths)):
			if(MODE == "capture"):
				break

		k = cv2.waitKey(1)
		if k == 27:
			break

	cap.release()
	cv2.destroyAllWindows()

if __name__=='__main__':
	main(sys.argv[1:])