"""Main module."""
import os.path
import numpy as np
from astropy.io import fits
import configparser
import copy

from . import correlation_item, data, utils
from vega.scale_parameters import ScaleParameters
from vega.model import Model
from vega.minimizer import Minimizer
from vega.analysis import Analysis
from vega.output import Output
from vega.parameters.param_utils import get_default_values
from vega.plots.plot import VegaPlots

class VegaInterface:
    """Main Vega class.

    Parse the main config and initialize a correlation item for each component.

    If there is data, initialize data and model objects for each component.

    Handle the parameter config and call the analysis class.
    """
    _blind = None

    def __init__(self, main_path):
        """

        Parameters
        ----------
        main_path : string
            Path to main.ini config file
        """
        # Read the main config file
        self.main_config = configparser.ConfigParser()
        self.main_config.optionxform = lambda option: option
        self.main_config.read(utils.find_file(main_path))

        # Read the fiducial pk file
        self.fiducial = self._read_fiducial(self.main_config['fiducial'])

        # Read the effective redshift and the data config paths
        self.fiducial['z_eff'] = self.main_config['data sets'].getfloat('zeff')
        write_cf = self.main_config['output'].getboolean('write_cf', False)
        write_pk = self.main_config['output'].getboolean('write_pk', False)
        self.fiducial['save-components'] = write_cf or write_pk
        ini_files = self.main_config['data sets'].get('ini files').split()

        # Initialize the individual components
        self.corr_items = {}
        for path in ini_files:
            config = configparser.ConfigParser()
            config.optionxform = lambda option: option
            config.read(utils.find_file(os.path.expandvars(path)))

            name = config['data'].get('name')
            self.corr_items[name] = correlation_item.CorrelationItem(config)

        # Check if all correlations have data files
        self.data = {}
        self._has_data = True
        for name, corr_item in self.corr_items.items():
            if not corr_item.has_data:
                self._has_data = False

        # Initialize the data
        for name, corr_item in self.corr_items.items():
            if self._has_data:
                self.data[name] = data.Data(corr_item)
            else:
                self.data[name] = None

        # Initialize scale parameters
        self.scale_params = ScaleParameters(self.main_config['cosmo-fit type'])

        # initialize the models
        self.models = {}
        if self._has_data:
            for name, corr_item in self.corr_items.items():
                self.models[name] = Model(corr_item, self.fiducial, self.scale_params,
                                          self.data[name])

        # Read parameters
        self.params = self._read_parameters(self.corr_items, self.main_config['parameters'])
        self.sample_params = self._read_sample(self.main_config['sample'])

        # Check blinding
        self._scale_par_names = ['ap', 'at', 'ap_sb', 'at_sb', 'phi', 'gamma', 'alpha',
                                 'phi_smooth', 'gamma_smooth', 'alpha_smooth', 'aiso', 'epsilon']
        if self._has_data:
            self._blind = False
            for data_obj in self.data.values():
                if data_obj.blind:
                    self._blind = True
            if self._blind:
                for par in self.sample_params['limits'].keys():
                    if par in self._scale_par_names:
                        raise ValueError('Running on blind data, please fix scale parameters')

        # Get priors
        self.priors = {}
        if 'priors' in self.main_config:
            self.priors = self._init_priors(self.main_config['priors'])
            for param in self.priors.keys():
                if param not in self.sample_params['limits'].keys():
                    print('Warning: Prior specified for a parameter that is'
                          ' not sampled!')

        # Read the monte carlo parameters
        self.mc_config = None
        if 'monte carlo' in self.main_config:
            self.mc_config = {}
            config = self.main_config['monte carlo']

            self.mc_config['params'] = copy.deepcopy(self.params)
            mc_params = self.main_config['mc parameters']
            for param, value in mc_params.items():
                self.mc_config['params'][param] = float(value)

            self.mc_config['sample'] = self._read_sample(config)

        # Initialize the minimizer and the analysis objects
        if not self.sample_params['limits']:
            self.minimizer = None
        else:
            self.minimizer = Minimizer(self.chi2, self.sample_params)
        self.analysis = Analysis(Minimizer(self.chi2, self.sample_params),
                                 self.main_config, self.mc_config)

        # Check for sampler
        self.has_sampler = False
        if 'control' in self.main_config:
            self.has_sampler = self.main_config['control'].getboolean('sampler', False)
            if self.has_sampler:
                if 'Polychord' not in self.main_config:
                    raise RuntimeError('run_sampler called, but no sampler initialized')

        self.output = Output(self.main_config['output'], self.data, self.corr_items, self.analysis)

        self.monte_carlo = False
        self.plots = None
        if self._has_data:
            self.plots = VegaPlots(vega_data=self.data)

    def compute_model(self, params=None, run_init=True, direct_pk=None):
        """Compute correlation function model using input parameters.

        Parameters
        ----------
        params : dict, optional
            Computation parameters, by default None
        run_init: boolean, optional
            Whether to run model.init() before computing the model, by default True
        direct_pk: 1D array or None, optional
            If not None, the full Pk (e.g. from CLASS/CAMB) to be used directly, by default None

        Returns
        -------
        dict
            Dictionary of cf models for each component
        """
        # Overwrite computation parameters
        local_params = copy.deepcopy(self.params)
        if params is not None:
            for par, val in params.items():
                local_params[par] = val

        # Go through each component and compute the model cf
        model_cf = {}
        if run_init:
            self.models = {}
        for name, corr_item in self.corr_items.items():
            if run_init:
                self.models[name] = Model(corr_item, self.fiducial, self.scale_params,
                                          self.data[name])

            if direct_pk is None:
                model_cf[name] = self.models[name].compute(local_params, self.fiducial['pk_full'],
                                                           self.fiducial['pk_smooth'])
            else:
                model_cf[name] = self.models[name].compute_direct(local_params, direct_pk)

        return model_cf

    def chi2(self, params=None, direct_pk=None):
        """Compute full chi2 for all components.

        Parameters
        ----------
        params : dict, optional
            Computation parameters, by default None
        direct_pk: 1D array or None, optional
            If not None, the full Pk (e.g. from CLASS/CAMB) to be used directly, by default None

        Returns
        -------
        float
            chi^2
        """
        assert self._has_data

        # Check if blinding is initialized
        if self._blind is None:
            self._blind = False
            for data_obj in self.data.values():
                if data_obj.blind:
                    self._blind = True

        # Overwrite computation parameters
        local_params = copy.deepcopy(self.params)
        if params is not None:
            for par, val in params.items():
                local_params[par] = val

        # Enforce blinding
        if self._blind:
            for par, val in local_params.items():
                if par in self._scale_pars:
                    local_params[par] = 1.

        # Go trough each component and compute the chi^2
        chi2 = 0
        for name in self.corr_items:
            try:
                if direct_pk is None:
                    model_cf = self.models[name].compute(local_params, self.fiducial['pk_full'],
                                                         self.fiducial['pk_smooth'])
                else:
                    model_cf = self.models[name].compute_direct(local_params, direct_pk)
            except utils.VegaBoundsError:
                self.models[name].PktoXi.cache_pars = None
                return 1e100

            if self.monte_carlo:
                diff = self.data[name].masked_mc_mock - model_cf[self.data[name].mask]
                chi2 += diff.T.dot(self.data[name].scaled_inv_masked_cov.dot(diff))
            else:
                diff = self.data[name].masked_data_vec - model_cf[self.data[name].mask]
                chi2 += diff.T.dot(self.data[name].inv_masked_cov.dot(diff))

        # Add priors
        for param, prior in self.priors.items():
            chi2 += self._gaussian_chi2_prior(local_params[param], prior[0], prior[1])

        assert isinstance(chi2, float)
        return chi2

    def log_lik(self, params=None, direct_pk=None):
        """Compute full log likelihood for all components.

        Parameters
        ----------
        params : dict, optional
            Computation parameters, by default None
        direct_pk: 1D array or None, optional
            If not None, the full Pk (e.g. from CLASS/CAMB) to be used directly, by default None

        Returns
        -------
        float
            log Likelihood
        """
        assert self._has_data

        # Get the full chi2
        chi2 = self.chi2(params, direct_pk)

        # Compute the normalization for each component
        log_norm = 0
        for name in self.corr_items:
            log_norm -= 0.5 * self.data[name].data_size * np.log(2 * np.pi)

            if self.monte_carlo:
                log_norm -= 0.5 * self.data[name].scaled_log_cov_det
            else:
                log_norm -= 0.5 * self.data[name].log_cov_det

        # Compute log lik
        log_lik = log_norm - 0.5 * chi2

        # Add priors normalization
        for param, prior in self.priors.items():
            log_lik += self._gaussian_lik_prior(prior[1])

        return log_lik

    def monte_carlo_sim(self, params=None, scale=None, seed=0, forecast=False):
        """Compute Monte Carlo simulations for each Correlation item.

        Parameters
        ----------
        params : dict, optional
            Computation parameters, by default None
        scale : float/dict, optional
            Scaling for the covariance, by default 1.
        seed : int, optional
            Seed for the random number generator, by default 0
        forecast : boolean, optional
            Forecast option. If true, we don't add noise to the mock,
            by default False

        Returns
        -------
        dict
            Dictionary with MC mocks for each item
        """
        assert self._has_data

        # Overwrite computation parameters
        local_params = copy.deepcopy(self.params)
        if params is not None:
            for par, val in params.items():
                local_params[par] = val

        mocks = {}
        for name in self.corr_items:
            # Compute fiducial model
            fiducial_model = self.models[name].compute(
                local_params, self.fiducial['pk_full'],
                self.fiducial['pk_smooth'])

            # Get scale
            if scale is None:
                item_scale = self.corr_items[name].cov_rescale
            elif type(scale) is float or type(scale) is int:
                item_scale = scale
            elif name in scale:
                item_scale = scale[name]
            else:
                item_scale = 1.

            # Create the mock
            mocks[name] = self.data[name].create_monte_carlo(fiducial_model, item_scale, seed,
                                                             forecast)

        self.monte_carlo = True
        return mocks

    def minimize(self):
        """Minimize the chi2 over the sampled parameters.
        """
        if self.minimizer is None:
            print("No sampled parameters. Skipping minimization.")
            return

        # if not self.fiducial['save-components']:
            # self.set_fast_metals()

        self.minimizer.minimize()

    @property
    def bestfit(self):
        """Access the bestfit results from iminuit.

        Returns
        -------
        Minimizer
            Returns the Minimizer class which stores the bestfit values
        """
        return self.minimizer

    def set_fast_metals(self):
        """Activate fast metals. This is automatically called when
        running the minimizer or the sampler.
        """
        print('Warning! Activating fast metals for minimizing/sampling.')
        for name in self.corr_items:
            if self.models[name].metals is not None:
                self.models[name].metals.fast_metals = True

    @staticmethod
    def _read_fiducial(fiducial_config):
        """Read the fiducial pk file and get the configs.

        Parameters
        ----------
        fiducial_config : ConfigParser
            fiducial section from the main config file

        Returns
        -------
        dict
            dictionary with the fiducial data and config
        """
        # First check the path and replace with the right model if necessary
        path = fiducial_config.get('filename')
        path = utils.find_file(os.path.expandvars(path))
        # if not os.path.isfile(path):
        # path = resource_filename('vega', 'models') + '/{}'.format(path)
        print('INFO: reading input Pk {}'.format(path))

        fiducial = {}

        # Open the fits file and get what we need
        hdul = fits.open(path)
        fiducial['z_fiducial'] = hdul[1].header['ZREF']
        fiducial['Omega_m'] = hdul[1].header['OM']
        fiducial['Omega_de'] = hdul[1].header['OL']
        fiducial['k'] = hdul[1].data['K']
        fiducial['pk_full'] = hdul[1].data['PK']
        fiducial['pk_smooth'] = hdul[1].data['PKSB']
        hdul.close()

        return fiducial

    @staticmethod
    def _read_parameters(corr_items, parameters_config):
        """Read computation parameters.

        If a parameter is specified multiple times,
        the parameters in the main config file have priority.

        Parameters
        ----------
        corr_items : dict
            Dictionary of correlation items
        parameters_config : ConfigParser
            parameters section from main config

        Returns
        -------
        dict
            Computation parameters
        """
        params = {}

        # First get the parameters from each component config
        for name, corr_item in corr_items.items():
            if 'parameters' in corr_item.config:
                for param, value in corr_item.config.items('parameters'):
                    params[param] = float(value)

        # Next get the parameters in the main config
        for param, value in parameters_config.items():
            params[param] = float(value)

        return params

    def _read_sample(self, sample_config):
        """Read sample parameters.

        These must be of the form:

        param = min max / for sampler only
        or
        param = min max val err / for both sampler and fitter.

        Fitter accepts None for min/max, but the sampler does not.

        Parameters
        ----------
        sample_config : ConfigParser
            sample section from main config

        Returns
        -------
        dict
            Config for the sampled parameters
        """
        # Initialize the dictionaries we need
        sample_params = {}
        sample_params['limits'] = {}
        sample_params['values'] = {}
        sample_params['errors'] = {}
        sample_params['fix'] = {}

        default_values = get_default_values()

        def check_param(param):
            if param not in default_values:
                raise ValueError('Default values not found for: %s. Please add'
                                 ' them to default_values.txt, or provide the'
                                 ' full sampling specification.' % param)

        for param, values in sample_config.items():
            if param not in self.params:
                print('Warning: You tried sampling the parameter: %s.'
                      ' As this parameter was not specified under'
                      ' [parameters], it will be skipped.' % param)
                continue

            values_list = values.split()

            # Get the prior limits
            # ! Sampler needs actual values (no None)
            if len(values_list) > 1:
                lower_limit = None
                upper_limit = None
                if values_list[0] != 'None':
                    lower_limit = float(values_list[0])
                if values_list[1] != 'None':
                    upper_limit = float(values_list[1])
                sample_params['limits'][param] = (lower_limit, upper_limit)
            else:
                if values_list[0] not in ['True', 'true', 't', 'y', 'yes']:
                    continue
                check_param(param)
                sample_params['limits'][param] = default_values[param]['limits']

            # Get the values and errors for the fitter
            if len(values_list) > 2:
                sample_params['values'][param] = float(values_list[2])
            else:
                check_param(param)
                sample_params['values'][param] = self.params[param]

            if len(values_list) > 3:
                assert len(values_list) == 4
                sample_params['errors'][param] = float(values_list[3])
            else:
                check_param(param)
                sample_params['errors'][param] = default_values[param]['error']

            # Populate the fix values
            sample_params['fix'][param] = False

        return sample_params

    @staticmethod
    def _gaussian_chi2_prior(value, mean, sigma):
        return (value - mean)**2 / sigma**2

    @staticmethod
    def _gaussian_lik_prior(sigma):
        return -0.5 * np.log(2 * np.pi) - np.log(sigma)

    @staticmethod
    def _init_priors(prior_config):
        """Initialize the priors. Only gaussian priors are currently supported

        Parameters
        ----------
        prior_config : ConfigParser
            priors section from main config

        Returns
        -------
        dict
            Dictionary of priors (mean, sigma) with the keys as parameter names
        """
        prior_dict = {}
        for param, prior in prior_config.items():
            prior_list = prior.split()
            if len(prior_list) != 3:
                raise ValueError('Prior configuration must have the format:'
                                 ' "<param> = gaussian <mean> <sigma>"')
            if prior_list[0] not in ['gaussian', 'Gaussian']:
                raise ValueError('Only gaussian priors are supported.')

            prior_dict[param] = np.array(prior_list[1:]).astype(float)

        return prior_dict


