from datetime import datetime

import numpy as np
import matplotlib.pyplot as plt
np.seterr(over='ignore')

from .Continuum import continuum
from .Ion import ion
import ChiantiPy.tools.data as chdata
import ChiantiPy.tools.constants as const
import ChiantiPy.tools.util as util
import ChiantiPy.Gui as chGui
from ChiantiPy.base import specTrails
from ChiantiPy.base import ionTrails

class radLoss(ionTrails, specTrails):
    '''
    Calculate the emission spectrum as a function of temperature and density.

    includes elemental abundances or ionization equilibria

    temperature and density can be arrays but, unless the size of either is one (1),
    the two must have the same size

    the returned spectrum will be convolved with a filter of the specified width on the
    specified wavelength array

    the default filter is gaussianR with a resolving power of 1000.  Other filters,
    such as gaussian, box and lorentz, are available in ChiantiPy.filters.  When using the box filter,
    the width should equal the wavelength interval to keep the units of the continuum and line
    spectrum the same.

    A selection of ions can be make with ionList containing the names of
    the desired lines in Chianti notation, i.e. C VI = c_6

    a minimum abundance can be specified so that the calculation can be speeded up by excluding
    elements with a low abundance. With solar photospheric abundances -

    setting minAbund = 1.e-4 will include H, He, C, O, Ne
    setting minAbund = 2.e-5 adds  N, Mg, Si, S, Fe
    setting minAbund = 1.e-6 adds  Na, Al, Ar, Ca, Ni

    Setting em will multiply the spectrum at each temperature by the value of em.

    em [for emission measure], can be a float or an array of the same length as the
    temperature/density.

    abundance: to select a particular set of abundances, set abundance to the name of a CHIANTI abundance file,
        without the '.abund' suffix, e.g. 'sun_photospheric_1998_grevesse'
        If set to a blank (''), a gui selection menu will popup and allow the selection of an set of abundances
    '''
    def __init__(self, temperature, eDensity, elementList=None, ionList = None, minAbund=None, doContinuum=1, abundance=None, verbose=0, allLines=1, keepIons=0):
        t1 = datetime.now()
        masterlist = chdata.MasterList
        # use the ionList but make sure the ions are in the database
        if ionList:
            alist=[]
            for one in ionList:
                if masterlist.count(one):
                    alist.append(one)
                else:
                    if verbose:
                        pstring = ' %s not in CHIANTI database'%(one)
                        print(pstring)
            masterlist = alist
        self.Defaults=chdata.Defaults
        self.argCheck(temperature=temperature, eDensity=eDensity, pDensity=None, em=None)

        #
        if abundance is not None:
            if type(abundance) == str:
                if abundance in chdata.AbundanceList:
                    self.AbundanceName = abundance
                else:
                    abundChoices = chdata.AbundanceList
                    abundChoice = chGui.gui.selectorDialog(abundChoices,label='Select Abundance name',multiChoice=False)
                    abundChoice_idx = abundChoice.selectedIndex
                    self.AbundanceName = abundChoices[abundChoice_idx[0]]
                    if verbose:
                        print(' %s abundances chosen'%(self.AbundanceName))
            else:
                print(' keyword abundance must be a string, either a blank (\'\') or the name of an abundance file')
                return
        else:
            self.AbundanceName = self.Defaults['abundfile']
        if hasattr(self,'AbundanceName'):
            self.Abundance = chdata.Abundance[self.AbundanceName]['abundance']
        #
#        abundAll = chdata.Abundance[self.AbundanceName]['abundance']
#        # needed by ionGate
        self.AbundAll = self.Abundance
        #
        nonzed = self.Abundance > 0.
        minAbundAll = self.Abundance[nonzed].min()
        # if minAbund is even set
        if minAbund:
            if minAbund < minAbundAll:
                minAbund = minAbundAll
        self.MinAbund = minAbund
#        ionInfo = util.masterListInfo()
        #
        nTempDens = self.NTempDens
        freeFreeLoss = np.zeros_like(self.Temperature)
        freeBoundLoss = np.zeros_like(self.Temperature)
        twoPhotonLoss = np.zeros_like(self.Temperature)
        boundBoundLoss = np.zeros_like(self.Temperature)
        twoPhotonLoss = np.zeros_like(self.Temperature)
        #
        self.IonsCalculated = []
        if keepIons:
            self.IonInstances = {}
        self.Finished = []
        #
        self.WvlRange = [0., 1.e+30]
        #
        self.ionGate(elementList = elementList, ionList = ionList, minAbund=minAbund, doContinuum=doContinuum, doWvlTest=0, verbose=verbose)
        #
        #
        for akey in sorted(self.Todo.keys()):
            zStuff = util.convertName(akey)
            Z = zStuff['Z']
            ionstage = zStuff['Ion']
            dielectronic = zStuff['Dielectronic']
            abundance = self.Abundance[Z - 1]
            if verbose:
                print(' %5i %5s abundance = %10.2e '%(Z, const.El[Z-1],  abundance))
            if verbose:
                print(' doing ion %s for the following processes %s'%(akey, self.Todo[akey]))
            if ionstage != 1:
                if verbose:
                    print(' calculating ff continuum for :  %s'%(akey))
                if 'ff' in self.Todo[akey]:
                    # need to skip the neutral
                    cont = continuum(akey, temperature, abundance=abundance)
                    cont.freeFreeLoss()
                    freeFreeLoss += cont.FreeFreeLoss['rate']
                if 'fb' in self.Todo[akey]:
                    if verbose:
                        print(' calculating fb continuum for :  %s'%(akey))
                    cont = continuum(akey, temperature, abundance=abundance)
                    cont.freeBoundLoss(verbose=verbose)
                    if 'errorMessage' not in cont.FreeBoundLoss.keys():
                        freeBoundLoss += cont.FreeBoundLoss['rate']
            if 'line' in self.Todo[akey]:
                if verbose:
                    print(' calculating spectrum for  :  %s'%(akey))
                thisIon = ion(akey, temperature, eDensity, abundance=abundance)
                thisIon.intensity(allLines=allLines)
                self.IonsCalculated.append(akey)
                if 'errorMessage' not in  thisIon.Intensity.keys():
                    self.Finished.append(akey)
                    thisIon.boundBoundLoss()
                    boundBoundLoss += thisIon.BoundBoundLoss['rate']
                else:
                    if verbose:
                        print(thisIon.Intensity['errorMessage'])
                # get 2 photon emission for H and He sequences
                if (Z - ionstage) in [0, 1] and not dielectronic:
                    thisIon.twoPhotonLoss()
                    twoPhotonLoss += thisIon.TwoPhotonLoss['rate']
        self.FreeFreeLoss = freeFreeLoss
        self.FreeBoundLoss = freeBoundLoss
        self.BoundBoundLoss = boundBoundLoss
        self.TwoPhotonLoss = twoPhotonLoss
        #
        total = freeFreeLoss + freeBoundLoss + boundBoundLoss + twoPhotonLoss
        t2 = datetime.now()
        dt=t2-t1
        print(' elapsed seconds = %10.2e'%(dt.seconds))
        xlabel = 'Temperature (K)'
        ylabel = r'erg  s$^{-1}$ cm$^{3}$'
        self.RadLoss = {'rate':total, 'temperature':self.Temperature, 'density':self.EDensity, 'minAbund':minAbund, 'abundance':self.AbundanceName, 'ylabel':ylabel, 'xlabel':xlabel}
    #
    # -------------------------------------------------------------------
    #
    def radLossPlot(self, title=0):
        '''
        to plot the radiative losses vs temperature
        '''
        fontsize = 16
        temp = self.RadLoss['temperature']
        rate = self.RadLoss['rate']
        plt.loglog(temp, rate)
#        plt.ylabel(r'erg  s$^{-1}$  ($\int\,$ N$_e\,$N$_H\,$d${\it l}$)$^{-1}$',fontsize=fontsize)
        plt.xlabel(self.RadLoss['xlabel'],fontsize=fontsize)
        plt.ylabel(self.RadLoss['ylabel'],fontsize=fontsize)
        if title:
            title = 'Radiative loss rate,  minAbund = %10.2e'%(self.MinAbund)
            if self.EDensity.size == 1:
                title += ', density = %10.2e'%(self.EDensity)
            plt.title(title, fontsize=fontsize)
        plt.tight_layout()


'''
.. module:: controller
   :synopsis: Control algorithms for lateral control of vehicle dynamics
.. moduleauthor:: Yongkyun Shin <github.com/yongkyuns>

This module is a collection of algorithms for control of lateral vehicle dynamics.
Currently implemented algorithms are PID controller, Model-Predictive controller,
and Deep-Deterministic Policy Gradient (DDPG) controller.
'''

