# This file is part of QuTiP: Quantum Toolbox in Python.
#
#    Copyright (c) 2011 and later, Paul D. Nation and Robert J. Johansson.
#    All rights reserved.
#
#    Redistribution and use in source and binary forms, with or without
#    modification, are permitted provided that the following conditions are
#    met:
#
#    1. Redistributions of source code must retain the above copyright notice,
#       this list of conditions and the following disclaimer.
#
#    2. Redistributions in binary form must reproduce the above copyright
#       notice, this list of conditions and the following disclaimer in the
#       documentation and/or other materials provided with the distribution.
#
#    3. Neither the name of the QuTiP: Quantum Toolbox in Python nor the names
#       of its contributors may be used to endorse or promote products derived
#       from this software without specific prior written permission.
#
#    THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS
#    "AS IS" AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT
#    LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A
#    PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT
#    HOLDER OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL,
#    SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT
#    LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE,
#    DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY
#    THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT
#    (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
#    OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
###############################################################################
from collections.abc import Iterable
import warnings
from copy import deepcopy

import numpy as np
from scipy.interpolate import CubicSpline

from qutip.qobj import Qobj
from qutip.qobjevo import QobjEvo
from qutip.operators import identity
from qutip.qip.operations.gates import expand_operator, globalphase
from qutip.tensor import tensor
from qutip.mesolve import mesolve
from qutip.mcsolve import mcsolve
from qutip.qip.circuit import QubitCircuit
from qutip.qip.noise import (
    Noise, RelaxationNoise, DecoherenceNoise,
    ControlAmpNoise, RandomNoise, process_noise)
from qutip.qip.pulse import Pulse, Drift, _merge_qobjevo, _fill_coeff

__all__ = ['Processor']