# ---
# jupyter:
#   jupytext:
#     text_representation:
#       extension: .py
#       format_name: light
#       format_version: '1.5'
#       jupytext_version: 1.11.3
#   kernelspec:
#     display_name: Python 3
#     language: python
#     name: python3
# ---

# + [markdown] id="view-in-github" colab_type="text"
#

# + [markdown] id="qwWX1CpCAvwp"
# # Tabular Sarsa Algorithm and the Taxi Environment
#
# Authors: Fred Amouzgar <fred.amouzgar@mq.edu.au> and Kevin Murphy <murphyk@gmail.com>
#
#
# Tabular methods are suitable for small and discrete state space and discrete action space environments. So, the state-action function (Q) can be represented by a table of values. For large state space environments, we prefer to use approximation methods such as neural networks. However, the simplicity of tabular methods' implementation is helpful to demonstrate RL method's functionality.  In this notebook, we train a SARSA agent for [OpenAI's Taxi Gym](https://gym.openai.com/envs/Taxi-v2/) environment (an example originally proposed by Tom Dietterich).

# + [markdown] id="IphbO1Y2BQ-J"
# ## 1- Installations

# + id="IYHrnYXTBuXe"

try:
    # # %tensorflow_version only exists in Colab.
    # %tensorflow_version 2.x
    IS_COLAB = True
except Exception:
    IS_COLAB = False

if not(IS_COLAB):
  print('not a colab')
  # !pip -q install gym numpy matplotlib

# + [markdown] id="vrJGjKw-B6X7"
# ## 2- Setting up the Environment

# + colab={"base_uri": "https://localhost:8080/"} id="NA58-PytB44R" outputId="6a1ab6ba-6fa9-454a-d3cc-173d7eba1f1f"
import gym

# Finding the Taxi environment
for env in gym.envs.registry.all():
    if env.id.startswith("Taxi"):
        env_name = env.id
##

print("Environment Name:", env_name)
env = gym.make(env_name)
env.reset()
env.render()

# + [markdown] id="N2q3JJNoIcdP"
# Here is a description of the taxi environment from the [docstring](https://github.com/openai/gym/blob/master/gym/envs/toy_text/taxi.py).
#
#
# ### Description:
#
# There are four designated locations in the grid world indicated by R(ed), G(reen), Y(ellow), and B(lue). When the episode starts, the taxi starts off at a random square and the passenger is at a random location. The taxi drives to the passenger's location, picks up the passenger, drives to the passenger's destination (another one of the four specified locations), and then drops off the passenger. Once the passenger is dropped off, the episode ends.
#
# ### States/ Observations:
#
# State space is (taxi_row, taxi_col, passenger_location, destination).
#
# There are 500 discrete states since there are 25 taxi positions, 5 possible locations of the passenger (including the case when the passenger is in the taxi), and 4 destination locations.
#
#
# ### Passenger locations:
# - 0: R(ed)
# - 1: G(reen)
# - 2: Y(ellow)
# - 3: B(lue)
# - 4: in taxi
#
# ### Destinations:
# - 0: R(ed)
# - 1: G(reen)
# - 2: Y(ellow)
# - 3: B(lue)
#
# ### Actions:
#
# There are 6 discrete deterministic actions:
# - 0: move south
# - 1: move north
# - 2: move east
# - 3: move west
# - 4: pickup passenger
# - 5: drop off passenger
#
#
# Notice in this environment the taxi cannot perform certain actions in certain states due to walls. In environment's code, we will simply provide a -1 penalty for every wall hit and the taxi won't move anywhere. This will just rack up penalties causing the taxi to consider going around the wall.
#
# ### Rewards:
#
# There is a default per-step reward of -1,
# except for delivering the passenger, which is +20,
# or executing "pickup" and "drop-off" actions illegally, which is -10.
#
# ### Rendering:
# - blue: passenger
# - magenta: destination
# - yellow: empty taxi
# - green: full taxi
# - other letters (R, G, Y and B): locations for passengers and destinations
#
#

# + [markdown] id="2U7EBEtzCyhH"
# ## 3- Developing the SARSA agent
#
#
#
# Here's the full update formula covered in line 7 and 8:
#

# + id="FFghKMePC4tl"
import numpy as np
import pickle

class Sarsa_Agent:
    def __init__(self, states_n, actions_n, learning_rate=0.2, epsilon=0.1, gamma=0.95, epsilon_decay=True,
                 epsilon_decay_factor=0.01):
        self.learning_rate = learning_rate
        self.epsilon = epsilon
        self.gamma = gamma
        self.states_n = states_n
        self.actions_n = actions_n
        self.Q = np.zeros((states_n, actions_n))
        self.new_a = None
        self.epsilon_decay = epsilon_decay
        self.epsilon_decay_factor = epsilon_decay_factor

    def act(self, state):
        """The act method implements the epsilon-greedy policy"""
        if np.random.rand() < self.epsilon:
            act = np.random.choice(np.arange(self.actions_n))
        else:
            act = np.argmax(self.Q[int(state), :])
        return act

    def decay_epsilon(self, factor):
        """Decaying the epsilon, so it gradualy reduces the exploration and exploits more"""
        self.epsilon -= factor if self.epsilon >= 0 else 0

    def update(self, new_s, r, s, a, done):
        """The update method updates the agent for one step"""
        self.new_a = self.act(new_s)
        mask = 0 if done else 1
        s, a, self.new_a, new_s = int(s), int(a), int(self.new_a), int(new_s)
        self.Q[s, a] += self.learning_rate * (r + self.gamma * self.Q[new_s, self.new_a] * mask - self.Q[s, a])
        if done and self.epsilon_decay:
            self.decay_epsilon(self.epsilon_decay_factor)
        return self.new_a

    def save(self, file_name="taxi.pkl"):
        """The save method saves (pickles) the agent's Q table"""
        with open(file_name, mode="wb") as f:
            pickle.dump(self.Q, f)

    def load(self, file_name="taxi.pkl"):
        """The load method loads a pickled Q table"""
        with open(file_name, mode="rb") as f:
            self.Q = pickle.load(f)