import numpy as np
from scipy import sparse
import cvxpy as cp
import torch
from .learning.ddpg_agent import Agent
from .learning.model import Actor
import sys, os
# import acado

# DDPG Params
STATE_SIZE = 4
ACTION_SIZE = 1
RANDOM_SEED = 1
CHECK_POINT = 'checkpoint_actor.pth'

POS = 0
ANG = 1
NOW = 0
FUTURE = 4

class Controller:
    '''
    Generic controller parent class with some common variables (e.g. sample time, number of states, etc.)
    '''

    def __init__(self, NU=1, NY=1, T=1, dt=0.01):
        self.reference = np.zeros((NY, T))
        self.measured_output = np.zeros((NY, T))
        self.control_input = np.zeros(NU)
        self.dt = dt  # Controller sample time
        self.T = T   # Number of control points along the path

    def control(self):
        pass

class PathController(Controller):
    '''
    PathController has projection horizon in extension to generic Controller class.
    Projection horizon defines the preview, or lookahead distances w.r.t. vehicle
    '''
    def __init__(self, NU=1, NY=1, T=1, dt=0.01, proj_horizon=0.5):
        super().__init__(NU=NU, NY=NY, T=T, dt=dt)
        # [sec] Use T control points along desired path for proj_horizon time
        self.proj_horizon = proj_horizon

    @property
    def proj_step(self):
        return self.proj_horizon / self.T

class DDPG(PathController):
    '''
    Experimental controller which uses reinforcement learning algorithm.
    '''
    def __init__(self, check_point=CHECK_POINT, dt=0.01, T=1, NY=1, proj_horizon=0.5):
        super().__init__(dt=dt, T=T, proj_horizon=proj_horizon)
        device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
        if torch.cuda.is_available():
            trained_model = torch.load(CHECK_POINT)
        else:
            trained_model = torch.load(
                CHECK_POINT, map_location={'cuda:0': 'cpu'})

        self.agent = Agent(state_size=STATE_SIZE,
                           action_size=ACTION_SIZE, random_seed=RANDOM_SEED)
        self.agent.actor_local = Actor(
            STATE_SIZE, ACTION_SIZE, RANDOM_SEED).to(device)
        self.agent.actor_local.load_state_dict(trained_model)

    def get_state(self, yref):
        # current/future pos reference, current/future heading reference
        states = np.array([yref[POS, NOW], yref[POS, FUTURE],
                           yref[ANG, NOW], yref[ANG, FUTURE]])
        states = np.reshape(states, [ACTION_SIZE, STATE_SIZE])
        return states

    def control(self, yref):
        states = self.get_state(yref)
        action = self.agent.act(states, add_noise=False)
        return float(action)

class PID(PathController):
    '''
    PID controller with variable number of control points along desired path.
    Control parameters are lateral position and heading. For each control point & parameter,
    there is a corresponding normalized weight which adds up to 1. Therefore, a single PID
    controller is used to control all of the controlled states and points along desired path.

    ================  ==================================================
    **Arguments:**
    NU                (int) Number of controlled states
    NY                (int) Number of outputs
    T                 (int) Number of reference points along trajectory
    dt                (float) Sample time [sec] of the controller
    P_Gain            (float) Proportional gain of the controller
    I_Gain            (float) Integral gain of the controller
    D_Gain            (float) Derivative gain of the controller
    weight_split_T    (list) Determine how to split weight on reference points (e.g. [0.5, 0.5] for equal split on all reference points)
    weight_split_Y    (list) Determine how to split weight on measured output (e.g. [0.5, 0.5] for equal split on position and heading control)

    **Variables:**
    weight_table      (numpy array) NY-by-T array of normalized weight for each controlled point & parameter
    ================  ==================================================
    '''

    def __init__(self, NU=1,
                       NY=1,
                       T=1,
                       dt=0.01,
                       P_Gain=1,
                       I_Gain=0.1,
                       D_Gain=0.01,
                       weight_split_T=None,
                       weight_split_Y=None,
                       proj_horizon=0.5):
        super().__init__(T=T, dt=dt, proj_horizon=proj_horizon)
        self.P_Gain = P_Gain
        self.I_Gain = I_Gain
        self.D_Gain = D_Gain

        self.weight_table = self.__calc_weight_table(
            T, NY, weight_split_T, weight_split_Y)

        self._integral = np.zeros(T)
        self._error_old = np.zeros(T)

    def __calc_weight_table(self, T, NY, T_split, Y_split):
        '''
        Update normalized weight table for PID control.
        '''
        if T_split == None and Y_split == None:
            table = np.ones((NY, T)) / (NY*T)
        else:
            table = np.zeros((NY, T))
            vec_T = np.linspace(T_split[0], T_split[1], num=T)
            vec_Y = Y_split
            for i in range(NY):
                for j in range(T):
                    table[i, j] = vec_Y[i] * vec_T[j]
            if np.sum(table) > 0:
                table /= np.sum(table)
        return table

    def control(self, error_table):
        '''
        Calculate control input
        '''
        u = np.zeros(self.T)
        for i in range(self.T):
            e = 0
            for j in range(self.weight_table.shape[0]):
                e += self.weight_table[j, i] * error_table[j, i]
            self._integral[i] += e * self.dt
            u[i] = (self.P_Gain * e) + (self.I_Gain * self._integral[i]
                                        ) + (self.D_Gain * (e-self._error_old[i])/self.dt)
            self._error_old[i] = e
        return np.sum(u)