class Processor(object):
    """
    A simulator of a quantum device based on the QuTiP solver
    :func:`qutip.mesolve`.  It is defined by the available driving Hamiltonian
    and the decoherence time for each component systems.  The processor can
    simulate the evolution under the given control pulses. Noisy evolution is
    supported by :class:`.Noise` and can be added to the processor.

    Parameters
    ----------
    N: int
        The number of component systems.

    t1: list or float, optional
        Characterize the decoherence of amplitude damping for
        each qubit. A list of size `N` or a float for all qubits.

    t2: list of float, optional
        Characterize the decoherence of dephasing for
        each qubit. A list of size `N` or a float for all qubits.

    dims: list, optional
        The dimension of each component system.
        Default value is a
        qubit system of ``dim=[2,2,2,...,2]``

    spline_kind: str, optional
        Type of the coefficient interpolation. Default is "step_func"
        Note that they have different requirement for the length of ``coeff``.

        - "step_func":
          The coefficient will be treated as a step function.  E.g.
          ``tlist=[0,1,2]`` and ``coeff=[3,2]``, means that the coefficient is
          3 in t=[0,1) and 2 in t=[2,3). It requires
          ``len(coeff)=len(tlist)-1`` or ``len(coeff)=len(tlist)``, but in the
          second case the last element of `coeff` has no effect.

        - "cubic": Use cubic interpolation for the coefficient. It requires
          ``len(coeff)=len(tlist)``

    Attributes
    ----------
    N: int
        The number of component systems.

    pulses: list of :class:`.Pulse`
        A list of control pulses of this device

    t1: float or list
        Characterize the decoherence of amplitude damping of
        each qubit.

    t2: float or list
        Characterize the decoherence of dephasing for
        each qubit.

    noise: :class:`.Noise`, optional
        A list of noise objects. They will be processed when creating the
        noisy :class:`qutip.QobjEvo` from the processor or run the simulation.

    drift: :class:`qutip.qip.pulse.Drift`
        A `Drift` object representing the drift Hamiltonians.

    dims: list
        The dimension of each component system.
        Default value is a
        qubit system of ``dim=[2,2,2,...,2]``

    spline_kind: str
        Type of the coefficient interpolation.
        See parameters of :class:`.Processor` for details.
    """
    def __init__(self, N, t1=None, t2=None,
                 dims=None, spline_kind="step_func"):
        self.N = N
        self.pulses = []
        self.t1 = t1
        self.t2 = t2
        self.noise = []
        self.drift = Drift()
        if dims is None:
            self.dims = [2] * N
        else:
            self.dims = dims
        self.pulse_mode = "discrete"
        self.spline_kind = spline_kind

    @property
    def num_qubits(self):
        return self.N

    @num_qubits.setter
    def num_qubits(self, value):
        self.N = value

    def add_drift(self, qobj, targets, cyclic_permutation=False):
        """
        Add a drift Hamiltonians. The drift Hamiltonians are intrinsic
        of the quantum system and cannot be controlled by external field.

        Parameters
        ----------
        qobj: :class:`qutip.Qobj`
            The drift Hamiltonian.
        targets: list
            The indices of the target qubits
            (or subquantum system of other dimensions).
        """
        if not isinstance(qobj, Qobj):
            raise TypeError("The drift Hamiltonian must be a qutip.Qobj.")
        if not qobj.isherm:
            raise ValueError("The drift Hamiltonian must be Hermitian.")

        num_qubits = len(qobj.dims[0])
        if targets is None:
            targets = list(range(num_qubits))
        if not isinstance(targets, list):
            targets = [targets]
        if cyclic_permutation:
            for i in range(self.N):
                temp_targets = [(t + i) % self.N for t in targets]
                self.drift.add_drift(qobj, temp_targets)
        else:
            self.drift.add_drift(qobj, targets)

    def add_control(self, qobj, targets=None, cyclic_permutation=False,
                    label=None):
        """
        Add a control Hamiltonian to the processor. It creates a new
        :class:`.Pulse`
        object for the device that is turned off
        (``tlist = None``, ``coeff = None``). To activate the pulse, one
        can set its `tlist` and `coeff`.

        Parameters
        ----------
        qobj: :class:`qutip.Qobj`
            The Hamiltonian for the control pulse..

        targets: list, optional
            The indices of the target qubits
            (or subquantum system of other dimensions).

        cyclic_permutation: bool, optional
            If true, the Hamiltonian will be expanded for
            all cyclic permutation of the target qubits.

        label: str, optional
            The label (name) of the pulse
        """
        # Check validity of ctrl
        if not isinstance(qobj, Qobj):
            raise TypeError("The control Hamiltonian must be a qutip.Qobj.")
        if not qobj.isherm:
            raise ValueError("The control Hamiltonian must be Hermitian.")

        num_qubits = len(qobj.dims[0])
        if targets is None:
            targets = list(range(num_qubits))
        if not isinstance(targets, list):
            targets = [targets]
        if cyclic_permutation:
            for i in range(self.N):
                temp_targets = [(t + i) % self.N for t in targets]
                if label is not None:
                    temp_label = label + "_" + str(temp_targets)
                temp_label = label
                self.pulses.append(
                    Pulse(qobj, temp_targets, spline_kind=self.spline_kind,
                          label=temp_label))
        else:
            self.pulses.append(
                Pulse(qobj, targets, spline_kind=self.spline_kind, label=label)
                )

    def find_pulse(self, pulse_name):
        if isinstance(pulse_name, str):
            try:
                return self.pulses[self.pulse_dict[pulse_name]]
            except (KeyError):
                raise KeyError(
                    "Pulse name {} undefined. "
                    "Please define it in the attribute "
                    "`pulse_dict`.".format(pulse_name))
        elif isinstance(pulse_name, int):
            return self.pulses[pulse_name]
        else:
            raise TypeError(
                "pulse_name is either a string or an integer, not "
                "{}".format(type(pulse_name))
                )

    @property
    def ctrls(self):
        """
        A list of Hamiltonians of all pulses.
        """
        result = []
        for pulse in self.pulses:
            result.append(pulse.get_ideal_qobj(self.dims))
        return result

    @property
    def coeffs(self):
        """
        A list of the coefficients for all control pulses.
        """
        if not self.pulses:
            return None
        coeffs_list = [pulse.coeff for pulse in self.pulses]
        return coeffs_list

    @coeffs.setter
    def coeffs(self, coeffs_list):
        for i, coeff in enumerate(coeffs_list):
            self.pulses[i].coeff = coeff

    @property
    def pulse_mode(self):
        if self.spline_kind == "step_func":
            return "discrete"
        elif self.spline_kind == "cubic":
            return "continuous"
        else:
            raise ValueError(
                "Saved spline_kind not understood.")

    @pulse_mode.setter
    def pulse_mode(self, mode):
        if mode == "discrete":
            spline_kind = "step_func"
        elif mode == "continuous":
            spline_kind = "cubic"
        else:
            raise ValueError(
                "Pulse mode must be either discrete or continuous.")

        self.spline_kind = spline_kind
        for pulse in self.pulses:
            pulse.spline_kind = spline_kind

    def get_full_tlist(self, tol=1.0e-10):
        """
        Return the full tlist of the ideal pulses.
        If different pulses have different time steps,
        it will collect all the time steps in a sorted array.

        Returns
        -------
        full_tlist: array-like 1d
            The full time sequence for the ideal evolution.
        """
        full_tlist = [pulse.tlist
                      for pulse in self.pulses if pulse.tlist is not None]
        if not full_tlist:
            return None
        full_tlist = np.unique(np.sort(np.hstack(full_tlist)))
        # account for inaccuracy in float-point number
        full_tlist = np.concatenate(
            (full_tlist[:1], full_tlist[1:][np.diff(full_tlist) > tol]))
        return full_tlist

    def get_full_coeffs(self, full_tlist=None):
        """
        Return the full coefficients in a 2d matrix form.
        Each row corresponds to one pulse. If the `tlist` are
        different for different pulses, the length of each row
        will be same as the `full_tlist` (see method
        `get_full_tlist`). Interpolation is used for
        adding the missing coefficient according to `spline_kind`.

        Returns
        -------
        coeffs: array-like 2d
            The coefficients for all ideal pulses.
        """
        # TODO add tests
        self._is_pulses_valid()
        if not self.pulses:
            return np.array((0, 0), dtype=float)
        if full_tlist is None:
            full_tlist = self.get_full_tlist()
        coeffs_list = []
        for pulse in self.pulses:
            if pulse.tlist is None and pulse.coeff is None:
                coeffs_list.append(np.zeros(len(full_tlist)))
                continue
            if not isinstance(pulse.coeff, (bool, np.ndarray)):
                raise ValueError(
                    "get_full_coeffs only works for "
                    "NumPy array or bool coeff.")
            if isinstance(pulse.coeff, bool):
                if pulse.coeff:
                    coeffs_list.append(np.ones(len(full_tlist)))
                else:
                    coeffs_list.append(np.zeros(len(full_tlist)))
                continue
            if self.spline_kind == "step_func":
                arg = {"_step_func_coeff": True}
                coeffs_list.append(
                    _fill_coeff(pulse.coeff, pulse.tlist, full_tlist, arg))
            elif self.spline_kind == "cubic":
                coeffs_list.append(
                    _fill_coeff(pulse.coeff, pulse.tlist, full_tlist, {}))
            else:
                raise ValueError("Unknown spline kind.")
        return np.array(coeffs_list)

    def set_all_tlist(self, tlist):
        """
        Set the same `tlist` for all the pulses.

        Parameters
        ----------
        tlist: array-like, optional
            A list of time at which the time-dependent coefficients are
            applied. See :class:`.Pulse` for detailed information`
        """
        if isinstance(tlist, list) and len(tlist) == len(self.pulses):
            for i, pulse in enumerate(self.pulses):
                pulse.tlist = tlist[i]
        else:
            for pulse in self.pulses:
                pulse.tlist = tlist

    def add_pulse(self, pulse):
        """
        Add a new pulse to the device.

        Parameters
        ----------
        pulse: :class:`.Pulse`
            `Pulse` object to be added.
        """
        if isinstance(pulse, Pulse):
            if pulse.spline_kind is None:
                pulse.spline_kind = self.spline_kind
            self.pulses.append(pulse)
        else:
            raise ValueError("Invalid input, pulse must be a Pulse object")

    def remove_pulse(self, indices=None, label=None):
        """
        Remove the control pulse with given indices.

        Parameters
        ----------
        indices: int or list of int
            The indices of the control Hamiltonians to be removed.
        label: str
            The label of the pulse
        """
        if indices is not None:
            if not isinstance(indices, Iterable):
                indices = [indices]
            indices.sort(reverse=True)
            for ind in indices:
                del self.pulses[ind]
        else:
            for ind, pulse in enumerate(self.pulses):
                if pulse.label == label:
                    del self.pulses[ind]

    def _is_pulses_valid(self):
        """
        Check if the pulses are in the correct shape.

        Returns: bool
            If they are valid or not
        """
        for i, pulse in enumerate(self.pulses):
            if pulse.coeff is None or isinstance(pulse.coeff, bool):
                # constant pulse
                continue
            if pulse.tlist is None:
                raise ValueError(
                    "Pulse id={} is invalid. "
                    "Please define a tlist for the pulse.".format(i))
            if pulse.tlist is not None and pulse.coeff is None:
                raise ValueError(
                    "Pulse id={} is invalid. "
                    "Please define a coeff for the pulse.".format(i))
            coeff_len = len(pulse.coeff)
            tlist_len = len(pulse.tlist)
            if pulse.spline_kind == "step_func":
                if coeff_len == tlist_len-1 or coeff_len == tlist_len:
                    pass
                else:
                    raise ValueError(
                        "The length of tlist and coeff of the pulse "
                        "labelled {} is invalid. "
                        "It's either len(tlist)=len(coeff) or "
                        "len(tlist)-1=len(coeff) for coefficients "
                        "as step function".format(i))
            else:
                if coeff_len == tlist_len:
                    pass
                else:
                    raise ValueError(
                        "The length of tlist and coeff of the pulse "
                        "labelled {} is invalid. "
                        "It should be either len(tlist)=len(coeff)".format(i))
        return True

    def add_noise(self, noise):
        """
        Add a noise object to the processor

        Parameters
        ----------
        noise: :class:`.Noise`
            The noise object defined outside the processor
        """
        if isinstance(noise, Noise):
            self.noise.append(noise)
        else:
            raise TypeError("Input is not a Noise object.")

    def save_coeff(self, file_name, inctime=True):
        """
        Save a file with the control amplitudes in each timeslot.

        Parameters
        ----------
        file_name: string
            Name of the file.

        inctime: bool, optional
            True if the time list should be included in the first column.
        """
        self._is_pulses_valid()
        coeffs = np.array(self.get_full_coeffs())
        if inctime:
            shp = coeffs.T.shape
            data = np.empty((shp[0], shp[1] + 1), dtype=np.float64)
            data[:, 0] = self.get_full_tlist()
            data[:, 1:] = coeffs.T
        else:
            data = coeffs.T

        np.savetxt(file_name, data, delimiter='\t', fmt='%1.16f')

    def read_coeff(self, file_name, inctime=True):
        """
        Read the control amplitudes matrix and time list
        saved in the file by `save_amp`.

        Parameters
        ----------
        file_name: string
            Name of the file.

        inctime: bool, optional
            True if the time list in included in the first column.

        Returns
        -------
        tlist: array_like
            The time list read from the file.

        coeffs: array_like
            The pulse matrix read from the file.
        """
        data = np.loadtxt(file_name, delimiter='\t')
        if not inctime:
            self.coeffs = data.T
            return self.coeffs
        else:
            tlist = data[:, 0]
            self.set_all_tlist(tlist)
            self.coeffs = data[:, 1:].T
            return self.get_full_tlist, self.coeffs

    def get_noisy_pulses(self, device_noise=False, drift=False):
        """
        It takes the pulses defined in the `Processor` and
        adds noise according to `Processor.noise`. It does not modify the
        pulses saved in `Processor.pulses` but returns a new list.
        The length of the new list of noisy pulses might be longer
        because of drift Hamiltonian and device noise. They will be
        added to the end of the pulses list.

        Parameters
        ----------
        device_noise: bool, optional
            If true, include pulse independent noise such as single qubit
            Relaxation. Default is False.
        drift: bool, optional
            If true, include drift Hamiltonians. Default is False.

        Returns
        -------
        noisy_pulses: list of :class:`.Pulse`
            A list of noisy pulses.
        """
        pulses = deepcopy(self.pulses)
        noisy_pulses = process_noise(
            pulses, self.noise, self.dims, t1=self.t1, t2=self.t2,
            device_noise=device_noise)
        if drift:
            noisy_pulses += [self.drift]
        return noisy_pulses

    def get_qobjevo(self, args=None, noisy=False):
        """
        Create a :class:`qutip.QobjEvo` representation of the evolution.
        It calls the method `get_noisy_pulses` and create the `QobjEvo`
        from it.

        Parameters
        ----------
        args: dict, optional
            Arguments for :class:`qutip.QobjEvo`
        noisy: bool, optional
            If noise are included. Default is False.

        Returns
        -------
        qobjevo: :class:`qutip.QobjEvo`
            The :class:`qutip.QobjEvo` representation of the unitary evolution.
        c_ops: list of :class:`qutip.QobjEvo`
            A list of lindblad operators is also returned. if ``noisy==Flase``,
            it is always an empty list.
        """
        # TODO test it for non array-like coeff
        # check validity
        self._is_pulses_valid()

        if args is None:
            args = {}
        else:
            args = args
        # set step function

        if not noisy:
            dynamics = self.pulses
        else:
            dynamics = self.get_noisy_pulses(
                device_noise=True, drift=True)

        qu_list = []
        c_ops = []
        for pulse in dynamics:
            if noisy:
                qu, new_c_ops = pulse.get_noisy_qobjevo(dims=self.dims)
                c_ops += new_c_ops
            else:
                qu = pulse.get_ideal_qobjevo(dims=self.dims)
            qu_list.append(qu)

        final_qu = _merge_qobjevo(qu_list)
        final_qu.args.update(args)

        # bring all c_ops to the same tlist, won't need it in QuTiP 5
        full_tlist = self.get_full_tlist()
        temp = []
        for c_op in c_ops:
            temp.append(_merge_qobjevo([c_op], full_tlist))
        c_ops = temp

        if noisy:
            return final_qu, c_ops
        else:
            return final_qu, []

    def run_analytically(self, init_state=None, qc=None):
        """
        Simulate the state evolution under the given `qutip.QubitCircuit`
        with matrice exponentiation. It will calculate the propagator
        with matrix exponentiation and return a list of :class:`qutip.Qobj`.
        This method won't include noise or collpase.

        Parameters
        ----------
        qc: :class:`.QubitCircuit`, optional
            Takes the quantum circuit to be implemented. If not given, use
            the quantum circuit saved in the processor by ``load_circuit``.

        init_state: :class:`qutip.Qobj`, optional
            The initial state of the qubits in the register.

        Returns
        -------
        U_list: list
            A list of propagators obtained for the physical implementation.
        """
        if init_state is not None:
            U_list = [init_state]
        else:
            U_list = []
        tlist = self.get_full_tlist()
        coeffs = self.get_full_coeffs()

        # Compute drift Hamiltonians
        H_drift = 0
        for drift_ham in self.drift.drift_hamiltonians:
            H_drift += drift_ham.get_qobj(self.dims)

        # Compute control Hamiltonians
        for n in range(len(tlist)-1):
            H = H_drift + sum(
                [coeffs[m, n] * self.ctrls[m]
                    for m in range(len(self.ctrls))])
            dt = tlist[n + 1] - tlist[n]
            U = (-1j * H * dt).expm()
            U = self.eliminate_auxillary_modes(U)
            U_list.append(U)

        try:  # correct_global_phase are defined for ModelProcessor
            if self.correct_global_phase and self.global_phase != 0:
                U_list.append(globalphase(
                    self.global_phase, N=self.num_qubits)
                )
        except AttributeError:
            pass

        return U_list

    def run(self, qc=None):
        """
        Calculate the propagator of the evolution by matrix exponentiation.
        This method won't include noise or collpase.

        Parameters
        ----------
        qc: :class:`.QubitCircuit`, optional
            Takes the quantum circuit to be implemented. If not given, use
            the quantum circuit saved in the processor by `load_circuit`.

        Returns
        -------
        U_list: list
            The propagator matrix obtained from the physical implementation.
        """
        if qc:
            self.load_circuit(qc)
        return self.run_analytically(qc=qc, init_state=None)

    def run_state(self, init_state=None, analytical=False, states=None,
                  noisy=True, solver="mesolve", **kwargs):
        """
        If `analytical` is False, use :func:`qutip.mesolve` to
        calculate the time of the state evolution
        and return the result. Other arguments of mesolve can be
        given as keyword arguments.

        If `analytical` is True, calculate the propagator
        with matrix exponentiation and return a list of matrices.
        Noise will be neglected in this option.

        Parameters
        ----------
        init_state: Qobj
            Initial density matrix or state vector (ket).

        analytical: bool
            If True, calculate the evolution with matrices exponentiation.

        states: :class:`qutip.Qobj`, optional
            Old API, same as init_state.

        solver: str
            "mesolve" or "mcsolve"

        **kwargs
            Keyword arguments for the qutip solver.

        Returns
        -------
        evo_result: :class:`qutip.Result`
            If ``analytical`` is False,  an instance of the class
            :class:`qutip.Result` will be returned.

            If ``analytical`` is True, a list of matrices representation
            is returned.
        """
        if states is not None:
            warnings.warn(
                "states will be deprecated and replaced by init_state",
                DeprecationWarning)
        if init_state is None and states is None:
            raise ValueError("Qubit state not defined.")
        elif init_state is None:
            # just to keep the old parameters `states`,
            # it is replaced by init_state
            init_state = states
        if analytical:
            if kwargs or self.noise:
                raise warnings.warn(
                    "Analytical matrices exponentiation"
                    "does not process noise or"
                    "any keyword arguments.")
            return self.run_analytically(init_state=init_state)

        # kwargs can not contain H or tlist
        if "H" in kwargs or "tlist" in kwargs:
            raise ValueError(
                "`H` and `tlist` are already specified by the processor "
                "and can not be given as a keyword argument")

        # construct qobjevo for unitary evolution
        if "args" in kwargs:
            noisy_qobjevo, sys_c_ops = self.get_qobjevo(
                    args=kwargs["args"], noisy=noisy)
        else:
            noisy_qobjevo, sys_c_ops = self.get_qobjevo(noisy=noisy)

        # add collpase operators into kwargs
        if "c_ops" in kwargs:
            if isinstance(kwargs["c_ops"], (Qobj, QobjEvo)):
                kwargs["c_ops"] += [kwargs["c_ops"]] + sys_c_ops
            else:
                kwargs["c_ops"] += sys_c_ops
        else:
            kwargs["c_ops"] = sys_c_ops

        # choose solver:
        if solver == "mesolve":
            evo_result = mesolve(
                H=noisy_qobjevo, rho0=init_state,
                tlist=noisy_qobjevo.tlist, **kwargs)
        elif solver == "mcsolve":
            evo_result = mcsolve(
                H=noisy_qobjevo, psi0=init_state,
                tlist=noisy_qobjevo.tlist, **kwargs)

        return evo_result

    def load_circuit(self, qc):
        """
        Translate an :class:`.QubitCircuit` to its
        corresponding Hamiltonians. (Defined in subclasses)
        """
        raise NotImplementedError("Use the function in the sub-class")

    def eliminate_auxillary_modes(self, U):
        """
        Eliminate the auxillary modes like the cavity modes in cqed.
        (Defined in subclasses)
        """
        return U

    def get_operators_labels(self):
        """
        Get the labels for each Hamiltonian.
        It is used in the method``plot_pulses``.
        It is a 2-d nested list, in the plot,
        a different color will be used for each sublist.
        """
        label_list = []
        for pulse in self.pulses:
            label_list.append(pulse.label)
        return [label_list]

    def plot_pulses(
            self, title=None, figsize=(12, 6), dpi=None,
            show_axis=False, rescale_pulse_coeffs=True,
            num_steps=1000):
        """
        Plot the ideal pulse coefficients.

        Parameters
        ----------
        title: str, optional
            Title for the plot.

        figsize: tuple, optional
            The size of the figure.

        dpi: int, optional
            The dpi of the figure.

        show_axis: bool, optional
            If the axis are shown.

        rescale_pulse_coeffs: bool, optional
            Rescale the hight of each pulses.

        num_steps: int, optional
            Number of time steps in the plot.

        Returns
        -------
        fig: matplotlib.figure.Figure
            The `Figure` object for the plot.

        ax: matplotlib.axes._subplots.AxesSubplot
            The axes for the plot.

        Notes
        -----
        ``plot_pulses`` only works for array_like coefficients
        """
        import matplotlib.pyplot as plt
        import matplotlib.gridspec as gridspec
        color_list = plt.rcParams['axes.prop_cycle'].by_key()['color']

        # create a axis for each pulse
        fig = plt.figure(figsize=figsize, dpi=dpi)
        grids = gridspec.GridSpec(len(self.pulses), 1)
        grids.update(wspace=0., hspace=0.)

        tlist = np.linspace(0., self.get_full_tlist()[-1], num_steps)
        dt = tlist[1] - tlist[0]

        # make sure coeffs start and end with zero, for ax.fill
        tlist = np.hstack(([-dt*1.e-20], tlist, [tlist[-1] + dt*1.e-20]))
        coeffs = []
        for pulse in self.pulses:
            coeffs.append(_pulse_interpolate(pulse, tlist))

        pulse_ind = 0
        axis = []
        for i, label_group in enumerate(self.get_operators_labels()):
            for j, label in enumerate(label_group):
                grid = grids[pulse_ind]
                ax = plt.subplot(grid)
                axis.append(ax)
                ax.fill(tlist, coeffs[pulse_ind], color_list[i], alpha=0.7)
                ax.plot(tlist, coeffs[pulse_ind], color_list[i])
                if rescale_pulse_coeffs:
                    ymax = np.max(np.abs(coeffs[pulse_ind])) * 1.1
                else:
                    ymax = np.max(np.abs(coeffs)) * 1.1
                if ymax != 0.:
                    ax.set_ylim((-ymax, ymax))

                # disable frame and ticks
                if not show_axis:
                    ax.set_xticks([])
                    ax.spines['bottom'].set_visible(False)
                ax.spines['top'].set_visible(False)
                ax.spines['right'].set_visible(False)
                ax.spines['left'].set_visible(False)
                ax.set_yticks([])
                ax.set_ylabel(label, rotation=0)
                pulse_ind += 1
                if i == 0 and j == 0 and title is not None:
                    ax.set_title(title)
        fig.tight_layout()
        return fig, axis