# + [markdown] id="mmSjC-diEcQS"
# ## 4- Defining the Training Loop

# + id="oMOOyLDdEhtJ"
from IPython.display import clear_output
import matplotlib.pyplot as plt
from time import sleep
import numpy as np

def train_taxi(env, agent, episodes=150):
    if env is None:
        raise ValueError("No Environment is given.")
    if agent is None:
        raise ValueError("No agent is given.")

    steps = []
    returns = []
    for episode in range(episodes):
        state = env.reset()
        action = agent.act(state)
        done = False
        step_n = 0
        return_episode = 0
        while not done:
            new_state, reward, done, _ = env.step(action)
            return_episode += reward
            new_action = agent.update(new_state,reward,state,action,done)
            state, action = new_state, new_action
            step_n += 1
            if done:
                steps.append(step_n)
                returns.append(return_episode)
                clear_output(wait=True)
                plt.title("Steps:" + str(step_n) + " Return:"+str(return_episode))
                plt.plot(list(range(len(steps))),steps)
                plt.plot(list(range(len(steps))),returns)
                plt.legend(["Steps", "Returns"])
                plt.show()

# + [markdown] id="-AEYmZgQFB_w"
# ## 5- Let's train our agent for 1500 episodes (takes ~5 minutes)

# + colab={"base_uri": "https://localhost:8080/", "height": 281} id="1GV0owZVEoQ0" outputId="5bb38dc9-bcfe-4a61-bfda-7c70313a7cc8"
agent = Sarsa_Agent(env.observation_space.n, env.action_space.n,epsilon_decay=True)

train_taxi(env, agent, episodes=1500)
agent.save()

# + [markdown] id="XS2D27_CmdpM"
# ## 6- Methods for Playing and Rendering the Taxi environment in the notebook

# + id="HLiUntsMmPU0"
def taxi_print_frames(frames, wait_btw_frames, episode):
    for i, frame in enumerate(frames):
        clear_output(wait=True)
        print(frame['frame'])
        print(f"Passenger #: {episode + 1}")
        print("-----------")
        print(f"Timestep: {i + 1}")
        print(f"State: {frame['state']}")
        print(f"Action: {frame['action']}")
        print(f"Reward: {frame['reward']}")
        sleep(wait_btw_frames)

def play_taxi(env, agent, passengers=2, wait_btw_frames=1):
    for episode in range(passengers):
        state = env.reset()
        frames = []
        done = False
        step = 0
        while not done:
            action = agent.act(state)
            new_state, reward, done, _ = env.step(action)
            frames.append({
                'frame': env.render(mode='ansi'),
                'state': state,
                'action': action,
                'reward': reward
            })
            step += 1
            state = new_state
        taxi_print_frames(frames, wait_btw_frames=wait_btw_frames, episode=episode)

# + [markdown] id="qUqhWVoeng_-"
# ## 7- Watch a Trained SARSA Cab Driver
#
# Note: You can change the number of passengers if you want to move more than 3. Change the wait_btw_frames if you want to see the game running faster or slower.

# + colab={"base_uri": "https://localhost:8080/"} id="Zmyw9mKsmRZs" outputId="8dfe9b5a-4084-4779-95c7-26431ba2f4c5"
play_taxi(env, agent, passengers=3, wait_btw_frames=1)


r"""
HPF Spectrum
---------------

A container for an HPF spectrum of :math:`M=28` total total orders :math:`m`, each with vectors for wavelength flux and uncertainty, e.g. :math:`F_m(\lambda)`.  HPF additionally has a sky fiber and optionally a Laser Frequency Comb fiber.  Our experimental API currently ignores the LFC fiber.  The sky fiber can be accessed by passing the `sky=True` kwarg when retrieving the

HPFSpectrum
##############
"""

import warnings
import logging
from muler.echelle import EchelleSpectrum, EchelleSpectrumList
import numpy as np
import astropy
from astropy.io import fits
from astropy import units as u
from astropy.wcs import WCS, FITSFixedWarning
from astropy.nddata import StdDevUncertainty
from scipy.interpolate import InterpolatedUnivariateSpline
from astropy.constants import R_jup, R_sun, G, M_jup, R_earth, c
from astropy.time import Time
import copy
from importlib_resources import files
from specutils.manipulation import LinearInterpolatedResampler
from scipy.ndimage import binary_dilation
from . import templates
import pandas as pd

log = logging.getLogger(__name__)

for category in [
    astropy.utils.exceptions.AstropyDeprecationWarning,
    FITSFixedWarning,
    RuntimeWarning,
]:
    warnings.filterwarnings("ignore", category=category)

# Convert FITS running index number to echelle order m
grating_order_offsets = {"Goldilocks": 0, "HPF": 0}  # Not implemented yet

# Science-to-sky fiber throughput ratio template
static_sky_ratio_file = files(templates).joinpath("HPF_sci_to_sky_ratio_beta.csv")
STATIC_SKY_RATIO_DATAFRAME = pd.read_csv(static_sky_ratio_file)

# Blaze function template
static_blaze_file = files(templates).joinpath("HPF_blaze_templates.csv")
STATIC_BLAZE_DATAFRAME = pd.read_csv(static_blaze_file)

# A0V template
static_A0V_file = files(templates).joinpath("PHOENIX_10kK_hpf_template.csv")
STATIC_A0V_DATAFRAME = pd.read_csv(static_A0V_file)

# TelFit template
static_telfit_file = files(templates).joinpath("telfit_HPFtemplate_temp286_hum050.csv")
STATIC_TELFIT_DATAFRAME = pd.read_csv(static_telfit_file)