class MPC(PathController):
    '''
    Generic model predictive controller class. Uses linearized state-space model.
    Convex optimization is solved through cvxpy module.

    ================  ==================================================
    **Arguments:**
    getModel          (function) Function which should return linearized state-space dynamics model as numpy array
    NU                (int) Number of controlled states
    NY                (int) Number of outputs
    T                 (int) Number of reference points along trajectory
    dt                (float) Sample time [sec] of the controller
    output_weight     (float) Weight for control of the output
    input_weight      (float) Weight for input constraint control
    input_lim         (float) Control input limit
    input_rate_lim    (float) Control input rate limit
    ================  ==================================================
    '''

    def __init__(self, getModel,
                 T=5,
                 dt=0.01,
                 output_weight=None,
                 input_weight=None,
                 input_lim=None,
                 input_rate_lim=None,
                 proj_horizon=0.5):

        super().__init__(NU=1, NY=1, T=T, dt=dt, proj_horizon=proj_horizon)
        self.getModel = getModel
        self.output_weight = output_weight
        self.input_weight = input_weight
        self.input_lim = input_lim
        self.input_rate_lim = input_rate_lim

        self.setup_problem()

    def setup_problem(self):
        A, B, C, _ = self.getModel(sample_time=self.dt)
        NU = B.shape[1]  # Number of inputs to state-space model
        NX = A.shape[1]  # Number of states
        NY = C.shape[0]  # Number of outputs
        T = self.T  # Prediction horizon
        Q_DEFAULT = 5
        R_DEFAULT = 1

        # Constraints
        if self.input_lim is not None:
            umax = np.array(self.input_lim)
            umin = -1 * umax
        if self.input_rate_lim is not None:
            urmax = np.array([[1], [1]]) * 1
            urmin = -1 * urmax

        # Objective function
        if self.output_weight is None:
            Q = sparse.eye(NY) * Q_DEFAULT
        else:
            Q = sparse.diags(self.output_weight)

        if self.input_weight is None:
            R = sparse.eye(NU) * R_DEFAULT
        else:
            R = sparse.diags(self.input_weight)

        # Define problem
        self.u = cp.Variable((NU, T+1)) # T has one extra buffer element at the end
        x = cp.Variable((NX, T+1))
        y = cp.Variable((NY, T+1))
        self.y_ref = cp.Parameter((NY, T))
        # self.y_ref.value = np.zeros((NY, T))
        self.x_init = cp.Parameter(NX)
        # self.x_init.value = np.zeros((NX,1))
        objective = 0
        constraints = [x[:, 0] == self.x_init]

        for k in range(T):
            objective += cp.quad_form(y[:, k] -
                                         self.y_ref[:, k], Q) + cp.quad_form(self.u[:, k], R)

            constraints += [x[:, k+1] == A*x[:, k] + B*self.u[:, k]]
            constraints += [y[:, k] == C*x[:, k]]

            if self.input_rate_lim is not None:
                constraints += [urmin <= self.u[:, k+1] -
                                self.u[:, k], self.u[:, k+1] - self.u[:, k] <= urmax]
            if self.input_lim is not None:
                constraints += [umin <= self.u[:, k], self.u[:, k] <= umax]

        self.prob = cp.Problem(cp.Minimize(objective), constraints)

    def update_values(self, x_init, y_ref):
        self.y_ref.value = y_ref
        self.x_init.value = x_init

    def control(self,y_desired, init_state):
        self.update_values(init_state, y_desired)
        # self.prob.solve(solver=cp.OSQP, max_iter=100,
                #    verbose=False, warm_start=True)
        self.prob.solve(cp.ECOS)
        return self.u[:, 0].value[0]

    def control_with_acado(self, y_desired, init_state):
        A, B, C, _ = self.getModel(sample_time=self.dt)
        NU = B.shape[1]    # Number of inputs to state-space model
        NX = A.shape[1]    # Number of states
        NYN = NX#C.shape[0]   # Number of outputs
        NY = NX+NU        # Control problem
        T = self.T         # Prediction horizon

        x0=np.zeros((1,NX))
        X=np.zeros((T+1,NX))
        X[0:T,1] = y_desired[0,:]
        X[0:T,3] = y_desired[1,:]
        U=np.zeros((T,NU))
        Y=np.concatenate((X[0:T,:],np.zeros((T,NU))),axis=1)
        # Y = np.zeros((T,NY))
        yN=np.zeros((1,NYN))
        x0[0,:]=init_state  # initial state
        yN[0,:]=Y[-1,:NYN]         # reference terminal state

        # Objective function
        Q_DEFAULT = 5
        R_DEFAULT = 1
        # Q = sparse.eye(NY) * Q_DEFAULT
        # Qf = sparse.eye(NYN) * Q_DEFAULT
        # R = sparse.eye(NU) * R_DEFAULT
        Q = np.diag([Q_DEFAULT,0,Q_DEFAULT,0,R_DEFAULT])
        Qf = np.diag([Q_DEFAULT,0,Q_DEFAULT,0])

        X, U = acado.mpc(0, 1, x0, X,U,Y,yN, np.transpose(np.tile(Q,T)), Qf, 0)

        if U[0] != 0 or U[4] != 0:
            print('NOT ZERO!!!!')
        return 0

    # def control2(self, y_desired, init_state):
    #     A, B, C, _ = self.getModel(sample_time=self.dt)

    #     NU = B.shape[1]  # Number of inputs to state-space model
    #     NX = A.shape[1]  # Number of states
    #     NY = C.shape[0]  # Number of outputs
    #     T = self.T  # Prediction horizon
    #     Q_DEFAULT = 5
    #     R_DEFAULT = 1

    #     # Constraints
    #     if self.input_lim is not None:
    #         umax = np.array(self.input_lim)
    #         umin = -1 * umax
    #     if self.input_rate_lim is not None:
    #         urmax = np.array([[1], [1]]) * 1
    #         urmin = -1 * urmax

    #     # Objective function
    #     if self.output_weight is None:
    #         Q = sparse.eye(NY) * Q_DEFAULT
    #     else:
    #         Q = sparse.diags(self.output_weight)

    #     if self.input_weight is None:
    #         R = sparse.eye(NU) * R_DEFAULT
    #     else:
    #         R = sparse.diags(self.input_weight)

    #     # Define problem
    #     u = cp.Variable((NU, T+1)) # T has one extra buffer element at the end
    #     x = cp.Variable((NX, T+1))
    #     y = cp.Variable((NY, T+1))
    #     yref = cp.Parameter((NY, T))
    #     x_init = cp.Parameter(NX)
    #     objective = 0
    #     constraints = [x[:, 0] == x_init]

    #     yref.value = y_desired
    #     x_init.value = init_state

    #     for k in range(T):
    #         objective += cp.quad_form(y[:, k] -
    #                                      yref[:, k], Q) + cp.quad_form(u[:, k], R)

    #         constraints += [x[:, k+1] == A*x[:, k] + B*u[:, k]]
    #         constraints += [y[:, k] == C*x[:, k]]

    #         if self.input_rate_lim is not None:
    #             constraints += [urmin <= u[:, k+1] -
    #                             u[:, k], u[:, k+1] - u[:, k] <= urmax]
    #         if self.input_lim is not None:
    #             constraints += [umin <= u[:, k], u[:, k] <= umax]

    #     prob = cp.Problem(cp.Minimize(objective), constraints)
    #     prob.solve(solver=cp.OSQP, max_iter=100,
    #                verbose=False, warm_start=True)

    #     return u[:, 0].value[0]

def main():
    pass

if __name__ == '__main__':
    main()


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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

        q_target = rewards + self.gamma * selected_q_next

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

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

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

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


#!/usr/bin/env python
"""
The :mod:`hotspots.calculation` handles the main Fragment Hotspot Maps algorithm. In addition, an alternative pocket burial method, Ghecom, is provided.

The main classes of the :mod:`hotspots.calculation` module are:

- :class:`hotspots.calculation.Buriedness`
- :class:`hotspots.calculation.Runner`

More information about the Fragment Hotspot Maps method is available from:
    - Radoux, C.J. et. al., Identifying the Interactions that Determine Fragment Binding at Protein Hotspots J. Med. Chem. 2016, 59 (9), 4314-4325 [dx.doi.org/10.1021/acs.jmedchem.5b01980]

More information about the Ghecom method is available from:
    - Kawabata T, Go N. Detection of pockets on protein surfaces using small and large probe spheres to find putative ligand binding sites. Proteins 2007; 68: 516-529
"""
from __future__ import print_function, division

import multiprocessing
import operator
import random
import shutil
import sys
import tempfile
import time
from functools import reduce
from os import system, environ
from os.path import join

import numpy as np
import pkg_resources

from scipy import ndimage
from skimage.morphology import ball
from skimage.transform import resize
# from hotspots.protoss import Protoss
from tqdm import tqdm

# Development note. We import the CCDC modules after 3rd part modules where possible as
# we generally find that the ccdc package is less fussy about the underlying compiled
# libraries used.
#

from ccdc.cavity import Cavity
from ccdc.io import MoleculeWriter, MoleculeReader
from ccdc.molecule import Molecule, Coordinates
from ccdc.protein import Protein
from ccdc.utilities import PushDir

from hotspots.atomic_hotspot_calculation import _AtomicHotspot
from hotspots.grid_extension import Grid
from hotspots.hs_utilities import Helper
from hotspots.wrapper_pdb import PDBResult
from hotspots.wrapper_ghecom import Ghecom
from hotspots.result import Results

class ExpBuriedness(object):

    def __init__(self, prot, out_grid, max_probe_radius=11):
        self.prot = prot
        self.out_grid = out_grid
        self.max_probe = max_probe_radius
        self.probe_selem_dict = self._generate_probe_selem()
        self.protein_grid = self._from_molecule(prot)
        if self.out_grid is None:
            self.out_grid = self.protein_grid.copy_and_clear()
        print("buriedness init")

    def _generate_probe_selem(self):
        """
        Generates a set of structuring elements (i.e. probe spheres) starting from a 3x3 cross (0.5 A radius) probe up
        to the max probe radius
        :return:
        """
        probe_selem_dict = {}

        # Small probe

        probe_selem_dict[2] = ball(2)
        # r= probe radius
        for r in range(3,self.max_probe +1):
            probe_selem_dict[r] = ball(r)
        print("generating probes")
        return probe_selem_dict

    def _from_molecule(self, mol, scaling=1):
        """
        generate a molecule mask where gp within the vdw radius of the molecule heavy atoms are set to 1.0
        :param mol: `ccdc.molecule.Molecule`
        :param padding: int
        :param scaling: float
        :return: `hotspots.grid_extension.Grid`
        """

        coords = [a.coordinates for a in mol.atoms]
        g = Grid.initalise_grid(coords=coords, padding= 10, spacing=1)

        for probe in sorted(self.probe_selem_dict.keys(), reverse=True):
            for a in mol.heavy_atoms:
                g.set_sphere(point=a.coordinates,
                             radius=probe * scaling,
                             value=probe,
                             mode='replace',
                             scaling='None')

        for a in mol.heavy_atoms:
            g.set_sphere(point=a.coordinates,
                         radius=a.vdw_radius,
                         value=100,
                         mode='replace',
                         scaling='None')

        try:
            out_bound_box = self.out_grid.bounding_box
            print(out_bound_box)
            origin_indices = g.point_to_indices(out_bound_box[0])
            far_indices = g.point_to_indices(out_bound_box[1])
            region = origin_indices + far_indices
            print(region)
            return g.sub_grid(region)
        except AttributeError:
            return g

    def _close_grid(self, g, probe):

        g_array = g.get_array().astype(int)
        closed_array = ndimage.binary_erosion(g_array, structure=self.probe_selem_dict[probe])
        return Grid.array_to_grid(closed_array.astype(int), g)

    def _multiscale_closing(self,g):

        #probe_sizes.reverse()
        all_g = None
        prot_mask = g > 90

        for probe in sorted(self.probe_selem_dict.keys())[1:]:
            if all_g is None:
                all_g = prot_mask.copy()

            probe_mask = ((g < probe+0.1) & (g>0)) | (prot_mask)

            probe_mask = self._close_grid(probe_mask,probe)
            novel = (-all_g) * (g>0)
            all_g += (probe_mask * novel * (self.max_probe - probe))

        return all_g * (prot_mask < 1)

    def _open_grid(self, g, probe):

        g_array = g.get_array().astype(int)
        opened_array = ndimage.binary_opening(g_array, structure=self.probe_selem_dict[probe])
        return Grid.array_to_grid(opened_array.astype(int), g)

    def buriedness_grid(self):

        closed_g = self._multiscale_closing(self.protein_grid)
        print("close_g")
        out_g = self._open_grid(closed_g,2) * closed_g
        print("out_g")
        out_array = out_g.get_array()
        print("out array")
        scaled_g = Grid.initalise_grid(self.out_grid.bounding_box, padding=0, spacing=0.5)
        print("scaled_g")
        scaled_array = resize(out_array, scaled_g.nsteps ,anti_aliasing=False)
        print(scaled_array)

        # Future tweaking here
        final_array = scaled_array
        print("array to grid")
        return Grid.array_to_grid(final_array.astype(int), scaled_g)