def _pulse_interpolate(pulse, tlist):
    """
    A function that calls Scipy interpolation routine. Used for plotting.
    """
    if pulse.tlist is None and pulse.coeff is None:
        coeff = np.zeros(len(tlist))
        return coeff
    if isinstance(pulse.coeff, bool):
        if pulse.coeff:
            coeff = np.ones(len(tlist))
        else:
            coeff = np.zeros(len(tlist))
        return coeff
    coeff = pulse.coeff
    if len(coeff) == len(pulse.tlist)-1:  # for discrete pulse
        coeff = np.concatenate([coeff, [0]])

    from scipy import interpolate
    if pulse.spline_kind == "step_func":
        kind = "previous"
    else:
        kind = "cubic"
    inter = interpolate.interp1d(
        pulse.tlist, coeff, kind=kind,
        bounds_error=False, fill_value=0.0)
    return inter(tlist)


#!/usr/bin/env python
# -*- coding: utf-8 -*-

"""
Deprecated Colour Models Transformations
========================================

Defines various deprecated colour models transformations:

-   :func:`RGB_to_HSV`
-   :func:`HSV_to_RGB`
-   :func:`RGB_to_HSL`
-   :func:`HSL_to_RGB`
-   :func:`RGB_to_CMY`
-   :func:`CMY_to_RGB`
-   :func:`CMY_to_CMYK`
-   :func:`CMYK_to_CMY`
-   :func:`RGB_to_HEX`
-   :func:`HEX_to_RGB`

These colour models are stated as deprecated because they trade off perceptual
relevance for computation speed. They should not be used in the colour science
domain although they are useful for image analysis and provide end user
software colour selection tools.

They are provided for convenience and completeness.

Warning
-------
Don't use that! Seriously...

References
----------
.. [1]  http://en.wikipedia.org/wiki/HSL_and_HSV
        (Last accessed 10 August 2014)
.. [2]  `Color Gamut Transform Pairs
        <http://alvyray.com/Papers/CG/color78.pdf>`_,
        DOI: http://dx.doi.org/10.1145/800248.807361
        (Last accessed 10 August 2014)
"""