class HPFSpectrum(EchelleSpectrum):
    r"""
    A container for HPF spectra

    Args:
        file (str): A path to a reduced HPF spectrum from Goldilocks *or* the HPF instrument team
        order (int): which spectral order to read
        cached_hdus (list) :
            A pre-loaded HDU to reduce file I/O for multiorder access.
            If provided, must give both HDUs.  Optional, default is None.
    """

    def __init__(self, *args, file=None, order=19, cached_hdus=None, **kwargs):

        self.site_name = "mcdonald"
        self.ancillary_spectra = ["sky", "lfc"]
        self.noisy_edges = (3, 2045)
        self.instrumental_resolution = 55_000.0

        if file is not None:
            if "Goldilocks" in file:
                pipeline = "Goldilocks"
            elif "Slope" in file:
                pipeline = "HPF"
            else:
                raise NameError("Cannot identify file as an HPF spectrum")
            grating_order = grating_order_offsets[pipeline] + order

            if cached_hdus is not None:
                hdus = cached_hdus[0]
            else:
                hdus = fits.open(str(file))
            hdr = hdus[0].header

            ## Target Spectrum
            lamb = hdus[7].data[order].astype(np.float64) * u.AA
            flux = hdus[1].data[order].astype(np.float64) * u.ct
            unc = hdus[4].data[order].astype(np.float64) * u.ct
            if pipeline == "HPF":
                unc = np.sqrt(unc.value) * u.ct

            meta_dict = {
                "x_values": np.arange(0, 2048, 1, dtype=np.int),
                "pipeline": pipeline,
                "m": grating_order,
                "header": hdr,
            }

            uncertainty = StdDevUncertainty(unc)
            mask = (
                np.isnan(flux) | np.isnan(uncertainty.array) | (uncertainty.array <= 0)
            )

            super().__init__(
                spectral_axis=lamb,
                flux=flux,
                mask=mask,
                wcs=None,
                uncertainty=uncertainty,
                meta=meta_dict,
                **kwargs,
            )

            ## Sky Spectrum
            lamb = hdus[8].data[order].astype(np.float64) * u.AA
            flux = hdus[2].data[order].astype(np.float64) * u.ct
            unc = hdus[5].data[order].astype(np.float64) * u.ct
            if pipeline == "HPF":
                unc = np.sqrt(unc.value) * u.ct
            uncertainty = StdDevUncertainty(unc)
            mask = (
                np.isnan(flux) | np.isnan(uncertainty.array) | (uncertainty.array <= 0)
            )
            sky_spectrum = HPFSpectrum(
                spectral_axis=lamb,
                flux=flux,
                mask=mask,
                wcs=None,
                uncertainty=uncertainty,
                meta=meta_dict.copy(),
                **kwargs,
            )

            ## LFC Spectrum
            lamb = hdus[9].data[order].astype(np.float64) * u.AA
            flux = hdus[3].data[order].astype(np.float64) * u.ct
            unc = hdus[6].data[order].astype(np.float64) * u.ct
            if pipeline == "HPF":
                unc = np.sqrt(unc.value) * u.ct
            uncertainty = StdDevUncertainty(unc)
            mask = (
                np.isnan(flux) | np.isnan(uncertainty.array) | (uncertainty.array <= 0)
            )
            lfc_spectrum = HPFSpectrum(
                spectral_axis=lamb,
                flux=flux,
                mask=mask,
                wcs=None,
                uncertainty=uncertainty,
                meta=meta_dict.copy(),
                **kwargs,
            )

            ## We could optionally enable lfc and sky metadata for these referece spectra
            ## That's slightly redundant, it enables antipatterns like:
            # `spectrum.sky.lfc` rather than simply `spectrum.lfc`

            # sky_spectrum.meta["lfc"] = lfc_spectrum
            # lfc_spectrum.meta["sky"] = sky_spectrum

            sky_spectrum.meta["provenance"] = "Sky fiber"
            lfc_spectrum.meta["provenance"] = "Laser Frequency Comb"
            self.meta["provenance"] = "Target fiber"

            self.meta["sky"] = sky_spectrum
            self.meta["lfc"] = lfc_spectrum

        else:
            super().__init__(*args, **kwargs)

    @property
    def provenance(self):
        """What is the provenance of each spectrum?"""
        return self.meta["provenance"]

    @property
    def pipeline(self):
        """Which pipeline does this spectrum originate from?"""
        return self.meta["pipeline"]

    @property
    def spectrographname(self):
        """What's the name of the spectrograph?"""
        return "HPF"

    @property
    def sky(self):
        """Sky fiber spectrum stored as its own HPFSpectrum object"""
        return self.meta["sky"]

    @property
    def lfc(self):
        """Sky fiber spectrum stored as its own HPFSpectrum object"""
        return self.meta["lfc"]

    @property
    def RA(self):
        """The right ascension from header files"""
        return self.meta["header"]["RA"] * u.hourangle

    @property
    def DEC(self):
        """The declination from header files"""
        return self.meta["header"]["DEC"] * u.deg

    @property
    def astropy_time(self):
        """The astropy time based on the header"""
        mjd = self.meta["header"]["DATE-OBS"]
        return Time(mjd, format="isot", scale="utc")

    def get_static_blaze_template(self, method="Goldilocks"):
        """Get the static blaze template for HPF, as estimated by Goldilocks

        Parameters
        ----------
        method : (Str)
            Either "Goldilocks" or "2021_median" (default: Goldilocks)
        """
        type_dict = {"Goldilocks": "blaze_Goldilocks", "2021_median": "blaze_2021"}
        assert method in type_dict.keys()
        blaze_type = type_dict[method]

        # Watch out! Some HPFSpectrum methods *will not work* on this calibration spectrum!
        return HPFSpectrum(
            spectral_axis=STATIC_BLAZE_DATAFRAME.wavelength_Angstrom.values
            * u.Angstrom,
            flux=STATIC_BLAZE_DATAFRAME[blaze_type].values * u.dimensionless_unscaled,
        )

    def get_static_sky_ratio_template(self):
        """Get the static sky ratio template for HPF, as estimated from twilight flats"""

        # Watch out! Some HPFSpectrum methods *will not work* on this calibration spectrum!
        return HPFSpectrum(
            spectral_axis=STATIC_SKY_RATIO_DATAFRAME.wave_Ang.values * u.Angstrom,
            flux=STATIC_SKY_RATIO_DATAFRAME.beta_estimator.values
            * u.dimensionless_unscaled,
        )

    def get_static_A0V_template(self, method="PHOENIX"):
        """Get the static A0V template for HPF, as estimated by either Vega or PHOENIX

        Parameters
        ----------
        method : (Str)
            What template to use.  Currently only a state PHOENIX model is supported.
            Other A0V templates may be added in the future, such as Vega.
            (default: PHOENIX)
        """
        if method == "PHOENIX":

            return HPFSpectrum(
                spectral_axis=STATIC_A0V_DATAFRAME.wave_ang.values * u.Angstrom,
                flux=STATIC_A0V_DATAFRAME.flux.values * u.dimensionless_unscaled,
            )
        else:
            raise NotImplementedError

    def get_static_TelFit_template(self):
        """Get the static TelFit template for HPF

        A convenience function for getting a quicklook Telluric template
        """

        return HPFSpectrum(
            spectral_axis=STATIC_TELFIT_DATAFRAME.wavelength_A.values * u.Angstrom,
            flux=STATIC_TELFIT_DATAFRAME.transmission.values * u.dimensionless_unscaled,
        )

    def _deblaze_by_template(self):
        """Deblazing with a template-based method"""
        blaze_template = self.get_static_blaze_template(method="Goldilocks")
        resampler = LinearInterpolatedResampler()
        resampled_blaze = resampler(blaze_template, self.wavelength)
        return self.divide(resampled_blaze, handle_meta="first_found")

    def deblaze(self, method="template"):
        """Override the default spline deblazing with HPF-custom blaze templates.

        Parameters
        ----------
        method : (Str)
            Either "template" or "spline" (default: template)
        """
        if method == "template":
            return self._deblaze_by_template()
        else:
            log.error("This method is deprecated!  Please use the new deblaze method")
            raise NotImplementedError

    def sky_subtract(self, method="scalar"):
        """Subtract sky spectrum from science spectrum, with refinements for sky throughput

        Note: This operation does not wavelength shift or scale the sky spectrum

        Parameters
        ----------
        method : (str)
            The method for sky subtraction: "naive", "scalar", or "vector", as described in
            Gully-Santiago et al. in prep.  Default is scalar.

        Returns
        -------
        sky_subtractedSpec : (HPFSpectrum)
            Sky subtracted Spectrum
        """
        spec = copy.deepcopy(self)
        if method == "naive":
            log.warning(
                "Naive sky subtraction method is known to oversubtract the sky, see GitHub Issues."
            )
            beta = 1.0 * u.dimensionless_unscaled
        elif method == "scalar":
            beta = 0.93 * u.dimensionless_unscaled
        elif method == "vector":
            beta_native_spectrum = spec.get_static_sky_ratio_template()
            resampler = LinearInterpolatedResampler(extrapolation_treatment="zero_fill")
            beta = resampler(beta_native_spectrum, spec.spectral_axis)
        else:
            log.error("Method must be one of 'naive', 'scalar' or 'vector'. ")
            raise NotImplementedError

        # These steps should propagate uncertainty?
        sky_estimator = spec.sky.multiply(beta, handle_meta="first_found")
        return spec.subtract(sky_estimator, handle_meta="first_found")

    def mask_tellurics(self, method="TelFit", threshold=0.999, dilation=5):
        """Mask known telluric lines based on a static TelFit template or heuristics

        Note: This method is for quicklook purpsoes, it misses many unknown tellurics

        Parameters
        ----------
        method : (str)
            The method for telluric masking: "TelFit" or "heuristics"
            Default is TelFit.

        dilation : (int)
            The number of pixels adjacent to the threshold mask to include in a
            dilated mask. This control parameter accounts for velocity offsets
            between the template and observed telluric spectrum.

        Returns
        -------
        sky_subtractedSpec : (HPFSpectrum)
            Sky subtracted Spectrum
        """
        spec = copy.deepcopy(self)

        if method == "TelFit":
            telfit_template = spec.get_static_TelFit_template()
            resampler = LinearInterpolatedResampler(extrapolation_treatment="nan_fill")
            telluric_estimate = resampler(telfit_template, spec.spectral_axis)

            assert (threshold < 1.0) & (threshold > 0.0), "Threshold must be a fraction"
            threshold_mask = telluric_estimate.flux.value < threshold

            # Dilate the binary mask to account for velocity offsets and edge effects
            dilated_mask = binary_dilation(threshold_mask, iterations=dilation)
            assert (
                ~dilated_mask
            ).sum() > 2, "You should have at least 2 pixels left after masking"
            spec = spec._copy(mask=dilated_mask)
            spec_out = spec.remove_nans()
            return spec_out
        else:
            log.error("Only the TelFit method is currently implemented")
            raise NotImplementedError

    def blaze_divide_flats(self, flat, order=19):
        """Remove blaze function from spectrum by dividing by flat spectrum

        Returns
        -------
        blaze corrrected spectrum using flat fields : (HPFSpectrum)

        """
        log.warning("This method is deprecated!  Please use the new deblaze method")
        raise NotImplementedError

class HPFSpectrumList(EchelleSpectrumList):
    r"""
    An enhanced container for a list of HPF spectral orders

    """

    def __init__(self, *args, **kwargs):
        self.normalization_order_index = 14
        super().__init__(*args, **kwargs)

    @staticmethod
    def read(file, precache_hdus=True):
        """Read in a SpectrumList from a file

        Parameters
        ----------
        file : (str)
            A path to a reduced HPF spectrum from plp
        """

        hdus = fits.open(file, memmap=False)
        cached_hdus = [hdus]

        n_orders, n_pix = hdus[7].data.shape

        list_out = []
        for i in range(n_orders):
            spec = HPFSpectrum(file=file, order=i, cached_hdus=cached_hdus)
            list_out.append(spec)
        return HPFSpectrumList(list_out)

    def deblaze(self):
        """Deblaze the entire spectrum"""
        spec_out = copy.copy(self)
        for i in range(len(spec_out)):
            spec_out[i] = spec_out[i].deblaze()

        return spec_out

    def sky_subtract(self, method="vector"):
        """Sky subtract the entire spectrum"""
        spec_out = copy.copy(self)
        for i in range(len(spec_out)):
            spec_out[i] = spec_out[i].sky_subtract(method=method)

        return spec_out

    # def sky_subtract(self):
    #     """Sky subtract all orders
    #     """
    #     flux = copy.deepcopy(self.flux)
    #     sky = copy.deepcopy(self.sky)
    #     for i in range(len(self)):
    #         self[i] = flux[i] - sky[i]

    #     return self

import numpy,from scipy,import astropy,from astropy


#!/usr/bin/python3

import numpy as np
import json
import requests
import csv
import pickle
import os
import sys
import re
import statistics
import scipy
import matplotlib.pyplot as plt
import matplotlib.cm as cm
from matplotlib.colors import ListedColormap, LinearSegmentedColormap
from scipy import cluster
from sklearn.cluster import DBSCAN
from sklearn.metrics import silhouette_samples, silhouette_score
from rdkit import Chem, DataStructs, RDConfig
from rdkit.Chem import rdFMCS, AllChem, Draw, Lipinski
from itertools import chain, combinations

# import ligand smiles strings
ligand_file = 'lig_1118.json'

with open(ligand_file) as ff:
    smiles_dict = json.load(ff)

# check for ligands with bad or missing smiles strings
bad_smiles = []

for key,value in smiles_dict.items():
    if 'smiles_cactus' not in value or Chem.MolFromSmiles(smiles_dict[key]['smiles_cactus']) is None:
        bad_smiles.append(key)

# for each viral protein, make list of residues in consensus pockets
def pocket_residues(consensus_pockets,all_consensus_residues,directory,protnow):
    file = open(directory+'clusters_'+protnow+'.txt','r')
    line_list = file.readlines()
    consensus_pockets[protnow] = {}
    all_consensus_residues[protnow] = []

    for line in line_list:
        pocket = line.split()[0].split(':')[0]
        residues = line.split()[1].split(',')[0:-1]
        consensus_pockets[protnow][pocket] = residues
        for res in residues:
            if res not in all_consensus_residues[protnow]:
                all_consensus_residues[protnow].append(res)

    file.close()
    return consensus_pockets, all_consensus_residues