class _WeightedResult(object):
    """
    private class

    A class to handle the weighted grid results

    :param str identifier: identifier, default is the probe identifier assigned at the Atomic Hotspot Calculation stage
    :param `ccdc.utilities.Grid` grid: result grid
    """

    def __init__(self, identifier, grid):
        self.identifier = identifier
        self.grid = grid

class _SampleGrid(object):
    """
    private class

    class to handle sampled grids

    :param name: str, name of probe (donor, acceptor, apolar, positive, negative)
    :param grid: a :class: `ccdc.utilities.Grid` instance
    :param atom_predicate: atom_predicate will be used to select atoms of a molecule for sampling
    """

    def __init__(self, name, grid, atom_predicate):
        self.name = name
        self.grid = grid
        self.atom_predicate = atom_predicate
        self.mol = None

    @staticmethod
    def add_coordinates(coord, trans):
        """
        provides a coordinate list of atoms to be scored be scored in the SampleGrid
        :param coord: tup, (float(x), float(y), float(z)), set of atomic coordinates for "active" coordinates
        :param trans: tup, (float(x), float(y), float(z)), set of translations to translate probes to points
        above threshold
        :return: list of tup
        """

        return [coord[i] + trans[i] for i in range(0, len(trans))]

    def sample(self, coordinate_list, trans):
        """
        score Molecule in grids for which it has active atoms
        :param coordinate_list: list, set of coordinates of translated and rotated probe atoms to be scored
        :param trans:
        :return:
        """

        try:
            return [self.grid._grid.value(self.add_coordinates(c, trans)) for c in coordinate_list]
        except RuntimeError:
            return [0]

    def set_molecule(self, mol, polar_contribution):
        """
        set which atoms match the grid
        probes with polar atoms contribute do not contribute to apolar maps as this leads to artefacts

        :param mol:
        :param polar_contribution:
        :return:
        """
        self.mol = mol
        if self.name == 'apolar' and not polar_contribution and len([a for a in mol.atoms
                                                                     if a.is_donor or a.is_acceptor]) > 0:
            self._active_atoms = []
        elif not hasattr(self, '_active_atoms'):
            self._active_atoms = [a for a in self.mol.atoms if self.atom_predicate(a)]

    @staticmethod
    def is_donor(a):
        """
        returns true if a given atom is a donor

        :param a: a `ccdc.molecule.Atom` instance
        :return: bool, true if the atom classification is "donor"
        """

        if a.is_donor and a.formal_charge == 0:
            return True
        else:
            return False

    @staticmethod
    def is_acceptor(a):
        """
        returns true if a given atom is a acceptor

        :param a: a `ccdc.molecule.Atom` instance
        :return: bool, true if the atom classification is "acceptor"
        """
        if a.is_acceptor and a.formal_charge == 0:
            return True
        else:
            return False

    @staticmethod
    def is_apolar(a):
        """
        returns true if a given atom is a apolar

        :param a: a `ccdc.molecule.Atom` instance
        :return: bool, true if the atom classification is "apolar"
        """
        if a.is_donor or a.is_acceptor or a.formal_charge != 0 or a.atomic_symbol == "Xe":
            return False
        else:
            return True

    @staticmethod
    def is_positive(a):
        """
        returns true if a given atom is a positively charged

        :param a: a `ccdc.molecule.Atom` instance
        :return: bool, true if the atom is positively charged
        """
        return a.formal_charge > 0

    @staticmethod
    def is_negative(a):
        """
        returns true if a given atom is a negatively charged

        :param a: a `ccdc.molecule.Atom` instance
        :return: bool, true if the atom is negatively charged
        """
        return a.formal_charge < 0

    @staticmethod
    def is_aromatic(a):
        """
        returns true if a given atom is aromatic

        :param a: a `ccdc.molecule.Atom` instance
        :return: bool, true if the atom is aromatic
        """
        return a.atomic_symbol == 'C' and a.is_cyclic and any(b.atom_type == 'aromatic' for b in a.bonds)