from __future__ import division, unicode_literals

import math
import numpy as np

__author__ = 'Colour Developers'
__copyright__ = 'Copyright (C) 2013 - 2014 - Colour Developers'
__license__ = 'New BSD License - http://opensource.org/licenses/BSD-3-Clause'
__maintainer__ = 'Colour Developers'
__email__ = 'colour-science@googlegroups.com'
__status__ = 'Production'

__all__ = ['RGB_to_HSV',
           'HSV_to_RGB',
           'RGB_to_HSL',
           'HSL_to_RGB',
           'RGB_to_CMY',
           'CMY_to_RGB',
           'CMY_to_CMYK',
           'CMYK_to_CMY',
           'RGB_to_HEX',
           'HEX_to_RGB']

def RGB_to_HSV(RGB):
    """
    Converts from *RGB* colourspace to *HSV* colourspace.

    Parameters
    ----------
    RGB : array_like, (3,)
        *RGB* colourspace matrix.

    Returns
    -------
    ndarray, (3,)
        *HSV* matrix.

    Notes
    -----
    -   Input *RGB* colourspace matrix is in domain [0, 1].
    -   Output *HSV* colourspace matrix is in domain [0, 1].

    References
    ----------
    .. [3]  http://www.easyrgb.com/index.php?X=MATH&H=20#text20
            (Last accessed 18 May 2014)

    Examples
    --------
    >>> RGB = np.array([0.49019608, 0.98039216, 0.25098039])
    >>> RGB_to_HSV(RGB)  # doctest: +ELLIPSIS
    array([ 0.2786738...,  0.744     ,  0.98039216])
    """

    R, G, B = np.ravel(RGB)

    minimum = min(R, G, B)
    maximum = max(R, G, B)
    delta = maximum - minimum

    V = maximum

    if delta == 0:
        H = 0
        S = 0
    else:

        S = delta / maximum

        delta_R = (((maximum - R) / 6) + (delta / 2)) / delta
        delta_G = (((maximum - G) / 6) + (delta / 2)) / delta
        delta_B = (((maximum - B) / 6) + (delta / 2)) / delta

        if R == maximum:
            H = delta_B - delta_G
        elif G == maximum:
            H = (1 / 3) + delta_R - delta_B
        elif B == maximum:
            H = (2 / 3) + delta_G - delta_R

        if H < 0:
            H += 1
        if H > 1:
            H -= 1

    return np.array([H, S, V])