# for each viral protein pocket, make list of pocket ligands each residue contacts
def pocket_residue_ligand_pairs(directory,filenames,consensus_pockets,consensus_pocket_ligands,protnow):
    consensus_pocket_reslig_pairs[protnow] = {}

    for pocket,residues in consensus_pockets[protnow].items():
        consensus_pocket_reslig_pairs[protnow][pocket] = {}
        for res in residues:
            consensus_pocket_reslig_pairs[protnow][pocket][res] = []

    for fl in filenames:
        file = open(directory+fl,'r')
        line_list = file.readlines()

        for line in line_list:
            interaction = line.split()[0].split(':')[0]
            binding_residues = line.split()[-1].split(',')[0:-1]
            ligand = line.split()[0].split('.')[6]

            # viral protein
            if line.split()[0].split('.')[0].split('_')[0]=='nCoV':
                protein = line.split()[0].split('.')[0].split('_')[1]
                if protein=='Spike':
                    protein = 'S'

                for pocket,residues in consensus_pockets[protnow].items():
                    for res in residues:
                        if protein==protnow and res in binding_residues and ligand in consensus_pocket_ligands[protnow][pocket] and ligand not in consensus_pocket_reslig_pairs[protnow][pocket][res]:
                            consensus_pocket_reslig_pairs[protnow][pocket][res].append(ligand)
                            consensus_pocket_reslig_pairs[protnow][pocket][res].sort()

        file.close()
    return consensus_pocket_reslig_pairs

# for each viral protein pocket, make list of filtered ligands that bind (require ligand size>=8)
def filtered_pocket_ligands(directory,filenames,consensus_pockets,protnow,ligs_leaveout):
    consensus_pocket_ligands[protnow] = {}

    for pocket,residues in consensus_pockets[protnow].items():
        consensus_pocket_ligands[protnow][pocket] = []

    for fl in filenames:
        file = open(directory+fl,'r')
        line_list = file.readlines()

        for line in line_list:
            interaction = line.split()[0].split(':')[0]
            binding_residues = line.split()[-1].split(',')[0:-1]
            ligand = line.split()[0].split('.')[6]
            lig_size = line.split()[0].split('.')[7]

            # viral protein
            if line.split()[0].split('.')[0].split('_')[0]=='nCoV':
                protein = line.split()[0].split('.')[0].split('_')[1]
                if protein=='Spike':
                    protein = 'S'

                for pocket,residues in consensus_pockets[protnow].items():
                    if protein==protnow and (set(binding_residues) & set(residues)) and ligand not in ligs_leaveout[protnow]:
                        if len(ligand)<4:
                            if float(lig_size)>=8:
                                if ligand not in consensus_pocket_ligands[protnow][pocket] and ligand in smiles_dict and ligand not in bad_smiles:
                                    consensus_pocket_ligands[protnow][pocket].append(ligand)
                                    consensus_pocket_ligands[protnow][pocket].sort()
        file.close()
    return consensus_pocket_ligands

# calculate normalized Tanimoto distance between all pairs of ligands
def calc_Tanimoto_dist_norm(fp_radius,nBits,chemtax_dict):
    Tdist_dict = {}
    Tdistnorm_dict = {}
    Tdistlist = []
    for i1 in range(0,len(pocket_ligs)):
        lig1 = pocket_ligs[i1]
        if lig1 not in bad_smiles:
            m1 = Chem.MolFromSmiles(smiles_dict[lig1]['smiles_cactus'])
            fp1 = AllChem.GetMorganFingerprintAsBitVect(m1,fp_radius,nBits)
            for i2 in range(i1+1,len(pocket_ligs)):
                lig2 = pocket_ligs[i2]
                if lig2 not in bad_smiles:
                    m2 = Chem.MolFromSmiles(smiles_dict[lig2]['smiles_cactus'])
                    fp2 = AllChem.GetMorganFingerprintAsBitVect(m2,fp_radius,nBits)
                    Tsim = DataStructs.FingerprintSimilarity(fp1,fp2)
                    Tdist = 1-Tsim
                    Tdistlist.append(Tdist)
                    Tdist_dict[(lig1,lig2)] = Tdist

                    if lig1 in chemtax_dict and lig2 in chemtax_dict:
                        if Tdist==0.0 and chemtax_dict[lig1]['kingdom']!='' and chemtax_dict[lig2]['kingdom']=='':
                            chemtax_dict[lig2] = chemtax_dict[lig1]
                        elif Tdist==0.0 and chemtax_dict[lig1]['kingdom']=='' and chemtax_dict[lig2]['kingdom']!='':
                            chemtax_dict[lig1] = chemtax_dict[lig2]

    Tdistavg = sum(Tdistlist)/float(len(Tdistlist))
    Tdiststd = statistics.stdev(Tdistlist)

    print(Tdistavg,Tdiststd,min(Tdistlist))

    for ligpair in Tdist_dict.keys():
        Tdistnorm_dict[ligpair] = ((Tdist_dict[ligpair]-Tdistavg)/Tdiststd)+15

    return Tdistnorm_dict,chemtax_dict,Tdistavg,Tdiststd

# calculate Tanimoto distance for pairs of ligands in each pocket
def get_Tanimoto_dist(Tdistlist_dict,consensus_pocket_ligands,protnow,Tdistnorm_dict):
    Tdistlist_dict[protnow] = {}
    for pocket,ligands in consensus_pocket_ligands[protnow].items():
        Tdistlist = []
        for i1 in range(0,len(ligands)):
            lig1 = ligands[i1]
            for i2 in range(i1+1,len(ligands)):
                lig2 = ligands[i2]
                if (lig1,lig2) in Tdistnorm_dict:
                    Tdist = Tdistnorm_dict[(lig1,lig2)]
                    Tdistlist.append(Tdist)
                elif (lig2,lig1) in Tdistnorm_dict:
                    Tdist = Tdistnorm_dict[(lig2,lig1)]
                    Tdistlist.append(Tdist)
        Tdistlist_dict[protnow][pocket] = Tdistlist
    return Tdistlist_dict

# calculate normalized chemical taxonomy distance between all pairs of ligands
def calc_chemtax_dist_norm():
    Cdist_dict = {}
    Cdistnorm_dict = {}
    Cdistlist = []
    for i1 in range(0,len(pocket_ligs)):
        lig1 = pocket_ligs[i1]
        if lig1 not in bad_smiles:
            for i2 in range(i1+1,len(pocket_ligs)):
                lig2 = pocket_ligs[i2]
                Csim = 0
                if lig2 not in bad_smiles:
                    if chemtax_dict[lig1]['kingdom']!='' and chemtax_dict[lig2]['kingdom']!='':
                        for level in ['kingdom','superclass','class','subclass']:
                            if chemtax_dict[lig1][level]==chemtax_dict[lig2][level]:
                                Csim = Csim + 1
                        Csim = float(Csim)/float(4)
                        Cdist = 1-Csim
                        Cdistlist.append(Cdist)
                        Cdist_dict[(lig1,lig2)] = Cdist

    Cdistavg = sum(Cdistlist)/float(len(Cdistlist))
    Cdiststd = statistics.stdev(Cdistlist)

    print(Cdistavg,Cdiststd,min(Cdistlist))

    for ligpair in Cdist_dict.keys():
        Cdistnorm_dict[ligpair] = ((Cdist_dict[ligpair]-Cdistavg)/Cdiststd)+15

    return Cdistnorm_dict,Cdistavg,Cdiststd

# calculate distance based on chemical taxonomy for pairs of ligands in each pocket
def get_chemtax_dist(Cdistlist_dict,consensus_pocket_ligands,protnow,Cdistnorm_dict):
    Cdistlist_dict[protnow] = {}
    for pocket,ligands in consensus_pocket_ligands[protnow].items():
        Cdistlist = []
        for i1 in range(0,len(ligands)):
            lig1 = ligands[i1]
            for i2 in range(i1+1,len(ligands)):
                lig2 = ligands[i2]
                if (lig1,lig2) in Cdistnorm_dict:
                    Cdist = Cdistnorm_dict[(lig1,lig2)]
                    Cdistlist.append(Cdist)
                elif (lig2,lig1) in Cdistnorm_dict:
                    Cdist = Cdistnorm_dict[(lig2,lig1)]
                    Cdistlist.append(Cdist)
                else:
                    Cdistlist.append('NA')
        Cdistlist_dict[protnow][pocket] = Cdistlist
    return Cdistlist_dict

# calculate normalized word context distance between all pairs of ligands
def calc_ligname_dist_norm(ligname_dist_dict):
    LNdistnorm_dict = {}
    LNdist_dict = {}
    LNdistlist = []
    for i1 in range(0,len(pocket_ligs)):
        lig1 = pocket_ligs[i1]
        if lig1 not in bad_smiles:
            for i2 in range(i1+1,len(pocket_ligs)):
                lig2 = pocket_ligs[i2]
                if lig2 not in bad_smiles:
                    if (lig1,lig2) in ligname_dist_dict or (lig2,lig1) in ligname_dist_dict:
                        if (lig1,lig2) in ligname_dist_dict:
                            LNdist = ligname_dist_dict[(lig1,lig2)]
                        elif (lig2,lig1) in ligname_dist_dict:
                            LNdist = ligname_dist_dict[(lig2,lig1)]
                        LNdistlist.append(LNdist)
                        LNdist_dict[(lig1,lig2)] = LNdist

    LNdistavg = sum(LNdistlist)/float(len(LNdistlist))
    LNdiststd = statistics.stdev(LNdistlist)

    print(LNdistavg,LNdiststd,min(LNdistlist))

    for ligpair in LNdist_dict.keys():
        LNdistnorm_dict[ligpair] = ((LNdist_dict[ligpair]-LNdistavg)/LNdiststd)+15

    return LNdistnorm_dict,LNdistavg,LNdiststd

# make distance matrix from ligand name distance dictionary
def get_ligname_dist(LNdistlist_dict,consensus_pocket_ligands,protnow,LNdistnorm_dict):
    LNdistlist_dict[protnow] = {}
    for pocket,ligands in consensus_pocket_ligands[protnow].items():
        LNdistlist = []
        for i1 in range(0,len(ligands)):
            lig1 = ligands[i1]
            for i2 in range(i1+1,len(ligands)):
                lig2 = ligands[i2]
                if (lig1,lig2) in LNdistnorm_dict:
                    LNdist = LNdistnorm_dict[(lig1,lig2)]
                    LNdistlist.append(LNdist)
                elif (lig2,lig1) in LNdistnorm_dict:
                    LNdist = LNdistnorm_dict[(lig2,lig1)]
                    LNdistlist.append(LNdist)
                else:
                    LNdistlist.append('NA')
        LNdistlist_dict[protnow][pocket] = LNdistlist
    return LNdistlist_dict