class Runner(object):
    """
    A class for running the Fragment Hotspot Map calculation
    """

    class Settings(object):
        """
        adjusts the default settings for the calculation

        :param int nrotations: number of rotations (keep it below 10**6)
        :param float apolar_translation_threshold: translate probe to grid points above this threshold. Give lower values for greater sampling. Default 15
        :param float polar_translation_threshold: translate probe to grid points above this threshold. Give lower values for greater sampling. Default 15
        :param bool polar_contributions: allow carbon atoms of probes with polar atoms to contribute to the apolar output map.
        :param bool return_probes: Generate a sorted list of molecule objects, corresponding to probe poses
        :param bool sphere_maps: When setting the probe score on the output maps, set it for a sphere (radius 1.5) instead of a single point.
        """

        def __init__(self, nrotations=3000, apolar_translation_threshold=15, polar_translation_threshold=15,
                     polar_contributions=False, return_probes=False, sphere_maps=False):
            self.nrotations = nrotations
            self.apolar_translation_threshold = apolar_translation_threshold
            self.polar_translation_threshold = polar_translation_threshold
            self.polar_contributions = polar_contributions
            self.return_probes = return_probes
            self.sphere_maps = sphere_maps

        @property
        def _num_gp(self):
            """
            number of grid point for a given volume
            :return:
            """
            return int(float(500) / self.nrotations ** 3)

    class _Sampler(object):
        """
        Samples one or more grids with a probe molecule
        """

        def __init__(self, *grids, **kw):
            """
            Settings used to run fragment-hotspot-maps script

            :param grids: list, list of :class:`ccdc.utilities.Grid` instances
            :param kw: 'settings'
            """

            self.grids = grids
            self.grid_dic = {}
            for g in grids:
                self.grid_dic[g.name] = g
            settings = kw.get('settings')
            if settings is None:
                settings = self.Settings()

            self.settings = settings
            self.probe_grids = [_SampleGrid(g.name, g.grid.copy_and_clear(), g.atom_predicate) for g in self.grids]

        def get_priority_atom(self, molecule):
            """
            Select priority atom. Select polar atom. If multiple polar atoms, select the one furthest from the centre of
            geometry. If no polar atoms, select atom furthest from centre of geometry

            :param molecule: a :class: `ccdc.molecule.Molecule` instance
            :return: a :class: `ccdc.molecule.Molecule` instance, str atom type
            """
            c = molecule.centre_of_geometry()
            polar_atoms = [a for a in molecule.atoms if a.is_donor or a.is_acceptor]
            atom_by_distance = {}
            if len(polar_atoms) > 0:
                for a in polar_atoms:
                    d = Helper.get_distance(c, a.coordinates)
                    atom_by_distance[d] = a
            else:
                for a in molecule.atoms:
                    d = Helper.get_distance(c, a.coordinates)
                    atom_by_distance[d] = a

            greatest_distance = sorted(atom_by_distance.keys())[0]
            priority_atom = atom_by_distance[greatest_distance]

            pa_type = None
            if priority_atom.formal_charge != 0:
                if priority_atom.formal_charge < 0:
                    pa_type = "acceptor"
                elif priority_atom.formal_charge > 0:
                    pa_type = "donor"
            else:
                if priority_atom.is_acceptor:
                    pa_type = "acceptor"
                elif priority_atom.is_donor:
                    pa_type = "donor"
                else:
                    pa_type = "apolar"

            return priority_atom, pa_type

        def get_translation_points(self, priority_atom_type):
            """
            returns a list of coordinates that are greater than the threshold, that the probe will be translated to

            :param priority_atom_type: str, atomic interaction type
            :return: list, list of :class: `ccdc.molecule.Molecule` instances
            """
            translate_probe = []
            wg = self.grid_dic[priority_atom_type]

            if priority_atom_type == 'apolar':
                translation_threshold = self.settings.apolar_translation_threshold
            else:
                translation_threshold = self.settings.polar_translation_threshold
            hs = wg.grid.islands(translation_threshold)
            for g in hs:
                nx, ny, nz = g.nsteps

                maxima = [g.indices_to_point(i, j, k)
                          for i in range(nx) for j in range(ny) for k in range(nz)
                          if g.value(i, j, k) >= translation_threshold]

                translate_probe = translate_probe + maxima
            return translate_probe

        def generate_rand_quaternions(self):
            """
            Returns a list of random quaternions. Length matches settings.nrotations

            :return: tup, (a,b,c,d)
            """
            quaternions = []
            i = 1
            if self.settings.nrotations > 1:
                while i <= self.settings.nrotations:
                    r1 = random.uniform(-1, 1)
                    r2 = random.uniform(-1, 1)
                    s1 = r1 * r1 + r2 * r2
                    if s1 < 1:
                        r3 = random.uniform(-1, 1)
                        r4 = random.uniform(-1, 1)

                        s2 = r3 * r3 + r4 * r4
                        if s2 < 1:
                            q = (r1, r2, r3 * (np.sqrt((1 - s1) / s2)), r4 * (np.sqrt((1 - s1) / s2)))
                            quaternions.append(q)

                            i += 1

            return quaternions

        @staticmethod
        def score(values):
            """
            Calculate geometric mean of scores

            :param values: float, scores of atoms in probe
            :return: float, geometric mean of probe atom scores
            """

            return reduce(operator.__mul__, values, 1.0) ** (1. / len(values))
            # return sum(values)

        def sample_pose(self, trans, active_atoms_dic, probe):
            """
            Return a pose score (as defined by the score(self,dic) function) and a dictionary of atom:scores

            :param trans: list of translations
            :param active_atoms_dic: dict {"interaction type": "atoms to be scored"}
            :param probe: str, interaction_type
            :return:
            """
            if probe == "negative" or probe == "positive":
                atoms = [g.mol.atoms for g in self.grids if g.name == "positive"][0]
                weight = int(6 / len([a for a in atoms if str(a.atomic_symbol) == "C"]))

                apolar_values = [g.sample(active_atoms_dic[g.name], trans) for g in self.grids if
                                 len(active_atoms_dic[g.name]) > 0 and g.name == "apolar"] * weight
                charged_values = [g.sample(active_atoms_dic[g.name], trans) for g in self.grids if
                                  len(active_atoms_dic[g.name]) > 0 and g.name != "apolar"]
                values = apolar_values + charged_values

            else:
                values = [g.sample(active_atoms_dic[g.name], trans) for g in self.grids if
                          len(active_atoms_dic[g.name]) > 0]
            scores = [item for sublist in values for item in sublist]
            return self.score(scores)

        def update_out_grids(self, score, active_coordinates_dic, trans):
            """
            For active atoms for a given grid, set closest grid point value to score, unless already set to a higher
            value

            :param score: float, score of a given probe
            :param active_coordinates_dic:
            :param trans: list of translations
            :return:
            """

            for pg in self.probe_grids:
                actives = active_coordinates_dic[pg.name]
                if len(actives) == 0:
                    continue
                for a in actives:
                    coords = pg.add_coordinates(a, trans)
                    i, j, k = pg.grid.point_to_indices(coords)
                    orig_value = pg.grid.value(i, j, k)
                    #
                    if self.settings.sphere_maps:
                        # pg.grid.set_sphere(coords, 2, score, mode='max')
                        if score > orig_value:
                            pg.grid.set_sphere(coords, 1.5, score, mode='max', scaling='None')
                    else:
                        pg.grid.set_value(i, j, k, max(score, orig_value))

        def get_active_coordinates(self):
            """
            Returns a dictionary of {grid_name:[Coordinates]}

            :return:
            """
            active_coords_dic = {
                g.name: [a.coordinates for a in g._active_atoms]
                for g in self.grids
            }

            return active_coords_dic

        def sample(self, molecule, probe, update_grids = True):
            """
            Sample the grids according to the settings

            :param molecule:
            :param probe: str, interaction type, (donor, acceptor, negative, positive, apolar)
            :return:
            """
            priority_atom, priority_atom_type = self.get_priority_atom(molecule)
            translate_points = self.get_translation_points(priority_atom_type)
            molecule.remove_hydrogens()
            quaternions = self.generate_rand_quaternions()
            high_scoring_probes = {}
            print("\n    nRotations:", len(quaternions), "nTranslations:", len(translate_points), "probename:", probe)

            for g in self.grids:
                g.set_molecule(molecule, True)

            for g in self.probe_grids:
                g.set_molecule(molecule, self.settings.polar_contributions)

            for q in tqdm(quaternions):
                molecule.apply_quaternion(q)
                priority_atom_coordinates = priority_atom.coordinates
                active_coordinates_dic = self.get_active_coordinates()

                for priority_atom_point in translate_points:

                    translation = [priority_atom_point[i] - priority_atom_coordinates[i]
                                   for i in range(0, len(priority_atom_coordinates))]

                    score = self.sample_pose(translation, active_coordinates_dic, probe)

                    try:
                        if score < 1:
                            continue
                    except TypeError:
                        continue
                    if update_grids:
                        self.update_out_grids(score, active_coordinates_dic, translation)

                    if self.settings.return_probes is True:
                        if score > 10:
                            m = molecule.copy()
                            m.translate(translation)
                            m.identifier = "{}".format(score)

                            try:
                                high_scoring_probes[score].append(m)
                            except KeyError:
                                high_scoring_probes[score] = [m]

            if self.settings.return_probes is True:
                sampled_probes = []
                for key in sorted(high_scoring_probes.keys(), reverse=True)[:100]:
                    sampled_probes.extend(high_scoring_probes[key])
                print('Returned probes = ', len(sampled_probes))
                return sampled_probes

                # if len(sampled_probes) > 100:
                #     return sampled_probes[:100]
                # else:
                #     return sampled_probes

    def __init__(self, settings=None):
        self.out_grids = {}
        self.super_grids = {}
        self.buriedness = None
        self.sampled_probes = {}

        if settings is None:
            self.sampler_settings = self.Settings()
        else:
            self.sampler_settings = settings

    @property
    def protein(self):
        """
        protein submitted for calculation
        :return:
        """
        return self._protein

    @protein.setter
    def protein(self, prot):
        """
        protein submitted for calculation
        :param `ccdc.protein.Protein` prot: protein
        :return:
        """
        if isinstance(prot, Protein):
            self._protein = prot
        else:
            raise TypeError("`ccdc.protein.Protein` must be supplied. Hint: Use Protein.from_file()")

    @property
    def charged_probes(self):
        """
        optional settings, True if charged features are to be calculated
        :return:
        """
        return self._charged_probes

    @charged_probes.setter
    def charged_probes(self, option):
        """
        optional settings, True if charged features are to be calculated
        :param bool option: (default = True)
        :return:
        """
        if type(option) is bool:
            self._charged_probes = option
        else:
            raise TypeError("Expecting a bool, got {} instead".format(type(option)))

    @property
    def probe_size(self):
        """
        optional settings, change probe size. NB: feature has not been tested.
        :return:
        """
        return self._probe_size

    @probe_size.setter
    def probe_size(self, size):
        """
        optional settings, change probe size. NB: feature has not been tested.
        :param int size: number of heavy atoms contain within the molecular probes (integer between 3 and 8)
        :return:
        """
        if size in range(3, 8):
            self._probe_size = size
        else:
            raise ValueError("Probe size must be an integer between 3-7")

    @property
    def buriedness_method(self):
        """
        optional settings, pocket detection method. (default = "ghecom")
        :return:
        """
        return self._buriedness_method

    @buriedness_method.setter
    def buriedness_method(self, method):
        """
        optional settings, pocket detection method. (default = "ghecom")
        :param str method: either 'ghecom' or 'ligsite'
        :return:
        """
        method = method.lower()
        if method == 'ghecom':
            if sys.platform == 'linux' or sys.platform == 'linux2':
                if 'GHECOM_EXE' in environ:
                    self._buriedness_method = method
                else:
                    raise EnvironmentError("Must set Ghecom environment variable")
            else:
                raise OSError('Ghecom is only supported on linux')

        elif method == 'ghecom_internal':
            self._buriedness_method = method

        elif method == 'ligsite':
            if sys.platform == 'linux' or sys.platform == 'linux2':
                print("RECOMMENDATION: you have chosen LIGSITE as buriedness method, ghecom is recommended")
            self._buriedness_method = method

        else:
            raise ValueError("Buriedness method must be 'ghecom' (default) or 'ligsite")

    @property
    def cavities(self):
        """
        optional settings, cavities supplied to the calculation
        :return:
        """
        return self._cavities

    @cavities.setter
    def cavities(self, obj):
        """
        optional settings, cavities supplied to the calculation
        :param list or Coordinate or `ccdc.molecule.Molecule` or `ccdc.cavity.Cavity`: cavity information provided
        :return:
        """
        if obj is not None:
            if isinstance(obj, list) or isinstance(obj, tuple):
                if isinstance(obj, Coordinates):
                    try:
                        print(obj.x)
                        self._cavities = [obj]
                    except AttributeError:
                        self._cavities = obj

                    self._cavities = [obj]
                elif isinstance(obj[0], Molecule):
                    self._cavities = [m.centre_of_geometry() for m in obj]
                elif isinstance(obj[0], Cavity):
                    self._cavities = [Helper.cavity_centroid(c) for c in obj]
                else:
                    try:
                        self._cavities = [Coordinates(x=obj[0], y=obj[1], z=obj[2])]
                    except RuntimeError:
                        print("WARNING! Failed to detected cavity, Atomic Hotspot detection to run on whole protein")
                        self._cavities = None

            elif isinstance(obj, Molecule):
                self._cavities = [obj.centre_of_geometry()]
            elif isinstance(obj, Cavity):
                self._cavities = [Helper.cavity_centroid(obj)]

            else:
                print("WARNING! Failed to detected cavity, Atomic Hotspot detection to run on whole protein")
                self._cavities = None
        else:
            self._cavities = None

    @property
    def nprocesses(self):
        """
        number of processes used in parallelisation
        :return:
        """
        return self._nprocesses

    @nprocesses.setter
    def nprocesses(self, num):
        """
        number of processes used in parallelisation
        :param int num: number of processor
        :return:
        """
        num = int(num)
        if num in range(0, int(multiprocessing.cpu_count())):
            self._nprocesses = num
        else:
            raise OSError("CPU count = {}".format(multiprocessing.cpu_count()))

    @property
    def sampler_settings(self):
        """
        sampler settings
        :return:
        """
        return self._sampler_settings

    @sampler_settings.setter
    def sampler_settings(self, settings):
        """
        sampler settings
        :param `hotspots.calculation.Runner.Settings` settings: adjust run settings
        :return:
        """
        if isinstance(settings, self.Settings):
            self._sampler_settings = settings
        else:
            self._sampler_settings = None

    def _get_weighted_maps(self):
        """
        private method

        weight superstar output by burriedness
        :return: a list of :class:`WeightedResult` instances
        """
        results = []
        for s in self.superstar_grids:
            g, b = Grid.common_grid([s.grid, self.buriedness], padding=1)
            weighted_grid = g * b
            results.append(_WeightedResult(s.identifier, weighted_grid))

        return results

    def _dock_probe(self,mol, grid_dict):

        donor_grid = _SampleGrid('donor', grid_dict['donor'], _SampleGrid.is_donor)
        acceptor_grid = _SampleGrid('acceptor', grid_dict['acceptor'], _SampleGrid.is_acceptor)
        apolar_grid = _SampleGrid('apolar', grid_dict['apolar'], _SampleGrid.is_apolar)

        kw = {'settings': self.sampler_settings}
        self.sampler = self._Sampler(apolar_grid, donor_grid, acceptor_grid, **kw)

        probes = self.sampler.sample(mol, probe='probe', update_grids=False)
        return probes

    def _get_out_maps(self, probe, grid_dict, return_probes=False):
        """
        private method

        organises the sampling of weighted superstar maps by molecular probes
        :param str probe: probe identifier set in the Atomic Hotspot calculation
        :param dict grid_dict: dictionary with key = probe identifier and value = `hotspots.grid_extension.Grid`
        :param bool return_probes: optional, bool indicating if probe molecules should be returned
        :return:
        """
        donor_grid = _SampleGrid('donor', grid_dict['donor'], _SampleGrid.is_donor)
        acceptor_grid = _SampleGrid('acceptor', grid_dict['acceptor'], _SampleGrid.is_acceptor)
        apolar_grid = _SampleGrid('apolar', grid_dict['apolar'], _SampleGrid.is_apolar)

        if self.charged_probes:
            negative_grid = _SampleGrid('negative', grid_dict['negative'], _SampleGrid.is_negative)
            positive_grid = _SampleGrid('positive', grid_dict['positive'], _SampleGrid.is_positive)

        kw = {'settings': self.sampler_settings}
        if self.charged_probes:
            self.sampler = self._Sampler(apolar_grid, donor_grid, acceptor_grid, negative_grid, positive_grid, **kw)
        else:
            self.sampler = self._Sampler(apolar_grid, donor_grid, acceptor_grid, **kw)

        probe_path = pkg_resources.resource_filename('hotspots', 'probes/')

        if self.charged_probes:
            if probe == "negative" or probe == "positive":
                mol = MoleculeReader(join(probe_path, "rotate-{}_{}_flat.mol2".format(probe, "test")))[0]
            else:
                mol = MoleculeReader(join(probe_path, "rotate-{}_{}_flat.mol2".format(probe, self.probe_size)))[0]
        else:
            mol = MoleculeReader(join(probe_path, "rotate-{}_{}_flat.mol2".format(probe, self.probe_size)))[0]

        probes = self.sampler.sample(mol, probe=probe)

        for pg in self.sampler.probe_grids:
            if pg.name.lower() == probe:
                try:
                    self.out_grids[pg.name].append(pg.grid)
                except KeyError:
                    self.out_grids[pg.name] = [pg.grid]

        if return_probes is True:
            return probes

    def _calc_hotspots(self, return_probes=False):
        """
        handles the organisation of the hotspot calculation
        :param return_probes: optional, bool indicating if probe molecules should be returned
        :return:
        """
        print("Start atomic hotspot detection\n        Processors: {}".format(self.nprocesses))
        a = _AtomicHotspot()
        a.settings.atomic_probes = {"apolar": "AROMATIC CH CARBON",
                                    "donor": "UNCHARGED NH NITROGEN",
                                    "acceptor": "CARBONYL OXYGEN"}
        if self.charged_probes:
            a.settings.atomic_probes = {"negative": "CARBOXYLATE OXYGEN", "positive": "CHARGED NH NITROGEN"}

        probe_types = a.settings.atomic_probes.keys()
        self.superstar_grids = a.calculate(protein=self.protein,
                                           nthreads=self.nprocesses,
                                           cavity_origins=self.cavities)

        if self.clear_tmp == True:
            shutil.rmtree(a.settings.temp_dir)

        print("Atomic hotspot detection complete\n")

        print("Start buriedness calculation")
        if self.buriedness_method.lower() == 'ghecom' and self.buriedness is None:
            print("    method: Ghecom")
            out_grid = self.superstar_grids[0].buriedness.copy_and_clear()

            ghecom = Ghecom()
            path = ghecom.run(self.protein)
            self.buriedness = ghecom.pdb_to_grid(path, out_grid)

            shutil.rmtree(ghecom.temp) # probs not needed

        elif self.buriedness_method.lower() == 'ghecom_internal' and self.buriedness is None:
            print("    method: Internal version Ghecom")
            out_grid = self.superstar_grids[0].buriedness.copy_and_clear()
            b = ExpBuriedness(prot=self.protein, out_grid=None)
            self.buriedness = b.buriedness_grid()

        elif self.buriedness_method.lower() == 'ligsite' and self.buriedness is None:
            print("    method: LIGSITE")
            self.buriedness = Grid.get_single_grid(grd_dict={s.identifier: s.buriedness for s in self.superstar_grids},
                                                   mask=False)

        self.weighted_grids = self._get_weighted_maps()

        print("Buriedness calcualtion complete\n")

        print("Start sampling")
        grid_dict = {w.identifier: w.grid for w in self.weighted_grids}

        for probe in probe_types:
            if return_probes is True:
                ps = self._get_out_maps(probe, grid_dict, return_probes=True)
                print(len(ps))
                self.sampled_probes.update({probe: ps})

            else:
                self._get_out_maps(probe, grid_dict)

        print("Sampling complete\n")

    def _prepare_protein(self, protoss=False):
        """
        default protein preparation settings on the protein
        :return:
        """
        self.protein.remove_all_waters()
        for lig in self.protein.ligands:
            self.protein.remove_ligand(lig.identifier)
        self.protein.remove_all_metals()

        if protoss is False:
            self.protein.add_hydrogens()

    def from_superstar(self, protein, superstar_grids, buriedness, charged_probes=False, probe_size=7,
                       settings=None, clear_tmp=False):
        """
        calculate hotspot maps from precalculated superstar maps. This enables more effective parallelisation and reuse
        of object such as the Buriedness grids

        :param protein: a :class:`ccdc.protein.Protein` instance
        :param superstar_grids: a :class:`hotspots.atomic_hotspot_calculation._AtomicHotspotResult` instance
        :param buriedness: a :class:`hotspots.grid_extension.Grid` instance
        :param bool charged_probes: If True, include positive and negative probes
        :param int probe_size: Size of probe in number of heavy atoms (3-8 atoms)
        :param settings: `hotspots.calculation.Runner.Settings` settings: holds the sampler settings
        :param bool clear_tmp: If True, clear the temporary directory
        :return:
        """
        start = time.time()
        self.super_grids = {}
        self.superstar_grids = superstar_grids
        self.probe_types = [p.identifier for p in self.superstar_grids]
        self.buriedness = buriedness

        self.protein = protein
        self.charged_probes = charged_probes
        self.probe_size = probe_size
        self.clear_tmp = clear_tmp

        if settings is None:
            self.sampler_settings = self.Settings()
        else:
            self.sampler_settings = settings

        self.weighted_grids = self._get_weighted_maps()

        print("Start sampling")
        grid_dict = {w.identifier: w.grid for w in self.weighted_grids}

        for probe in self.probe_types:
            self._get_out_maps(probe, grid_dict)

        self.super_grids = {p: g[0] for p, g in self.out_grids.items()}

        print("Sampling complete\n")
        print("Runtime = {}seconds".format(time.time() - start))

        return Results(super_grids=self.super_grids,
                       protein=self.protein,
                       buriedness=self.buriedness)

    def from_protein(self, protein, charged_probes=False, probe_size=7, buriedness_method='ghecom',
                     cavities=None, nprocesses=1, settings=None, buriedness_grid=None, clear_tmp=False):
        """
        generates a result from a protein

        :param protein: a :class:`ccdc.protein.Protein` instance
        :param bool charged_probes: If True include positive and negative probes
        :param int probe_size: Size of probe in number of heavy atoms (3-8 atoms)
        :param str buriedness_method: Either 'ghecom' or 'ligsite'
        :param cavities: Coordinate or `ccdc.cavity.Cavity` or `ccdc.molecule.Molecule` or list specifying the cavity or cavities on which the calculation should be run
        :param int nprocesses: number of CPU's used
        :param `hotspots.calculation.Runner.Settings` settings: holds the sampler settings
        :param `ccdc.utilities.Grid` buriedness_grid: pre-calculated buriedness grid
        :return: a :class:`hotspots.result.Results` instance

        >>> from ccdc.protein import Protein
        >>> from hotspots.calculation import Runner

        >>> protein = Protein.from_file(<path_to_protein>)

        >>> runner = Runner()
        >>> settings = Runner.Settings()
        >>> settings.nrotations = 1000  # fewer rotations increase speed at the expense of accuracy
        >>> runner.from_protein(protein, nprocesses=3, settings=settings)
        Result()

        """
        start = time.time()
        self.super_grids = {}
        self.buriedness = buriedness_grid
        self.protein = protein
        self.charged_probes = charged_probes
        self.probe_size = probe_size
        self.buriedness_method = buriedness_method
        self.cavities = cavities
        self.clear_tmp = clear_tmp

        print(self.cavities)
        self.nprocesses = nprocesses
        if settings is None:
            self.sampler_settings = self.Settings()
        else:
            self.sampler_settings = settings
        self._calc_hotspots()  # return probes = False by default
        self.super_grids = {p: g[0] for p, g in self.out_grids.items()}
        print("Runtime = {}seconds".format(time.time() - start))

        return Results(super_grids=self.super_grids,
                       protein=self.protein,
                       buriedness=self.buriedness,
                       superstar={x.identifier: x.grid for x in self.superstar_grids},
                       weighted_superstar={x.identifier: x.grid for x in self.weighted_grids})

    def from_pdb(self, pdb_code, charged_probes=False, probe_size=7, buriedness_method='ghecom', nprocesses=3,
                 cavities=False, settings=None, clear_tmp=False):
        """
        generates a result from a pdb code

        :param str pdb_code: PDB code
        :param bool charged_probes: If True include positive and negative probes
        :param int probe_size: Size of probe in number of heavy atoms (3-8 atoms)
        :param str buriedness_method: Either 'ghecom' or 'ligsite'
        :param int nprocesses: number of CPU's used
        :param `hotspots.calculation.Runner.Settings` settings: holds the calculation settings
        :return: a :class:`hotspots.result.Result` instance

        >>> from hotspots.calculation import Runner

        >>> runner = Runner()
        >>> runner.from_pdb("1hcl")
        Result()

        """
        protoss = False

        tmp = tempfile.mkdtemp()
        # if  protoss is True:
        #     protoss = Protoss(out_dir=tmp)
        #     self.protein = protoss.add_hydrogens(pdb_code).protein
        #
        # else:
        PDBResult(identifier=pdb_code).download(out_dir=tmp)
        fname = join(tmp, "{}.pdb".format(pdb_code))
        self.protein = Protein.from_file(fname)
        self._prepare_protein(protoss)
        self.charged_probes = charged_probes
        self.probe_size = probe_size
        self.buriedness_method = buriedness_method
        self.clear_tmp = clear_tmp
        self.cavities = None
        if cavities is True:
            self.cavities = Cavity.from_pdb_file(fname)
        self.nprocesses = nprocesses

        if settings is None:
            self.sampler_settings = self.Settings()
        else:
            self.sampler_settings = settings

        if self.sampler_settings.return_probes is True:
            print('here')
            self._calc_hotspots(return_probes=True)

        else:
            self._calc_hotspots()

        self.super_grids = {p: g[0] for p, g in self.out_grids.items()}

        if clear_tmp == True:
            shutil.rmtree(tmp)

        return Results(super_grids=self.super_grids,
                       protein=self.protein,
                       buriedness=self.buriedness,
                       superstar={x.identifier: x.grid for x in self.superstar_grids},
                       weighted_superstar={x.identifier: x.grid for x in self.weighted_grids})

    def global_dock(self, hr:Results, mol, settings = None):
        self.super_grids = hr.super_grids
        self.buriedness = hr.buriedness
        self.protein = hr.protein
        self.charged_probes = False
        self.probe_size = 7
        self.buriedness_method = 'ligsite'
        self.cavities = None
        self.clear_tmp = True

        if settings is None:
            self.sampler_settings = self.Settings()

        self.sampler_settings.return_probes = True
        self.sampler_settings.nrotations = 2000
        self.apolar_translation_threshold = 17
        self.polar_translation_threshold = 17

        print(hr.weighted_superstar)

        grid_dict = hr.weighted_superstar

        ps = self._dock_probe(mol, grid_dict)
        return ps