def HSV_to_RGB(HSV):
    """
    Converts from *HSV* colourspace to *RGB* colourspace.

    Parameters
    ----------
    HSV : array_like, (3,)
        *HSV* colourspace matrix.

    Returns
    -------
    ndarray, (3,)
        *RGB* colourspace matrix.

    Notes
    -----
    -   Input *HSV* colourspace matrix is in domain [0, 1].
    -   Output *RGB* colourspace matrix is in domain [0, 1].

    References
    ----------
    .. [4]  http://www.easyrgb.com/index.php?X=MATH&H=21#text21
            (Last accessed 18 May 2014)

    Examples
    --------
    >>> HSV = np.array([0.27867384, 0.744, 0.98039216])
    >>> HSV_to_RGB(HSV)  # doctest: +ELLIPSIS
    array([ 0.4901960...,  0.9803921...,  0.2509803...])
    """

    H, S, V = np.ravel(HSV)

    if S == 0:
        R = V
        G = V
        B = V
    else:
        h = H * 6
        if h == 6:
            h = 0

        i = math.floor(h)
        j = V * (1 - S)
        k = V * (1 - S * (h - i))
        l = V * (1 - S * (1 - (h - i)))
        if i == 0:
            R = V
            G = l
            B = j
        elif i == 1:
            R = k
            G = V
            B = j
        elif i == 2:
            R = j
            G = V
            B = l
        elif i == 3:
            R = j
            G = k
            B = V
        elif i == 4:
            R = l
            G = j
            B = V
        elif i == 5:
            R = V
            G = j
            B = k

    return np.array([R, G, B])

def RGB_to_HSL(RGB):
    """
    Converts from *RGB* colourspace to *HSL* colourspace.

    Parameters
    ----------
    RGB : array_like, (3,)
        *RGB* colourspace matrix.

    Returns
    -------
    ndarray, (3,)
        *HSL* matrix.

    Notes
    -----
    -   Input *RGB* colourspace matrix is in domain [0, 1].
    -   Output *HSL* colourspace matrix is in domain [0, 1].

    References
    ----------
    .. [5]  http://www.easyrgb.com/index.php?X=MATH&H=18#text18
            (Last accessed 18 May 2014)

    Examples
    --------
    >>> RGB = np.array([0.49019608, 0.98039216, 0.25098039])
    >>> RGB_to_HSL(RGB)  # doctest: +ELLIPSIS
    array([ 0.2786738...,  0.9489796...,  0.6156862...])
    """

    R, G, B = np.ravel(RGB)

    minimum = min(R, G, B)
    maximum = max(R, G, B)
    delta = maximum - minimum

    L = (maximum + minimum) / 2

    if delta == 0:
        H = 0
        S = 0
    else:

        S = delta / (maximum + minimum) if L < 0.5 else delta / (
            2 - maximum - minimum)

        delta_R = (((maximum - R) / 6) + (delta / 2)) / delta
        delta_G = (((maximum - G) / 6) + (delta / 2)) / delta
        delta_B = (((maximum - B) / 6) + (delta / 2)) / delta

        if R == maximum:
            H = delta_B - delta_G
        elif G == maximum:
            H = (1 / 3) + delta_R - delta_B
        elif B == maximum:
            H = (2 / 3) + delta_G - delta_R

        if H < 0:
            H += 1
        if H > 1:
            H -= 1

    return np.array([H, S, L])

def HSL_to_RGB(HSL):
    """
    Converts from *HSL* colourspace to *RGB* colourspace.

    Parameters
    ----------
    HSL : array_like, (3,)
        *HSL* colourspace matrix.

    Returns
    -------
    ndarray, (3,)
        *RGB* colourspace matrix.

    Notes
    -----
    -   Input *HSL* colourspace matrix is in domain [0, 1].
    -   Output *RGB* colourspace matrix is in domain [0, 1].

    References
    ----------
    .. [6]  http://www.easyrgb.com/index.php?X=MATH&H=19#text19
            (Last accessed 18 May 2014)

    Examples
    --------
    >>> HSL = np.array([0.27867384, 0.94897959, 0.61568627])
    >>> HSL_to_RGB(HSL)  # doctest: +ELLIPSIS
    array([ 0.4901960...,  0.9803921...,  0.2509803...])
    """

    H, S, L = np.ravel(HSL)

    if S == 1:
        R = L
        G = L
        B = L
    else:
        def H_to_RGB(vi, vj, vH):
            """
            Converts *hue* value to *RGB* colourspace.
            """

            if vH < 0:
                vH += 1
            if vH > 1:
                vH -= 1
            if 6 * vH < 1:
                return vi + (vj - vi) * 6 * vH
            if 2 * vH < 1:
                return vj
            if 3 * vH < 2:
                return vi + (vj - vi) * ((2 / 3) - vH) * 6
            return vi

        j = L * (1 + S) if L < 0.5 else (L + S) - (S * L)
        i = 2 * L - j

        R = H_to_RGB(i, j, H + (1 / 3))
        G = H_to_RGB(i, j, H)
        B = H_to_RGB(i, j, H - (1 / 3))

    return np.array([R, G, B])

def RGB_to_CMY(RGB):
    """
    Converts from *RGB* colourspace to *CMY* colourspace.

    Parameters
    ----------
    RGB : array_like, (3,)
        *RGB* colourspace matrix.

    Returns
    -------
    ndarray, (3,)
        *CMY* matrix.

    Notes
    -----
    -   Input *RGB* colourspace matrix is in domain [0, 1].
    -   Output *CMY* colourspace matrix is in domain [0, 1].

    References
    ----------
    .. [7]  http://www.easyrgb.com/index.php?X=MATH&H=11#text11
            (Last accessed 18 May 2014)

    Examples
    --------
    >>> RGB = np.array([0.49019608, 0.98039216, 0.25098039])
    >>> RGB_to_CMY(RGB)  # doctest: +ELLIPSIS
    array([ 0.5098039...,  0.0196078...,  0.7490196...])
    """

    R, G, B = np.ravel(RGB)
    return np.array([1 - R, 1 - G, 1 - B])