# take weighted average of distance matrices
def weighted_dist(protnow,consensus_pocket_ligands,Tdistlist_dict,Cdistlist_dict,LNdistlist_dict,Wdistlist_dict):
    Wdistlist_dict[protnow] = {}
    for pocket,ligands in consensus_pocket_ligands[protnow].items():
        Wdistlist_dict[protnow][pocket] = []
        if len(Tdistlist_dict[protnow][pocket])!=len(Cdistlist_dict[protnow][pocket]):
            print('distance problem')
        if len(Cdistlist_dict[protnow][pocket])!=len(LNdistlist_dict[protnow][pocket]):
            print('distance problem')
        for ind in range(0,len(Tdistlist_dict[protnow][pocket])):
            Tdist = Tdistlist_dict[protnow][pocket][ind]
            Cdist = Cdistlist_dict[protnow][pocket][ind]
            LNdist = LNdistlist_dict[protnow][pocket][ind]
            if Cdist!='NA' and LNdist!='NA':
                Wdist = (float(1)/float(3))*Tdist + (float(1)/float(3))*Cdist + (float(1)/float(3))*LNdist
            elif Cdist!='NA' and LNdist=='NA':
                Wdist = 0.5*Tdist + 0.5*Cdist
            elif Cdist=='NA' and LNdist!='NA':
                Wdist = 0.5*Tdist + 0.5*LNdist
            elif Cdist=='NA' and LNdist=='NA':
                Wdist = Tdist
            else:
                print(ind,'no Wdist assignment')
            Wdistlist_dict[protnow][pocket].append(Wdist)
    return Wdistlist_dict

# get distance matrix from distance list
def get_dist_matrix(Wdistlist_dict,Wdistmat_dict,consensus_pocket_ligands):
    Wdistmat_dict[protnow] = {}
    for pocket,ligands in consensus_pocket_ligands[protnow].items():
        Wdistmat = -1*np.ones((len(ligands),len(ligands)))
        ind = 0
        for i1 in range(0,len(ligands)):
            for i2 in range(i1,len(ligands)):
                if i1==i2:
                    Wdistmat[i1,i2] = 0
                else:
                    Wdistmat[i1,i2] = Wdistlist_dict[protnow][pocket][ind]
                    Wdistmat[i2,i1] = Wdistlist_dict[protnow][pocket][ind]
                    ind = ind + 1
        Wdistmat_dict[protnow][pocket] = Wdistmat
    return Wdistmat_dict

# cluster ligands in each pocket using DBSCAN
def dbscan_cluster_pocket_ligands(cluster_dict,protnow,consensus_pocket_ligands,Wdistmat_dict):
    cluster_dict[protnow] = {}
    for pocket,ligands in consensus_pocket_ligands[protnow].items():
        best_params = {}
        if len(ligands)>2:
            best_params['sscore'] = -1
            eps_vec = np.linspace(1,15,14*20)
            minsamp_vec = np.arange(3,11)
            for eps_val in eps_vec:
                for minsamp_val in minsamp_vec:
                    db = DBSCAN(eps=eps_val, min_samples=minsamp_val, metric='precomputed').fit(Wdistmat_dict[protnow][pocket])
                    labels = db.labels_
                    n_clusters = len(set(labels)) - (1 if -1 in labels else 0) # Number of clusters in labels, ignoring noise if present
                    n_noise = list(labels).count(-1)

                    if len(ligands)>n_clusters and n_clusters>=2:
                        ss = silhouette_score(Wdistmat_dict[protnow][pocket],labels,metric='precomputed')
                        if ss > best_params['sscore']:
                            best_params = {'eps': eps_val, 'min_samples': minsamp_val, 'sscore': ss, 'n_clusters': n_clusters, 'n_noise': n_noise}

            if 'eps' in best_params.keys():
                db = DBSCAN(eps=best_params['eps'], min_samples=best_params['min_samples'], metric='precomputed').fit(Wdistmat_dict[protnow][pocket])
                labels = db.labels_
                n_clusters = len(set(labels)) - (1 if -1 in labels else 0) # Number of clusters in labels, ignoring noise if present
                n_noise = list(labels).count(-1)

                labels_array = np.empty((len(consensus_pocket_ligands[protnow][pocket]),1),dtype=np.int64)
                for k,ligand in enumerate(consensus_pocket_ligands[protnow][pocket]):
                    labels_array[k] = np.empty((1,),dtype=np.int64)
                    labels_array[k][0] = np.int64(labels[k])

                cluster_dict[protnow][pocket] = labels_array

            else:
                labels_array = np.empty((len(consensus_pocket_ligands[protnow][pocket]),1),dtype=np.int64)
                for k,ligand in enumerate(consensus_pocket_ligands[protnow][pocket]):
                    labels_array[k] = np.empty((1,),dtype=np.int64)
                    labels_array[k][0] = np.int64(k)

                cluster_dict[protnow][pocket] = labels_array

        elif len(ligands)==2:
            cluster_dict[protnow][pocket] = [[0], [0]]

    return cluster_dict

# silhouette plot
def silhouette(Wdistmat_dict,cluster_dict,consensus_pocket_ligands):
    for pocket, clusters in cluster_dict[protnow].items():
        clusout=[(x,clusters[k][0]) for k,x in enumerate(consensus_pocket_ligands[protnow][pocket])]
        clustall=[]
        for k in range(max([x[1] for x in clusout])+1):
            clustall.append([x[0] for x in clusout if x[1]==k])
        n_clusters=len(clustall)
        dist_matrix = Wdistmat_dict[protnow][pocket]

        try:
            cluster_labels = np.empty((len(consensus_pocket_ligands[protnow][pocket]),))
            max_clust_ind = np.amax(cluster_dict[protnow][pocket])
            for k in range(0,len(cluster_dict[protnow][pocket])):
                if cluster_dict[protnow][pocket][k][0]==-1:
                    cluster_labels[k] = max_clust_ind+1
                    max_clust_ind = max_clust_ind+1
                else:
                    cluster_labels[k] = cluster_dict[protnow][pocket][k][0]

            # Create a subplot with 1 row and 1 column
            fig = plt.figure()
            fig.set_size_inches(9, 6)
            ax=fig.add_subplot(111)

            # The (n_clusters+1)*10 is for inserting blank space between silhouette plots of individual clusters, to demarcate them clearly.
            ax.set_ylim([0, len(cluster_labels) + (n_clusters + 1) * 10])

            # The silhouette_score gives the average value for all the samples.
            silhouette_avg = silhouette_score(dist_matrix, cluster_labels, metric="precomputed", sample_size=None)
            print("There are ",n_clusters," clusters and the average silhouette_score is : ",silhouette_avg)

            # Compute the silhouette scores for each sample
            sample_silhouette_values = silhouette_samples(dist_matrix, cluster_labels, metric="precomputed")

            y_lower = 10
            for i in range(n_clusters):
                # Aggregate the silhouette scores for samples belonging to cluster i, and sort them
                ith_cluster_silhouette_values = sample_silhouette_values[cluster_labels == i]

                ith_cluster_silhouette_values.sort()

                size_cluster_i = ith_cluster_silhouette_values.shape[0]
                y_upper = y_lower + size_cluster_i

                color = cm.nipy_spectral(float(i)/n_clusters)
                ax.fill_betweenx(np.arange(y_lower, y_upper),0, ith_cluster_silhouette_values,facecolor=color, edgecolor=color, alpha=0.7)

                # Compute the new y_lower for next plot
                y_lower = y_upper + 10  # 10 for the 0 samples

            plt.title(protnow+', Pocket '+pocket,fontsize=16)
            plt.xlabel("Silhouette coefficient",fontsize=16)
            plt.ylabel("Cluster label",fontsize=16)

            # vertical line for average silhouette score of all the values
            ax.axvline(x=silhouette_avg, color="red", linestyle="--")

            ax.set_yticks([])  # Clear the yaxis labels / ticks
            plt.xlim((-0.1,0.6))
            ax.set_xticks([-0.1, 0, 0.1, 0.2, 0.3, 0.4, 0.5, 0.6])
            plt.xticks(fontsize=14)

            plt.show()
            #plt.savefig('figures/silhouette_plot_'+protnow+'_'+pocket+'.png')

        except:
            if len(cluster_dict[protnow][pocket])==n_clusters:
                print('All ligands clustered separately')
            elif n_clusters==1:
                print('All ligands clustered together')

    return

# find max common substructure for ligand cluster in each pocket
def pocket_mcs(cluster_dict,consensus_pocket_ligands):
    os.system('rm images/CCC-15-10-'+gdccut+'-4-0-ligs-8-current-resall/'+protnow+'_pocket*')
    with open('ligand-cluster-key-CCC-15-10-'+gdccut+'-4-0-ligs-8-current-resall-'+protnow+'.csv','w') as f:
        writeCSV = csv.DictWriter(f,fieldnames=['Pocket','Cluster Index','Ligands in Cluster'])
        writeCSV.writeheader()
        for pocket, clusters in cluster_dict[protnow].items():
            clusout=[(x,clusters[k][0]) for k,x in enumerate(consensus_pocket_ligands[protnow][pocket])]
            clustall=[]
            for k in range(max([x[1] for x in clusout])+1):
                clustall.append([x[0] for x in clusout if x[1]==k])
            n_clusters=len(clustall)
            for cind,clust in enumerate(clustall,1):
                newrow = {'Pocket': pocket, 'Cluster Index': cind, 'Ligands in Cluster': clust}
                writeCSV.writerow(newrow)

                subclasses = []
                for lig in clust:
                    if chemtax_dict[lig]['subclass'] not in subclasses:
                        subclasses.append(chemtax_dict[lig]['subclass'])
                if len(clust)>1:
                    print('Pocket',pocket,'\t','Cluster',cind,'\t','# Ligands',len(clust),'\t',\
                          'Class',chemtax_dict[clust[0]]['class'],'\t','# Subclasses',len(subclasses))

                if len(clust)>=3:
                    molecules = []
                    for lig in clust:
                        if lig not in bad_smiles:
                            molecules.append(Chem.MolFromSmiles(smiles_dict[lig]['smiles_cactus']))
                    if len(molecules)>1:
                        mcs = Chem.rdFMCS.FindMCS(molecules)
                        if mcs.numAtoms>=4:
                            mcs_smiles = Chem.MolToSmiles(Chem.MolFromSmarts(mcs.smartsString))
                            mcs_mol = Chem.MolFromSmarts(mcs.smartsString)
                            mcs_coords = AllChem.Compute2DCoords(mcs_mol)
                            image_file = 'images/CCC-15-10-'+gdccut+'-4-0-ligs-8-current-resall/'\
                            +protnow+'_pocket'+pocket+'_cluster'+str(cind)+'_mcs.png'
                            Draw.MolToFile(mcs_mol,image_file)
    return