"""
Photon emission and absoprtion models.
"""
import numpy as np
import h5py

from soxs.spectra import ApecGenerator, \
    get_wabs_absorb, get_tbabs_absorb
from soxs.utils import parse_prng
from pyxsim.utils import mylog
from yt.funcs import get_pbar
from yt.units.yt_array import YTArray, YTQuantity
from yt.utilities.physical_constants import hcgs, clight

hc = (hcgs*clight).in_units("keV*angstrom").v
# NOTE: XSPEC has hc = 12.39854 keV*A, so there may be slight differences in
# placement of spectral lines due to the above
cl = clight.v

class TableApecModel(ApecGenerator):
    r"""
    Initialize a thermal gas emission model from the AtomDB APEC tables
    available at http://www.atomdb.org. This code borrows heavily from Python
    routines used to read the APEC tables developed by Adam Foster at the
    CfA (afoster@cfa.harvard.edu).

    Parameters
    ----------
    emin : float
        The minimum energy for the spectral model.
    emax : float
        The maximum energy for the spectral model.
    nchan : integer
        The number of channels in the spectral model. If one
        is thermally broadening lines, it is recommended that
        this value result in an energy resolution per channel
        of roughly 1 eV.
    var_elem : list of strings, optional
        The names of elements to allow to vary freely
        from the single abundance parameter. Default:
        None
    model_root : string, optional
        The directory root where the model files are stored. If
        not provided, the default SOXS-provided files are used.
    model_vers : string, optional
        The version identifier string for the APEC files, e.g.
        "3.0.3". Default: 3.0.8
    thermal_broad : boolean, optional
        Whether or not the spectral lines should be thermally
        broadened. Default: True
    nolines : boolean, optional
        Turn off lines entirely for generating spectra.
        Default: False
    abund_table : string or array_like, optional
        The abundance table to be used for solar abundances.
        Either a string corresponding to a built-in table or an array
        of 30 floats corresponding to the abundances of each element
        relative to the abundance of H. Default is "angr".
        Built-in options are:
        "angr" : from Anders E. & Grevesse N. (1989, Geochimica et
        Cosmochimica Acta 53, 197)
        "aspl" : from Asplund M., Grevesse N., Sauval A.J. & Scott
        P. (2009, ARAA, 47, 481)
        "wilm" : from Wilms, Allen & McCray (2000, ApJ 542, 914
        except for elements not listed which are given zero abundance)
        "lodd" : from Lodders, K (2003, ApJ 591, 1220)
    nei : boolean, optional
        If True, use the non-equilibrium ionization tables. These are
        not supplied with pyXSIM/SOXS but must be downloaded separately, in
        which case the *apec_root* parameter must also be set to their
        location. Default: False

    Examples
    --------
    >>> apec_model = TableApecModel(0.05, 50.0, 1000, apec_vers="3.0.3",
    ...                             thermal_broad=False)
    """
    def __init__(self, emin, emax, nchan, var_elem=None,
                 model_root=None, model_vers=None,
                 thermal_broad=True, nolines=False,
                 abund_table="angr", nei=False):
        super(TableApecModel, self).__init__(emin, emax, nchan, var_elem=var_elem,
                                             apec_root=model_root, apec_vers=model_vers,
                                             broadening=thermal_broad, nolines=nolines,
                                             abund_table=abund_table, nei=nei)
        self.nchan = self.nbins

    def prepare_spectrum(self, zobs):
        """
        Prepare the thermal model for execution given a redshift *zobs* for the spectrum.
        """
        cosmic_spec, metal_spec, var_spec = \
            self._get_table(list(range(self.nT)), zobs, 0.0)
        self.cosmic_spec = YTArray(cosmic_spec, "cm**3/s")
        self.metal_spec = YTArray(metal_spec, "cm**3/s")
        if var_spec is None:
            self.var_spec = var_spec
        else:
            self.var_spec = YTArray(var_spec, "cm**3/s")

    def get_spectrum(self, kT):
        """
        Get the thermal emission spectrum given a temperature *kT* in keV.
        """
        tindex = np.searchsorted(self.Tvals, kT)-1
        var_spec = None
        if tindex >= self.Tvals.shape[0]-1 or tindex < 0:
            cosmic_spec = YTArray(np.zeros(self.nchan), "cm**3/s")
            metal_spec = cosmic_spec
            if self.var_spec is not None:
                var_spec = YTArray(np.zeros((self.num_var_elem, self.nchan)), "cm**3/s")
        else:
            dT = (kT-self.Tvals[tindex])/self.dTvals[tindex]
            cspec_l = self.cosmic_spec[tindex, :]
            mspec_l = self.metal_spec[tindex, :]
            cspec_r = self.cosmic_spec[tindex+1, :]
            mspec_r = self.metal_spec[tindex+1, :]
            cosmic_spec = cspec_l*(1.-dT)+cspec_r*dT
            metal_spec = mspec_l*(1.-dT)+mspec_r*dT
            if self.var_spec is not None:
                vspec_l = self.var_spec[:, tindex, :]
                vspec_r = self.var_spec[:, tindex+1, :]
                var_spec = vspec_l*(1.-dT) + vspec_r*dT
        return cosmic_spec, metal_spec, var_spec

    def return_spectrum(self, temperature, metallicity, redshift, norm,
                        velocity=0.0, elem_abund=None):
        """
        Given the properties of a thermal plasma, return a spectrum.

        Parameters
        ----------
        temperature : float
            The temperature of the plasma in keV.
        metallicity : float
            The metallicity of the plasma in solar units.
        redshift : float
            The redshift of the plasma.
        norm : float
            The normalization of the model, in the standard Xspec units of
            1.0e-14*EM/(4*pi*(1+z)**2*D_A**2).
        velocity : float, optional
            Velocity broadening parameter in km/s. Default: 0.0
        elem_abund : dict of element name, float pairs
            A dictionary of elemental abundances to vary
            freely of the abund parameter. Default: None
        """
        spec = super(TableApecModel, self).get_spectrum(temperature, metallicity,
                                                        redshift, norm, velocity=velocity,
                                                        elem_abund=elem_abund)
        return YTArray(spec.flux*spec.de, "photons/s/cm**2")