def CMY_to_RGB(CMY):
    """
    Converts from *CMY* colourspace to *CMY* colourspace.

    Parameters
    ----------
    CMY : array_like, (3,)
        *CMY* colourspace matrix.

    Returns
    -------
    ndarray, (3,)
        *RGB* colourspace matrix.

    Notes
    -----
    -   Input *CMY* colourspace matrix is in domain [0, 1].
    -   Output *RGB* colourspace matrix is in domain [0, 1].

    References
    ----------
    .. [8]  http://www.easyrgb.com/index.php?X=MATH&H=12#text12
            (Last accessed 18 May 2014)

    Examples
    --------
    >>> CMY = np.array([0.50980392, 0.01960784, 0.74901961])
    >>> CMY_to_RGB(CMY)  # doctest: +ELLIPSIS
    array([ 0.4901960...,  0.9803921...,  0.2509803...])
    """

    C, M, Y = np.ravel(CMY)
    return np.array([1 - C, 1 - M, 1 - Y])

def CMY_to_CMYK(CMY):
    """
    Converts from *CMY* colourspace to *CMYK* colourspace.

    Parameters
    ----------
    CMY : array_like, (3,)
        *CMY* colourspace matrix.

    Returns
    -------
    ndarray, (4,)
        *CMYK* matrix.

    Notes
    -----
    -   Input *CMY* colourspace matrix is in domain [0, 1].
    -   Output*CMYK* colourspace matrix is in domain [0, 1].

    References
    ----------
    .. [9]  http://www.easyrgb.com/index.php?X=MATH&H=13#text13
            (Last accessed 18 May 2014)

    Examples
    --------
    >>> CMY = np.array([0.50980392, 0.01960784, 0.74901961])
    >>> CMY_to_CMYK(CMY)  # doctest: +ELLIPSIS
    array([ 0.5       ,  0.        ,  0.744     ,  0.0196078...])
    """

    C, M, Y = np.ravel(CMY)

    K = 1

    if C < K:
        K = C
    if M < K:
        K = M
    if Y < K:
        K = Y
    if K == 1:
        C = 0
        M = 0
        Y = 0
    else:
        C = (C - K) / (1 - K)
        M = (M - K) / (1 - K)
        Y = (Y - K) / (1 - K)

    return np.array([C, M, Y, K])

def CMYK_to_CMY(CMYK):
    """
    Converts from *CMYK* colourspace to *CMY* colourspace.

    Parameters
    ----------
    CMYK : array_like, (4,)
        *CMYK* colourspace matrix.

    Returns
    -------
    ndarray, (3,)
        *CMY* matrix.

    Notes
    -----
    -   Input *CMYK* colourspace matrix is in domain [0, 1].
    -   Output *CMY* colourspace matrix is in domain [0, 1].

    References
    ----------
    .. [10] http://www.easyrgb.com/index.php?X=MATH&H=14#text14
            (Last accessed 18 May 2014)

    Examples
    --------
    >>> CMYK = np.array([0.5, 0, 0.744, 0.01960784])
    >>> CMYK_to_CMY(CMYK)  # doctest: +ELLIPSIS
    array([ 0.5098039...,  0.0196078...,  0.7490196...])
    """

    C, M, Y, K = np.ravel(CMYK)

    return np.array(
        [C * (1 - K) + K, M * (1 - K) + K, Y * (1 - K) + K])

def RGB_to_HEX(RGB):
    """
    Converts from *RGB* colourspace to hex triplet representation.

    Parameters
    ----------
    RGB : array_like, (3,)
        *RGB* colourspace matrix.

    Returns
    -------
    unicode
        Hex triplet representation.

    Notes
    -----
    -   Input *RGB* colourspace matrix is in domain [0, 1].

    Examples
    --------
    >>> RGB = np.array([0.66666667, 0.86666667, 1])
    >>> # Doctests skip for Python 2.x compatibility.
    >>> RGB_to_HEX(RGB)  # doctest: +SKIP
    '#aaddff'
    """

    RGB = np.ravel(RGB)
    R, G, B = map(int, RGB * 255)
    return '#{0:02x}{1:02x}{2:02x}'.format(R, G, B)