def fraction_cluster_contacts_heatmap(cluster_dict,consensus_pocket_ligands,consensus_pocket_reslig_pairs,protnow,fraction_ligand_contacts_matrix_dict):
    fraction_ligand_contacts_matrix_dict[protnow] = {}
    for pocket, clusters in cluster_dict[protnow].items():
        sorted_residues = sorted(consensus_pocket_reslig_pairs[protnow][pocket].keys(), key = lambda r: int(r[1:-2]))
        clusout=[(x,clusters[k][0]) for k,x in enumerate(consensus_pocket_ligands[protnow][pocket])]
        clustall=[]
        for k in range(max([x[1] for x in clusout])+1):
            clustall.append([x[0] for x in clusout if x[1]==k])
        n_clusters=len(clustall)
        flc_matrix = -1*np.ones((n_clusters,len(consensus_pocket_reslig_pairs[protnow][pocket].keys())))
        for cind,clust in enumerate(clustall,1):
            for res in sorted_residues:
                overlap = set(consensus_pocket_reslig_pairs[protnow][pocket][res]).intersection(set(clust))
                resind = sorted_residues.index(res)
                flc_matrix[cind-1,resind] = len(overlap)/float(len(clust))

        fraction_ligand_contacts_matrix_dict[protnow][pocket] = (flc_matrix, sorted_residues)

        ## heatmap using matplotlib (colorbar has same range for all plots)
        cb_viridis = cm.get_cmap('viridis', 100)
        plt.figure()
        plt.pcolor(np.arange(len(sorted_residues)), np.arange(n_clusters), flc_matrix, cmap=cb_viridis, vmin=0, vmax=1, shading='auto')
        plt.title(protnow+', Pocket '+str(pocket))
        plt.xlabel('Residues')
        plt.ylabel('Clusters')
        plt.xticks(ticks=np.arange(len(sorted_residues)), labels=sorted_residues, rotation=90)
        plt.yticks(ticks=list(np.arange(n_clusters)), labels=list(np.arange(1,n_clusters+1)))
        plt.colorbar(label='Fraction of Cluster Ligands in Contact')
        plt.show()

    return fraction_ligand_contacts_matrix_dict

# save files with PDB IDs for ligands in each cluster
def save_ligand_clusters(cluster_dict,consensus_pocket_ligands):
    with open('ligand-cluster-key-CCC-15-10-'+gdccut+'-4-0-ligs-8-current-resall-'+protnow+'.csv','w') as f:
        writeCSV = csv.DictWriter(f,fieldnames=['Pocket','Cluster Index','Ligands in Cluster'])
        writeCSV.writeheader()
        for pocket, clusters in cluster_dict[protnow].items():
            clusout=[(x,clusters[k][0]) for k,x in enumerate(consensus_pocket_ligands[protnow][pocket])]
            clustall=[]
            for k in range(max([x[1] for x in clusout])+1):
                clustall.append([x[0] for x in clusout if x[1]==k])
            n_clusters=len(clustall)
            for cind,clust in enumerate(clustall,1):
                newrow = {'Pocket': pocket, 'Cluster Index': cind, 'Ligands in Cluster': clust}
                writeCSV.writerow(newrow)

# compare experimentally screened positive compounds with ligand clusters
def compare_exp_pos_compounds(protnow,smiles_dict,cluster_dict,consensus_pocket_ligands):
    cluster_distance_pos = {}
    if protnow=='nsp5':
        with open('./fret_crys_test1.csv','r') as pos_smiles_file:
            readCSV = csv.DictReader(pos_smiles_file)
            for row in readCSV:
                exp_smiles = str(row['smiles'])
                exp_id = row['compound_id']
                m1 = Chem.MolFromSmiles(exp_smiles)
                fp1 = AllChem.GetMorganFingerprintAsBitVect(m1,fp_radius,nBits)
                image_file = 'images/experimental_compounds_positive/'+exp_id+'.png'
                Draw.MolToFile(m1,image_file)
                cluster_distance_pos[exp_id] = {}
                for pocket, clusters in cluster_dict[protnow].items():
                    cluster_distance_pos[exp_id][pocket] = {}
                    clusout=[(x,clusters[k][0]) for k,x in enumerate(consensus_pocket_ligands[protnow][pocket])]
                    clustall=[]
                    for k in range(max([x[1] for x in clusout])+1):
                        clustall.append([x[0] for x in clusout if x[1]==k])
                    n_clusters=len(clustall)
                    for cind,clust in enumerate(clustall,1):
                        Tdist_avg = 0
                        for lig in clust:
                            # calculate Tanimoto distance
                            m2 = Chem.MolFromSmiles(smiles_dict[lig]['smiles_cactus'])
                            fp2 = AllChem.GetMorganFingerprintAsBitVect(m2,fp_radius,nBits)
                            Tsim = DataStructs.FingerprintSimilarity(fp1,fp2)
                            Tdist = 1-Tsim
                            Tdist_avg = Tdist_avg + Tdist
                        Tdist_avg = Tdist_avg/float(len(clust))
                        cluster_distance_pos[exp_id][pocket][cind] = Tdist_avg

    return cluster_distance_pos

# compare experimentally screened negative compounds with ligand clusters
def compare_exp_neg_compounds(protnow,smiles_dict,cluster_dict,consensus_pocket_ligands):
    cluster_distance_neg = {}
    if protnow=='nsp5':
        with open('./all_postera_match_neg_10132021.csv','r') as neg_smiles_file:
            readCSV = csv.DictReader(neg_smiles_file)
            for row in readCSV:
                exp_smiles = str(row['SMILES'])
                exp_id = row['compound_id']
                m1 = Chem.MolFromSmiles(exp_smiles)
                try:
                    fp1 = AllChem.GetMorganFingerprintAsBitVect(m1,fp_radius,nBits)
                    image_file = 'images/experimental_compounds_negative/'+exp_id+'.png'
                    Draw.MolToFile(m1,image_file)
                    cluster_distance_neg[exp_id] = {}
                    for pocket, clusters in cluster_dict[protnow].items():
                        cluster_distance_neg[exp_id][pocket] = {}
                        clusout=[(x,clusters[k][0]) for k,x in enumerate(consensus_pocket_ligands[protnow][pocket])]
                        clustall=[]
                        for k in range(max([x[1] for x in clusout])+1):
                            clustall.append([x[0] for x in clusout if x[1]==k])
                        n_clusters=len(clustall)
                        for cind,clust in enumerate(clustall,1):
                            Tdist_avg = 0
                            for lig in clust:
                                # calculate Tanimoto distance
                                m2 = Chem.MolFromSmiles(smiles_dict[lig]['smiles_cactus'])
                                fp2 = AllChem.GetMorganFingerprintAsBitVect(m2,fp_radius,nBits)
                                Tsim = DataStructs.FingerprintSimilarity(fp1,fp2)
                                Tdist = 1-Tsim
                                Tdist_avg = Tdist_avg + Tdist
                            Tdist_avg = Tdist_avg/float(len(clust))
                            cluster_distance_neg[exp_id][pocket][cind] = Tdist_avg
                except:
                    print(exp_id)
                    pass

    return cluster_distance_neg

# find most similar ligand cluster for experimentally screened compounds
def find_closest_cluster(cluster_distance,protnow,cluster_dict,consensus_pocket_ligands):
    closest_cluster = {}
    closest_dist_list = {}
    if protnow=='nsp5':
        for pocket, clusters in cluster_dict[protnow].items():
            closest_dist_list[pocket] = {}
            clusout=[(x,clusters[k][0]) for k,x in enumerate(consensus_pocket_ligands[protnow][pocket])]
            clustall=[]
            for k in range(max([x[1] for x in clusout])+1):
                clustall.append([x[0] for x in clusout if x[1]==k])
            n_clusters=len(clustall)
            for cind,clust in enumerate(clustall,1):
                closest_dist_list[pocket][cind] = []

        for exp_id,itm in cluster_distance.items():
            closest_cluster[exp_id] = {}
            for pocket,cind in cluster_distance[exp_id].items():
                min_dist = 1
                for cind,Tdist in cluster_distance[exp_id][pocket].items():
                    if Tdist<min_dist:
                        min_dist = Tdist
                        closest_cluster[exp_id][pocket] = (cind,Tdist)

        for pocket, clusters in cluster_dict[protnow].items():
            clusout=[(x,clusters[k][0]) for k,x in enumerate(consensus_pocket_ligands[protnow][pocket])]
            clustall=[]
            for k in range(max([x[1] for x in clusout])+1):
                clustall.append([x[0] for x in clusout if x[1]==k])
            n_clusters=len(clustall)
            for cind,clust in enumerate(clustall,1):
                for exp_id,itm in closest_cluster.items():
                    for pckt,tup in closest_cluster[exp_id].items():
                        if str(pckt)==str(pocket) and str(cind)==str(closest_cluster[exp_id][pckt][0]):
                            closest_dist_list[pocket][cind].append(closest_cluster[exp_id][pckt][1])

    return closest_cluster, closest_dist_list