thermal_models = {"apec": TableApecModel}

class AbsorptionModel(object):
    _name = ""
    def __init__(self, nH, energy, cross_section):
        self.nH = YTQuantity(nH*1.0e22, "cm**-2")
        self.emid = YTArray(energy, "keV")
        self.sigma = YTArray(cross_section, "cm**2")

    def get_absorb(self, e):
        """
        Get the absorption spectrum.
        """
        sigma = np.interp(e, self.emid, self.sigma, left=0.0, right=0.0)
        return np.exp(-sigma*self.nH)

    def absorb_photons(self, eobs, prng=None):
        r"""
        Determine which photons will be absorbed by foreground
        galactic absorption.

        Parameters
        ----------
        eobs : array_like
            The energies of the photons in keV.
        prng : integer, :class:`~numpy.random.RandomState` object or :mod:`~numpy.random`, optional
            A pseudo-random number generator. Typically will only be specified
            if you have a reason to generate the same set of random numbers, such as for a
            test. Default is the :mod:`numpy.random` module.
        """
        prng = parse_prng(prng)
        n_events = eobs.size
        if n_events == 0:
            return np.array([], dtype='bool')
        detected = np.zeros(n_events, dtype='bool')
        nchunk = n_events // 100
        if nchunk == 0:
            nchunk = n_events
        k = 0
        pbar = get_pbar("Absorbing photons", n_events)
        while k < n_events:
            absorb = self.get_absorb(eobs[k:k+nchunk])
            nabs = absorb.size
            randvec = prng.uniform(size=nabs)
            detected[k:k+nabs] = randvec < absorb
            k += nabs
            pbar.update(k)
        pbar.finish()
        return detected