def HEX_to_RGB(HEX):
    """
    Converts from hex triplet representation to *RGB* colourspace.

    Parameters
    ----------
    HEX : unicode
        Hex triplet representation.

    Returns
    -------
    ndarray, (3,)
        *RGB* colourspace matrix.

    Notes
    -----
    -   Output *RGB* colourspace matrix is in domain [0, 1].

    Examples
    --------
    >>> HEX = '#aaddff'
    >>> HEX_to_RGB(HEX)  # doctest: +ELLIPSIS
    array([ 0.6666666...,  0.8666666...,  1.        ])
    """

    HEX = HEX.lstrip('#')
    length = len(HEX)
    return np.array([int(HEX[i:i + length // 3], 16) for i in
                     range(0, length, length // 3)]) / 255


#!/usr/bin/env python

from __future__ import division
import mdtraj as md
import numpy
from Bython.Cheminfo import Bymol
from Bython.Structure.configstruc import NON_POLAR_ATOMS, PROTEIN_AROMATIC_RES, POLAR_ATOMS, PROTEIN_POSITIVE, PROTEIN_NEGATIVE, PROTEIN_POLAR_SC
from itertools import product, chain

def SIFT_7bit(traj_ob=None, res_index2=None, res_index=None, aro_cut=4.0, apolar_cut=4.5, hbond_cut=3.5,elec_cut=4.0,verbose=True, use_sc_only=False):
    #Make a list of all protein residue indexes if res_index not defined
    if res_index2 == None:
        res_index2 = [traj_ob.topology.atom(atom_index).residue.index for atom_index in traj_ob.topology.select(' chainid 0 1 and name CA')]

    if res_index == None:
        res_index = [traj_ob.topology.atom(atom_index).residue.index for atom_index in traj_ob.topology.select('chainid 2 3 and name CA')]

        print ("Fingerprint resolution: 7-bits per residue (Apolar, Aromatic Face to Face, Aromatic Edge to Face, "+
               "Hbond Protein as Donor, Hbond Protein as Acceptor, Electrostatic Protein +ve, Electrostatic Protein -ve)")
    #Initialize the array, No_residue X 7 X No_of_ligands
    #sift7bit= numpy.zeros((len(res_index)*7*len(lig_names)), dtype=numpy.int) #traj_ob.xyz.shape[0],
    sift7bit = {} #Key: Ligand index, value: list of bits = 7*no_residues
    #Go through each residue and compute its interaction with each ligands
    for q_res in res_index:
        #Get the residue heavy atom indices
        if not use_sc_only:
            res_heavy_idx = [a_idx for a_idx in traj_ob.topology.select("resid " + str(q_res))
                             if traj_ob.topology.atom(a_idx).element.symbol != 'H']
        else:
            #Do only if residue is not Gly
            if traj_ob.topology.residue(q_res).name.lower() == 'gly':
                continue
            res_heavy_idx = [a_idx for a_idx in traj_ob.topology.select("resid " + str(q_res))
                             if traj_ob.topology.atom(a_idx).element.symbol != 'H' and traj_ob.topology.atom(a_idx).is_sidechain == True]

        for q2_res in res_index2:
            #Initialize the residue bits for this ligand
            temp_bits = [0]*7

	    #Get the residue heavy atom indices
            if not use_sc_only:
                res2_heavy_idx = [a_idx for a_idx in traj_ob.topology.select("resid " + str(q2_res))
                             if traj_ob.topology.atom(a_idx).element.symbol != 'H']
            else:
                #Do only if residue is not Gly
                if traj_ob.topology.residue(q2_res).name.lower() == 'gly':
                    continue
                res2_heavy_idx = [a_idx for a_idx in traj_ob.topology.select("resid " + str(q2_res))
                             if traj_ob.topology.atom(a_idx).element.symbol != 'H' and traj_ob.topology.atom(a_idx).is_sidechain == True]

            #Create residue and ligand atom-pairs
            a_pairs_dist = list(product(numpy.array(res_heavy_idx), numpy.array(res2_heavy_idx)))
            #Compute all pairwise distances and covert to Angstrom; These distances are in nanometers
            pair_distances = md.compute_distances(traj_ob, a_pairs_dist)*10 #shape (1, no_of_a_pairs_dist)
            #Check If ligand is within Apolar interaction distance
            if numpy.any(pair_distances[pair_distances < apolar_cut]):
                #Check if interaction is apolar
                if IsApolar(traj_ob, numpy.array(a_pairs_dist), pair_distances[0],apolar_cut):
                    #Change the apolar bit
                    temp_bits[0] = 1
            #Check If ligand is within Aromatic interaction distance
            if numpy.any(pair_distances[pair_distances < aro_cut]):
                #Check if interaction is aromatic
                isf2f, ise2f = IsAromatic(traj_ob, numpy.array(a_pairs_dist)
                                                 , pair_distances[0], aro_cut)
                if isf2f:
                    #Change the f2f bit
                    temp_bits[1] = 1
                if ise2f:
                    #Change the f2f bit
                    temp_bits[2] = 1
            #Check If ligand is within Hbond interaction distance
            if numpy.any(pair_distances[pair_distances < hbond_cut]):
                #Check if interaction is Hbond type
                isprotdon, isprotacc = IsHbond(traj_ob, numpy.array(a_pairs_dist),pair_distances[0], hbond_cut)
                if isprotdon:
                    #Change the protein donor bit
                    temp_bits[3] = 1
                if isprotacc:
                    #Change the protein acceptor bit
                    temp_bits[4] = 1
            #Check If ligand is within Electrostatic interaction distance
            if numpy.any(pair_distances[pair_distances < elec_cut]):
                #Check if interaction is Electrostatic
                prot_pos, prot_neg = IsElectro(traj_ob, numpy.array(a_pairs_dist),
                                               pair_distances[0], elec_cut)

                if prot_pos:
                    #Change the protein positive bit
                    temp_bits[5] = 1
                if prot_neg:
                    #Change the protein negative bit
                    temp_bits[6] = 1

            #append temp_bits to sift7bits
            if not sift7bit.has_key(q2_res):
                sift7bit[q2_res] = temp_bits
            else:
                sift7bit[q2_res].extend(temp_bits)
    #Now Combine the bits for different ligands (also corresponding residue indexes) and return it
    if len(res_index2)> 1:
        combined_sift = [lig_bits for key in sorted(sift7bit) for lig_bits in sift7bit[key]]
        #return combined_sift, list(numpy.array(res_index)+1)*len(sift7bit) #+1 since res_index is 0-based
        return combined_sift, (res_index)*len(sift7bit)
    else:
        return sift7bit[0], res_index

def IsApolar(traj_object, pair_indices, pair_dist, apolar_cut):
    #Get the atom pair indices that are within cutoff
    prospective_pairs = pair_indices[pair_dist < apolar_cut]
    isapolar = False
    #Check if any of the pairs is non_polar_atom pair
    for a_pair in prospective_pairs:
        #Get the atom symbol
        a1_symbol = traj_object.topology.atom(a_pair[0]).element.symbol.lower()
        a2_symbol = traj_object.topology.atom(a_pair[1]).element.symbol.lower()
        if a1_symbol in NON_POLAR_ATOMS and a2_symbol in NON_POLAR_ATOMS:
            isapolar = True
            break
    return isapolar

def IsAromatic(traj_object, pair_indices, pair_dist, aro_cut):
    #Get the atom pair indices that are within cutoff
    prospective_pairs = pair_indices[pair_dist < aro_cut]
    ise2f = False #Interaction is edge to face
    isf2f = False #Interaction is face to face

    #Check if any of the pairs is aromatic pair
    for a_pair in prospective_pairs:
        #Get the residue name (first index is for protein atom)
        a1_res_name = traj_object.topology.atom(a_pair[0]).residue.name.lower()
	a2_res_name = traj_object.topology.atom(a_pair[1]).residue.name.lower()
        #Check if residue is aromatic type
        if a1_res_name in PROTEIN_AROMATIC_RES and a2_res_name in PROTEIN_AROMATIC_RES:
            #Check if protein atom is in aromatic ring
            a1_atom_name = traj_object.topology.atom(a_pair[0]).name.lower()
            if traj_object.topology.atom(a_pair[0]).is_sidechain and a1_atom_name not in ('cb', 'oh'): # oh for tyr
                #Check if the ligand atom is also aromatic
                a2_atom_name = traj_object.topology.atom(a_pair[1]).name.lower()
                if traj_object.topology.atom(a_pair[1]).is_sidechain and a2_atom_name not in ('cb', 'oh'): # oh for tyr
                    ###Check the angles to see if inteaction is f2f or e2f###

                    #Get the Neighbor atoms (atom objects) of protein atom
                    a1_atom_neighb = GetNeighbor(traj_object, traj_object.topology.atom(a_pair[0]))
                    #Get the coordinates for protein atom in Angstrom and two of its ring neighbor; Hydrogens don't matter
                    a1_atom_cord = traj_object.xyz[:,a_pair[0],:][0]*10
                    a1_atom_neighb1_cord = traj_object.xyz[:,a1_atom_neighb[0].index,:][0]*10
                    a1_atom_neighb2_cord = traj_object.xyz[:,a1_atom_neighb[1].index,:][0]*10
                    #Get the cross product for protein ring
                    prot_ring_normal = GetCross(numpy.array([a1_atom_cord, a1_atom_neighb1_cord, a1_atom_neighb2_cord]))
                    uni_norm_prot_ring = prot_ring_normal / numpy.sqrt((prot_ring_normal*prot_ring_normal).sum()) #Unit-vector

		    #Get the Neighbor atoms (atom objects) of 2nd protein atom
                    a2_atom_neighb = GetNeighbor(traj_object, traj_object.topology.atom(a_pair[1]))
                    #Get the coordinates for protein atom in Angstrom and two of its ring neighbor; Hydrogens don't matter
                    a2_atom_cord = traj_object.xyz[:,a_pair[1],:][0]*10
                    a2_atom_neighb1_cord = traj_object.xyz[:,a2_atom_neighb[0].index,:][0]*10
                    a2_atom_neighb2_cord = traj_object.xyz[:,a2_atom_neighb[1].index,:][0]*10
                    #Get the cross product for protein ring
                    prot_ring_normal = GetCross(numpy.array([a2_atom_cord, a2_atom_neighb1_cord, a2_atom_neighb2_cord]))
                    uni_norm_prot2_ring = prot_ring_normal / numpy.sqrt((prot_ring_normal*prot_ring_normal).sum()) #Unit-vector

		    #Get the angle between unit normal vectors
                    normal_thetha_deg = numpy.degrees(numpy.arccos(numpy.dot(uni_norm_prot2_ring, uni_norm_prot_ring)))  # In radians

                    if normal_thetha_deg <= 30.0 or normal_thetha_deg >= 150.0:
                        isf2f = True
                    if normal_thetha_deg > 30.0 and normal_thetha_deg < 150.0:
                        ise2f = True
                    #If both e2f and f2f have been assigned then no need to check further pairs for this residue
                    if isf2f and ise2f:
                        break
    return isf2f, ise2f

def IsHbond(traj_object, pair_indices, pair_dist, hbond_cut):
    #Get the atom pair indices that are within cutoff
    prospective_pairs = pair_indices[pair_dist < hbond_cut]
    hbond_prot_acceptor = False
    hbond_prot_donor = False
    #Check if any of the pairs is polar_atom pair
    for a_pair in prospective_pairs:
        #Get the atom symbol
        a1_symbol = traj_object.topology.atom(a_pair[0]).element.symbol.lower()
        a2_symbol = traj_object.topology.atom(a_pair[1]).element.symbol.lower()
        if a1_symbol in POLAR_ATOMS and a2_symbol in POLAR_ATOMS:
            ###Check if angle criteria is satisfied###
            #Get the Neighbor atoms (atom objects) of protein atom
            a1_atom_neighb = GetNeighbor(traj_object, traj_object.topology.atom(a_pair[0]))
            #Check if protein 1 atom is DONOR
            prot_don_hyd = [a1_neig for a1_neig in a1_atom_neighb if a1_neig.element.symbol == 'H']
            if prot_don_hyd:
                is_prot_donor = True
            else:
                is_prot_donor = False

            #Get the neighbors of the 2nd protein atom
            #Get the Neighbor atoms (atom objects) of protein atom
            a2_atom_neighb = GetNeighbor(traj_object, traj_object.topology.atom(a_pair[1]))
            #Check if protein 2 atom is DONOR
            prot2_don_hyd = [a2_neig for a2_neig in a2_atom_neighb if a2_neig.element.symbol == 'H']
            if prot2_don_hyd:
                is_prot2_donor = True
            else:
                is_prot2_donor = False
            #Get protein atom coordinate
            a1_atom_cord = traj_object.xyz[:,a_pair[0],:][0]*10
            #Get protein 2 atom coordinate
            a2_atom_cord = traj_object.xyz[:,a_pair[1],:][0]*10
            #Check angles if protein acceptor and protein 2 donor
            if not is_prot_donor and is_prot2_donor:
                #Go through each neighboring hydrogen atom
                for h_neig in prot2_don_hyd:
                    #Get the coordinates
                    prot2_hyd_cord = traj_object.xyz[:,h_neig.index,:][0]*10
                    Thetha_degree = GetAngle(numpy.array([a1_atom_cord, prot2_hyd_cord, a2_atom_cord]))
                    if Thetha_degree > 135.0:
                        hbond_prot_acceptor = True
                        break #No need to go through rest of the hydrogens
            #Check angles if protein donor and protein 2 acceptor
            if is_prot_donor and not is_prot2_donor:
                #Go through each neighboring hydrogen atom
                for h_neig in prot_don_hyd:
                    #Get the coordinates
                    prot_hyd_cord = traj_object.xyz[:,h_neig.index,:][0]*10
                    Thetha_degree = GetAngle(numpy.array([a2_atom_cord, prot_hyd_cord, a1_atom_cord]))
                    if Thetha_degree > 135.0:
                        hbond_prot_donor = True
                        break #No need to go through rest of the hydrogens
            #Check if both ligand and protein atom are donor, then Hbond still possible with one as acceptor and other as donor
            if is_prot_donor and is_prot2_donor:
                #Do Hbond determination for both to see which one is acceptor/donor
                #Assuming protein donor, go through each neighboring hydrogen atom of protein donor
                for h_neig in prot_don_hyd:
                    #Get the coordinates
                    prot_hyd_cord = traj_object.xyz[:,h_neig.index,:][0]*10
                    Thetha_degree = GetAngle(numpy.array([a2_atom_cord, prot_hyd_cord, a1_atom_cord]))
                    if Thetha_degree > 135.0:
                        hbond_prot_donor = True
                        break #No need to go through rest of the hydrogens
                #Assuming ligand donor, go through each neighboring hydrogen atom of ligand donor
                for h_neig in prot2_don_hyd:
                    #Get the coordinates
                    prot2_hyd_cord = traj_object.xyz[:,h_neig.index,:][0]*10
                    Thetha_degree = GetAngle(numpy.array([a1_atom_cord, prot2_hyd_cord, a2_atom_cord]))
                    if Thetha_degree > 135.0:
                        hbond_prot_acceptor = True
                        break #No need to go through rest of the hydrogens

            #If both donor and acceptor have been assigned then no need to check further pairs for this residue
            if hbond_prot_acceptor and hbond_prot_donor:
                break
    return hbond_prot_donor, hbond_prot_acceptor

def IsElectro(traj_object, pair_indices, pair_dist, elec_cut):
    #Get the atom pair indices that are within cutoff
    prospective_pairs = pair_indices[pair_dist < elec_cut]
    prot_pos = False
    prot_neg = False
    #Check if any of the pairs is polar_atom pair
    for a_pair in prospective_pairs:
        #Get the residue name (first index is for protein atom)
        a1_res_name = traj_object.topology.atom(a_pair[0]).residue.name.lower()
	a2_res_name = traj_object.topology.atom(a_pair[1]).residue.name.lower()
        #Check if residue is charged type
        if a1_res_name in PROTEIN_POSITIVE + PROTEIN_NEGATIVE and a2_res_name in PROTEIN_POSITIVE + PROTEIN_NEGATIVE:
            #Get the atom symbol
            a1_symbol = traj_object.topology.atom(a_pair[0]).element.symbol.lower()
            #Check if protein atom is in sidechain and is polar
            if traj_object.topology.atom(a_pair[0]).is_sidechain and a1_symbol in POLAR_ATOMS:
                #Check if protein 2 atom is in sidechain and polar
                a2_symbol = traj_object.topology.atom(a_pair[1]).element.symbol.lower()
            	if traj_object.topology.atom(a_pair[1]).is_sidechain and a2_symbol in POLAR_ATOMS:
                    #If protein 1 is negatively charged and protein 2 is postively charged
                    if a1_res_name in PROTEIN_NEGATIVE and a2_res_name in PROTEIN_POSITIVE:
                        prot_neg = True
                    if a1_res_name in PROTEIN_POSITIVE and a2_res_name in PROTEIN_NEGATIVE:
                        prot_pos = True
            if prot_neg and prot_pos:
                break
    return prot_pos, prot_neg

def GetNeighbor(traj_object, atom_name):
    '''
    Returns a tuple of Neighbors of atom (atom object is returned)
    atom_name (in the format ex. TYR86-CE1; mdtraj atom object)
    '''
    #print atom_name
    each_neighb = [list(each_bond) for each_bond in traj_object.topology.bonds if atom_name in each_bond]
    if not each_neighb:
        raise Exception("Topology Object doesn't contain bond information. Use PDB file to define top in Readtraj")
    neighbs = list(set([all_neighb for all_neighb in chain(*each_neighb)]))
    neighbs.remove(atom_name)
    return tuple(neighbs)

def GetCross(cord_array):
    '''
    Cross product of input 3x3 cord_array
    '''
    vec_12 = cord_array[1] - cord_array[0] #Vector tail at atm1 and head at atm2
    vec_13 = cord_array[2] - cord_array[0] #Vector tail at atm1 and head at atm3
    return numpy.cross(vec_12, vec_13)

def GetAngle(cord_array):
    '''
    Returns the angle formed by atom coordinates in input 3x3 cord_array
    '''
    vec_21 = cord_array[0] - cord_array[1] #Vector tail at atm2 and head at atm1
    vec_31 = cord_array[2] - cord_array[1] #Vector tail at atm2 and head at atm3
    dot = numpy.dot(vec_21, vec_31)
    vec_21mod = numpy.sqrt((vec_21*vec_21).sum()) #length of vector 21
    vec_31mod = numpy.sqrt((vec_31*vec_31).sum())
    return numpy.degrees(numpy.arccos(dot / vec_21mod / vec_31mod))