# for each nsp5 ligand cluster, make overlapping histograms of Tanimoto distances for positive and negative compounds closest to that cluster
def closest_cluster_hist(protnow,closest_dist_list_pos,closest_dist_list_neg,cluster_dict,consensus_pocket_ligands):
    if protnow=='nsp5':
        for pocket, clusters in cluster_dict[protnow].items():
            clusout=[(x,clusters[k][0]) for k,x in enumerate(consensus_pocket_ligands[protnow][pocket])]
            clustall=[]
            for k in range(max([x[1] for x in clusout])+1):
                clustall.append([x[0] for x in clusout if x[1]==k])
            n_clusters=len(clustall)
            for cind,clust in enumerate(clustall,1):
                print('ligand cluster index :',cind)
                plt.figure()

                if cind==1:
                    plt.figtext(0.14,0.8,'p-value = 0.702\nno difference')
                elif cind==2:
                    plt.figtext(0.14,0.8,'p-value = 4.87e-17\ndifferent')
                elif cind==4:
                    plt.figtext(0.14,0.8,'p-value = 1.83e-29\ndifferent')
                elif cind==5:
                    plt.figtext(0.14,0.8,'p-value = 0.0825\nno difference')
                elif cind==7:
                    plt.figtext(0.14,0.8,'p-value = 1.36e-32\ndifferent')
                elif cind==10:
                    plt.figtext(0.14,0.8,'p-value = 5.90e-30\ndifferent')

                if len(closest_dist_list_pos[pocket][cind])>0 and len(closest_dist_list_neg[pocket][cind])>0:
                    print('average distance positive',statistics.mean(closest_dist_list_pos[pocket][cind]))
                    print('average distance negative',statistics.mean(closest_dist_list_neg[pocket][cind]))
                    plt.hist(closest_dist_list_pos[pocket][cind],range=(0.5,1.0),bins=20,density=True,alpha=0.5,label='positive',color='red')
                    plt.hist(closest_dist_list_neg[pocket][cind],range=(0.5,1.0),bins=20,density=True,alpha=0.5,label='negative',color='blue')

                elif len(closest_dist_list_pos[pocket][cind])==0 and len(closest_dist_list_neg[pocket][cind])>0:
                    print('average distance negative',statistics.mean(closest_dist_list_neg[pocket][cind]))
                    plt.hist(closest_dist_list_neg[pocket][cind],range=(0.5,1.0),bins=20,density=True,alpha=0.5,label='negative',color='blue')

                elif len(closest_dist_list_pos[pocket][cind])>0 and len(closest_dist_list_neg[pocket][cind])==0:
                    print('average distance negative',statistics.mean(closest_dist_list_pos[pocket][cind]))
                    plt.hist(closest_dist_list_pos[pocket][cind],range=(0.5,1.0),bins=20,density=True,alpha=0.5,label='positive',color='red')

                plt.legend(loc='upper right')
                plt.title(protnow+', Pocket '+pocket+', Cluster '+str(cind),fontsize=16)
                plt.xlabel('Tanimoto distance',fontsize=16)
                plt.xlim(0.5,1.0)
                plt.ylim(0,30)
                plt.ylabel('Normalized count',fontsize=16)
                plt.xticks(fontsize=14)
                plt.yticks(fontsize=14)
                plt.show()
                #plt.savefig('figures/exp_compounds_distance_distr_pocket'+pocket+'_cluster'+str(cind)+'.png')

    return

# for each nsp5 ligand cluster, perform t-test for positive and negative compounds closest to that cluster
def closest_cluster_ttest(protnow,closest_dist_list_pos,closest_dist_list_neg,cluster_dict,consensus_pocket_ligands):
    if protnow=='nsp5':
        for pocket, clusters in cluster_dict[protnow].items():
            clusout=[(x,clusters[k][0]) for k,x in enumerate(consensus_pocket_ligands[protnow][pocket])]
            clustall=[]
            for k in range(max([x[1] for x in clusout])+1):
                clustall.append([x[0] for x in clusout if x[1]==k])
            n_clusters=len(clustall)
            for cind,clust in enumerate(clustall,1):
                ttest_output = scipy.stats.ttest_ind(closest_dist_list_pos[pocket][cind],closest_dist_list_neg[pocket][cind])
                pval = ttest_output[1]
                print(pocket,cind,ttest_output)
    return

# for each cluster, compile list of experimentally screened compounds that were found to be most similar to cluster
def list_closest_ligands(closest_cluster,cluster_dict):
    closest_ligands = {}
    for pocket, clusters in cluster_dict[protnow].items():
        closest_ligands[pocket] = {}
        clusout=[(x,clusters[k][0]) for k,x in enumerate(consensus_pocket_ligands[protnow][pocket])]
        clustall=[]
        for k in range(max([x[1] for x in clusout])+1):
            clustall.append([x[0] for x in clusout if x[1]==k])
        n_clusters=len(clustall)
        for cind,clust in enumerate(clustall,1):
            closest_ligands[pocket][cind] = []
            for exp_id in closest_cluster.keys():
                if closest_cluster[exp_id][pocket][0]==cind:
                    closest_ligands[pocket][cind].append(exp_id)

    return closest_ligands

prot_list_focus = ['S','nsp5','nsp12']

consensus_pockets = {}
all_consensus_residues = {}
consensus_pocket_ligands = {}
consensus_pocket_reslig_pairs = {}
fraction_lig_contacts = {}
fraction_ligand_contacts_matrix_dict = {}
Tdistlist_dict = {}
Cdistlist_dict = {}
LNdistlist_dict = {}
Wdistlist_dict = {}
Wdistmat_dict = {}
Tweight = 0.5
cluster_dict = {}
comout_dict = {}
fp_radius = 2
nBits = 1024
gdccut = '60'

ligs_leaveout = pickle.load(open('ligs_leaveout.p','rb'))
chemtax_dict = pickle.load(open('chemtax_dict.p', 'rb'))
ligname_dist_dict_notscaled = pickle.load(open('ligname_dist_dict_notscaled.p','rb'))

directory = 'cluster-output-ncov-residues-shortestpath-CCC-15-10-'+gdccut+'-4-0.ligs_8/date_current_resall/'
for protnow in prot_list_focus:
    consensus_pockets = pocket_residues(consensus_pockets,all_consensus_residues,directory,protnow)[0]
    all_consensus_residues = pocket_residues(consensus_pockets,all_consensus_residues,directory,protnow)[1]

directory = './'
filenames = ['CCC.confidence_centroid_contacts.15_10_'+gdccut+'_4_0.ligs_8.nCoV.current.resall']

pocket_ligs = []
for protnow in prot_list_focus:
    consensus_pocket_ligands = filtered_pocket_ligands(directory,filenames,consensus_pockets,protnow,ligs_leaveout)
    consensus_pocket_reslig_pairs = pocket_residue_ligand_pairs(directory,filenames,consensus_pockets,consensus_pocket_ligands,protnow)
    for pocket,ligands in consensus_pocket_ligands[protnow].items():
        for lig in ligands:
            if lig not in pocket_ligs:
                pocket_ligs.append(lig)

Tdist_output = calc_Tanimoto_dist_norm(fp_radius,nBits,chemtax_dict)
Tdistnorm_dict = Tdist_output[0]
pickle.dump(Tdistnorm_dict,open('normalized-Tanimoto-dist.p','wb'))
#Tdistnorm_dict = pickle.load(open('normalized-Tanimoto-dist.p','rb'))

# chemical taxonomy dictionary
chemtax_dict = Tdist_output[1]
pickle.dump(chemtax_dict,open('chemtax_dict_updated.p','wb'))
#chemtax_dict = pickle.load(open('chemtax_dict_updated.p','rb'))

Cdist_output = calc_chemtax_dist_norm()
Cdistnorm_dict = Cdist_output[0]
pickle.dump(Cdistnorm_dict,open('normalized-chemtax-dist.p','wb'))
#Cdistnorm_dict = pickle.load(open('normalized-chemtax-dist.p','rb'))

LNdist_output = calc_ligname_dist_norm(ligname_dist_dict_notscaled)
LNdistnorm_dict = LNdist_output[0]
pickle.dump(LNdistnorm_dict,open('normalized-wordvec-dist.p','wb'))
#LNdistnorm_dict = pickle.load(open('normalized-wordvec-dist.p','rb'))

for protnow in prot_list_focus:
    Tdistlist_dict = get_Tanimoto_dist(Tdistlist_dict,consensus_pocket_ligands,protnow,Tdistnorm_dict)
    Cdistlist_dict = get_chemtax_dist(Cdistlist_dict,consensus_pocket_ligands,protnow,Cdistnorm_dict)
    LNdistlist_dict = get_ligname_dist(LNdistlist_dict,consensus_pocket_ligands,protnow,LNdistnorm_dict)

    Wdistlist_dict = weighted_dist(protnow,consensus_pocket_ligands,Tdistlist_dict,Cdistlist_dict,LNdistlist_dict,Wdistlist_dict)
    Wdistmat_dict = get_dist_matrix(Wdistlist_dict,Wdistmat_dict,consensus_pocket_ligands)

    cluster_dict = dbscan_cluster_pocket_ligands(cluster_dict,protnow,consensus_pocket_ligands,Wdistmat_dict)

    save_ligand_clusters(cluster_dict,consensus_pocket_ligands)

    fraction_ligand_contacts_matrix_dict = fraction_cluster_contacts_heatmap(cluster_dict,consensus_pocket_ligands,consensus_pocket_reslig_pairs,protnow,fraction_ligand_contacts_matrix_dict)

    silhouette(Wdistmat_dict,cluster_dict,consensus_pocket_ligands)

    pocket_mcs(cluster_dict,consensus_pocket_ligands)

    cluster_distance_pos = compare_exp_pos_compounds(protnow,smiles_dict,cluster_dict,consensus_pocket_ligands)
    cluster_distance_neg = compare_exp_neg_compounds(protnow,smiles_dict,cluster_dict,consensus_pocket_ligands)
    closest_cluster_pos = find_closest_cluster(cluster_distance_pos,protnow,cluster_dict,consensus_pocket_ligands)[0]
    closest_dist_list_pos = find_closest_cluster(cluster_distance_pos,protnow,cluster_dict,consensus_pocket_ligands)[1]
    closest_cluster_neg = find_closest_cluster(cluster_distance_neg,protnow,cluster_dict,consensus_pocket_ligands)[0]
    closest_dist_list_neg = find_closest_cluster(cluster_distance_neg,protnow,cluster_dict,consensus_pocket_ligands)[1]

    #print(closest_dist_list_pos)
    #print(closest_dist_list_neg)

    closest_cluster_hist(protnow,closest_dist_list_pos,closest_dist_list_neg,cluster_dict,consensus_pocket_ligands)
    closest_cluster_ttest(protnow,closest_dist_list_pos,closest_dist_list_neg,cluster_dict,consensus_pocket_ligands)

    #closest_ligands_pos = list_closest_ligands(closest_cluster_pos,cluster_dict)
    #closest_ligands_neg = list_closest_ligands(closest_cluster_neg,cluster_dict)
    #print(closest_ligands_pos)