class TBabsModel(AbsorptionModel):
    r"""
    Initialize a Tuebingen-Boulder (Wilms, J., Allen, A., &
    McCray, R. 2000, ApJ, 542, 914) ISM absorption model.

    Parameters
    ----------
    nH : float
        The foreground column density in units of 10^22 cm^{-2}.

    Examples
    --------
    >>> tbabs_model = TBabsModel(0.1)
    """
    _name = "tbabs"
    def __init__(self, nH):
        self.nH = YTQuantity(nH, "1.0e22*cm**-2")

    def get_absorb(self, e):
        e = np.array(e)
        return get_tbabs_absorb(e, self.nH.v)

class WabsModel(AbsorptionModel):
    r"""
    Initialize a Wisconsin (Morrison and McCammon; ApJ 270, 119)
    absorption model.

    Parameters
    ----------
    nH : float
        The foreground column density in units of 10^22 cm^{-2}.

    Examples
    --------
    >>> wabs_model = WabsModel(0.1)
    """
    _name = "wabs"
    def __init__(self, nH):
        self.nH = YTQuantity(nH, "1.0e22*cm**-2")

    def get_absorb(self, e):
        e = np.array(e)
        return get_wabs_absorb(e, self.nH.v)

absorb_models = {"wabs": WabsModel,
                 "tbabs": TBabsModel}