import warnings
from typing import Any, Dict, List, Optional, Tuple

import numpy as np
import pandas as pd

from .unilateral import Unilateral
from .utils import HDFMixin, draw_diagnose_times, fast_binomial_pmf

# I chose not to make this one a child of System, since it is basically only a
# container for two System instances
class Bilateral(HDFMixin):
    """Class that models metastatic progression in a lymphatic system
    bilaterally by creating two :class:`Unilateral` instances that are
    symmetric in their connections. The parameters describing the spread
    probabilities however need not be symmetric.

    See Also:
        :class:`Unilateral`: Two instances of this class are created as
        attributes.
    """
    def __init__(
        self,
        graph: Dict[Tuple[str], List[str]] = {},
        base_symmetric: bool = False,
        trans_symmetric: bool = True,
        **kwargs
    ):
        """Initialize both sides of the network as a :class:`Unilateral`
        instance:

        Args:
            graph: Dictionary of the same kind as for initialization of
                :class:`Unilateral`. This graph will be passed to the
                constructors of two :class:`Unilateral` attributes of this
                class.
            base_symmetric: If ``True``, the spread probabilities of the two
                sides from the tumor(s) to the LNLs will be set symmetrically.
            trans_symmetric: If ``True``, the spread probabilities among the
                LNLs will be set symmetrically.
        """
        self.ipsi   = Unilateral(graph=graph)   # ipsilateral and...
        self.contra = Unilateral(graph=graph)   # ...contralateral network

        self.base_symmetric  = base_symmetric
        self.trans_symmetric = trans_symmetric

    def __str__(self):
        """Print info about the structure and parameters of the bilateral
        lymphatic system.
        """
        num_tumors = len(self.ipsi.tumors)
        num_lnls   = len(self.ipsi.lnls)
        string = (
            f"Bilateral lymphatic system with {num_tumors} tumor(s) "
            f"and 2 * {num_lnls} LNL(s).\n"
        )
        string += "Symmetry: "
        string += "base " if self.base_symmetric else ""
        string += "trans\n" if self.trans_symmetric else "\n"
        string += "Ipsilateral:\t" + " ".join([f"{e}" for e in self.ipsi.edges])
        string += "\n"
        string += "Contralateral:\t" + " ".join([f"{e}" for e in self.contra.edges])

        return string

    @property
    def graph(self) -> Dict[Tuple[str], List[str]]:
        """Return the (unilateral) graph that was used to create this network.
        """
        return self.ipsi.graph

    @property
    def system(self):
        """Return a dictionary with the ipsi- & contralateral side's
        :class:`Unilateral` under the keys ``"ipsi"`` and ``"contra"``
        respectively.

        This is needed since in some weak moment, I thought it would be a great
        idea if a class named ``BilateralSystem`` had an attriute called
        ``system`` which contained two instances of the ``System`` class under
        the keys ``"ipsi"`` and ``"contra"``...
        """
        return {
            "ipsi"  : self.ipsi,
            "contra": self.contra
        }

    @property
    def state(self) -> np.ndarray:
        """
        Return the currently state (healthy or involved) of all LNLs in the
        system.
        """
        ipsi_state = self.ipsi.state
        contra_state = self.contra.state
        return np.concatenate([ipsi_state, contra_state])

    @state.setter
    def state(self, newstate: np.ndarray):
        """
        Set the state of the system to ``newstate``.
        """
        self.ipsi.state = newstate[:len(self.ipsi.lnls)]
        self.contra.state = newstate[len(self.ipsi.lnls):]

    @property
    def base_probs(self) -> np.ndarray:
        """Probabilities of lymphatic spread from the tumor(s) to the lymph
        node levels. If the ipsi- & contralateral spread from the tumor is set
        to be symmetric (``base_symmetric = True``) this only returns the
        parameters of one side. So, the returned array is composed like so:

        +-----------------+--------------------+
        | base probs ipsi | base probs contra* |
        +-----------------+--------------------+

        *Only when ``base_symmetric = False``, which is the default.

        When setting these parameters, the length of the provided array only
        needs to be half as long if ``base_symmetric`` is ``True``, since both
        sides will be set to the same values.

        See Also:
            :attr:`Unilateral.base_probs`
        """
        if self.base_symmetric:
            return self.ipsi.base_probs
        else:
            return np.concatenate([self.ipsi.base_probs,
                                   self.contra.base_probs])

    @base_probs.setter
    def base_probs(self, new_base_probs: np.ndarray):
        """Set the base probabilities from the tumor(s) to the LNLs.
        """
        if self.base_symmetric:
            self.ipsi.base_probs = new_base_probs
            self.contra.base_probs = new_base_probs
        else:
            num_base_probs = len(self.ipsi.base_edges)
            self.ipsi.base_probs = new_base_probs[:num_base_probs]
            self.contra.base_probs = new_base_probs[num_base_probs:]

    @property
    def trans_probs(self) -> np.ndarray:
        """Probabilities of lymphatic spread among the lymph node levels. If
        this ipsi- & contralateral spread is set to be symmetric
        (``trans_symmetric = True``) this only returns the parameters of one
        side. Similiar to the :attr:`base_probs`, this array's shape is:

        +------------------+---------------------+
        | trans probs ipsi | trans probs contra* |
        +------------------+---------------------+

        *Only if ``trans_symmetric = False``.

        And correspondingly, if setting these transmission probability one only
        needs half as large an array if ``trans_symmetric`` is ``True``.

        See Also:
            :attr:`Unilateral.trans_probs`
        """
        if self.trans_symmetric:
            return self.ipsi.trans_probs
        else:
            return np.concatenate([self.ipsi.trans_probs,
                                   self.contra.trans_probs])

    @trans_probs.setter
    def trans_probs(self, new_trans_probs: np.ndarray):
        """Set the transmission probabilities (from LNL to LNL) of the network.
        """
        if self.trans_symmetric:
            self.ipsi.trans_probs = new_trans_probs
            self.contra.trans_probs = new_trans_probs
        else:
            num_trans_probs = len(self.ipsi.trans_edges)
            self.ipsi.trans_probs = new_trans_probs[:num_trans_probs]
            self.contra.trans_probs = new_trans_probs[num_trans_probs:]

    @property
    def spread_probs(self) -> np.ndarray:
        """The parameters representing the probabilities for lymphatic spread
        along a directed edge of the graph representing the lymphatic network.

        If the bilateral network is set to have symmetries, the length of the
        list/array of numbers that need to be provided will be shorter. E.g.,
        when the bilateral lymphatic network is completely asymmetric, it
        requires an array of length :math:`2n_b + 2n_t` where :math:`n_b` is
        the number of edges from the tumor to the LNLs and :math:`n_t` the
        number of edges among the LNLs.

        Similar to the :attr:`base_probs` and the :attr:`trans_probs`, we can
        describe its shape like this:

        +-----------------+--------------------+------------------+----------------------+
        | base probs ipsi | base probs contra* | trans probs ipsi | trans probs contra** |
        +-----------------+--------------------+------------------+----------------------+

        | *Only if ``base_symmetric = False``, which is the default.
        | **Only if ``trans_symmetric = False``.

        See Also:
            :attr:`Unilateral.spread_probs`
        """
        return np.concatenate([self.base_probs, self.trans_probs])

    @spread_probs.setter
    def spread_probs(self, new_spread_probs: np.ndarray):
        """Set the spread probabilities of the :class:`Edge` instances in the
        the network.
        """
        num_base_probs = len(self.ipsi.base_edges)

        if self.base_symmetric:
            self.base_probs = new_spread_probs[:num_base_probs]
            self.trans_probs = new_spread_probs[num_base_probs:]
        else:
            self.base_probs = new_spread_probs[:2*num_base_probs]
            self.trans_probs = new_spread_probs[2*num_base_probs:]

    @property
    def modalities(self):
        """Compute the two system's observation matrices
        :math:`\\mathbf{B}^{\\text{i}}` and :math:`\\mathbf{B}^{\\text{c}}`.

        See Also:
            :meth:`Unilateral.modalities`: Setting modalities in unilateral
            System.
        """
        ipsi_modality_spsn = self.ipsi.modalities
        if ipsi_modality_spsn != self.contra.modalities:
            msg = ("Ipsi- & contralaterally stored modalities are not the same")
            raise RuntimeError(msg)

        return ipsi_modality_spsn

    @modalities.setter
    def modalities(self, modality_spsn: Dict[str, List[float]]):
        """
        Given specificity :math:`s_P` & sensitivity :math:`s_N` of different
        diagnostic modalities, compute the system's two observation matrices
        :math:`\\mathbf{B}_i` and :math:`\\mathbf{B}_c`.
        """
        self.ipsi.modalities = modality_spsn
        self.contra.modalities = modality_spsn

    @property
    def patient_data(self):
        """Table with rows of patients. Columns should have three levels. The
        first column is ('info', 'tumor', 't_stage'). The rest of the columns
        are separated by modality names on the top level, then subdivided into
        'ipsi' & 'contra' by the second level and finally, in the third level,
        the names of the lymph node level are given. Here is an example of such
        a table:

        +---------+----------------------+----------------------+
        |  info   |         MRI          |         PET          |
        +---------+----------+-----------+----------+-----------+
        |  tumor  |   ipsi   |  contra   |   ipsi   |  contra   |
        +---------+----------+-----------+----------+-----------+
        | t_stage |    II    |    II     |    II    |    II     |
        +=========+==========+===========+==========+===========+
        | early   | ``True`` | ``None``  | ``True`` | ``False`` |
        +---------+----------+-----------+----------+-----------+
        | late    | ``None`` | ``None``  | ``None`` | ``None``  |
        +---------+----------+-----------+----------+-----------+
        | early   | ``True`` | ``False`` | ``True`` | ``True``  |
        +---------+----------+-----------+----------+-----------+
        """
        try:
            return self._patient_data
        except AttributeError:
            raise AttributeError(
                "No patient data has been loaded yet"
            )

    @patient_data.setter
    def patient_data(self, patient_data: pd.DataFrame):
        """Load the patient data. For now, this just calls the :meth:`load_data`
        method, but at a later point, I would like to write a function here
        that generates the pandas :class:`DataFrame` from the internal matrix
        representation of the data.
        """
        self._patient_data = patient_data.copy()
        self.load_data(patient_data)

    def load_data(
        self,
        data: pd.DataFrame,
        t_stages: Optional[List[int]] = None,
        modality_spsn: Optional[Dict[str, List[float]]] = None,
        mode: str = "HMM"
    ):
        """Load a dataset by converting it into internal representation as data
        matrix.

        Args:
            data: Table with rows of patients. Columns must have three levels.
                The first column is ('info', 'tumor', 't_stage'). The rest of
                the columns are separated by modality names on the top level,
                then subdivided into 'ipsi' & 'contra' by the second level and
                finally, in the third level, the names of the lymph node level
                are given. Here is an example of such a table:

                +---------+---------------------+-----------------------+
                |  info   |         MRI         |         PET           |
                +---------+----------+----------+-----------+-----------+
                |  tumor  |   ipsi   |  contra  |   ipsi    |  contra   |
                +---------+----------+----------+-----------+-----------+
                | t_stage |    II    |    II    |    II     |    II     |
                +=========+==========+==========+===========+===========+
                | early   | ``True`` | ``None`` | ``True``  | ``False`` |
                +---------+----------+----------+-----------+-----------+
                | late    | ``None`` | ``None`` | ``False`` | ``False`` |
                +---------+----------+----------+-----------+-----------+
                | early   | ``True`` | ``True`` | ``True``  | ``None``  |
                +---------+----------+----------+-----------+-----------+

        See Also:
            :meth:`Unilateral.load_data`: Data loading method of unilateral
            system.
        """
        # split the DataFrame into two, one for ipsi-, one for contralateral
        ipsi_data = data.drop(
            columns=["contra"], axis=1, level=1, inplace=False
        )
        ipsi_data = pd.DataFrame(
            ipsi_data.values,
            index=ipsi_data.index,
            columns=ipsi_data.columns.droplevel(1)
        )
        contra_data = data.drop(
            columns=["ipsi"], axis=1, level=1, inplace=False
        )
        contra_data = pd.DataFrame(
            contra_data.values,
            index=contra_data.index,
            columns=contra_data.columns.droplevel(1)
        )

        self.ipsi.load_data(
            ipsi_data,
            t_stages=t_stages,
            modality_spsn=modality_spsn,
            mode=mode
        )
        self.contra.load_data(
            contra_data,
            t_stages=t_stages,
            modality_spsn=modality_spsn,
            mode=mode
        )

    def _are_valid_(self, new_spread_probs: np.ndarray) -> bool:
        """Check that the spread probability (rates) are all within limits.
        """
        if new_spread_probs.shape != self.spread_probs.shape:
            msg = ("Shape of provided spread parameters does not match network")
            raise ValueError(msg)
        if np.any(np.greater(0., new_spread_probs)):
            return False
        if np.any(np.greater(new_spread_probs, 1.)):
            return False

        return True

    def _log_likelihood(
        self,
        t_stages: Optional[List[Any]] = None,
        diag_times: Optional[Dict[Any, int]] = None,
        max_t: Optional[int] = 10,
        time_dists: Optional[Dict[Any, np.ndarray]] = None
    ):
        """Compute the log-likelihood of data, using the stored spread probs.
        This method mainly exists so that the checking and assigning of the
        spread probs can be skipped.
        """
        llh = 0.

        if diag_times is not None:
            if len(diag_times) != len(t_stages):
                msg = ("One diagnose time must be provided for each T-stage.")
                raise ValueError(msg)

            for stage in t_stages:
                diag_time = np.around(diag_times[stage]).astype(int)
                if diag_time > max_t:
                    return -np.inf

                # probabilities for any hidden state (ipsi- & contralaterally)
                state_probs = {}
                state_probs["ipsi"] = self.ipsi._evolve(diag_time)
                state_probs["contra"] = self.contra._evolve(diag_time)

                # joint probs for ipsi- & contralateral hidden states
                joint_state_probs = np.outer(state_probs["ipsi"],
                                             state_probs["contra"])
                log_p = np.log(
                    np.sum(
                        self.ipsi.diagnose_matrices[stage]
                        * (joint_state_probs
                           @ self.contra.diagnose_matrices[stage]),
                        axis=0
                    )
                )
                llh += np.sum(log_p)

            return llh

        elif time_dists is not None:
            if len(time_dists) != len(t_stages):
                msg = ("One distribution over diagnose times must be provided "
                       "for each T-stage.")
                raise ValueError(msg)

            # subtract 1, to also consider healthy starting state (t = 0)
            max_t = len(time_dists[t_stages[0]]) - 1

            state_probs = {}
            state_probs["ipsi"] = self.ipsi._evolve(t_last=max_t)
            state_probs["contra"] = self.contra._evolve(t_last=max_t)

            for stage in t_stages:
                joint_state_probs = (
                    state_probs["ipsi"].T
                    @ np.diag(time_dists[stage])
                    @ state_probs["contra"]
                )
                log_p = np.log(
                    np.sum(
                        self.ipsi.diagnose_matrices[stage]
                        * (joint_state_probs
                           @ self.contra.diagnose_matrices[stage]),
                        axis=0
                    )
                )
                llh += np.sum(log_p)

            return llh

        else:
            msg = ("Either provide a list of diagnose times for each T-stage "
                   "or a distribution over diagnose times for each T-stage.")
            raise ValueError(msg)

    def log_likelihood(
        self,
        spread_probs: np.ndarray,
        t_stages: Optional[List[Any]] = None,
        diag_times: Optional[Dict[Any, int]] = None,
        max_t: Optional[int] = 10,
        time_dists: Optional[Dict[Any, np.ndarray]] = None
    ):
        """Compute log-likelihood of (already stored) data, given the spread
        probabilities and either a discrete diagnose time or a distribution to
        use for marginalization over diagnose times.

        Args:
            spread_probs: Spread probabiltites from the tumor to the LNLs, as
                well as from (already involved) LNLs to downsream LNLs. Includes
                both sides of the neck. The composition of this array is:

                +----------------------------+-----------------------------+
                | base probs (ipsi & contra) | trans probs (ipsi & contra) |
                +----------------------------+-----------------------------+

                If certain symmetries are chosen, only one set of base or
                transmission probabilities might have to be provided.

            t_stages: List of T-stages that are also used in the data to denote
                how advanced the primary tumor of the patient is. This does not
                need to correspond to the clinical T-stages 'T1', 'T2' and so
                on, but can also be more abstract like 'early', 'late' etc.

            diag_times: For each T-stage, one can specify with what time step
                the likelihood should be computed. If this is set to `None`,
                and a distribution over diagnose times `time_dists` is provided,
                the function marginalizes over diagnose times.

            max_t: Latest possible diagnose time. This is only used to return
                `-np.inf` in case one of the `diag_times` exceeds this value.

            time_dists: Distribution over diagnose times that can be used to
                compute the likelihood of the data, given the spread
                probabilities, but marginalized over the time of diagnosis. If
                set to `None`, a diagnose time must be explicitly set for each
                T-stage.

        Returns:
            The log-likelihood :math:`\\log{p(D \\mid \\theta)}` where :math:`D`
            is the data and :math:`\\theta` is the tuple of spread probabilities
            and diagnose times or distributions over diagnose times.

        See Also:
            :attr:`spread_probs`: Property for getting and setting the spread
            probabilities, of which a lymphatic network has as many as it has
            :class:`Edge` instances (in case no symmetries apply).

            :meth:`Unilateral.log_likelihood`: The log-likelihood function of
            the unilateral system.
        """
        if not self._are_valid_(spread_probs):
            return -np.inf

        self.spread_probs = spread_probs

        if t_stages is None:
            t_stages = list(self.ipsi.f.keys())

        return self._log_likelihood(
            t_stages=t_stages,
            diag_times=diag_times,
            max_t=max_t,
            time_dists=time_dists,
        )

    def marginal_log_likelihood(
        self,
        theta: np.ndarray,
        t_stages: Optional[List[Any]] = None,
        time_dists: dict = {}
    ) -> float:
        """
        Compute the likelihood of the (already stored) data, given the spread
        parameters, marginalized over time of diagnosis via time distributions.
        Wraps the :meth:`log_likelihood` method.

        Args:
            theta: Spread probabiltites from the tumor to the LNLs, as well as
                from (already involved) LNLs to downsream LNLs. Includes both
                sides of the neck. The composition of this array is:

                +----------------------------+-----------------------------+
                | base probs (ipsi & contra) | trans probs (ipsi & contra) |
                +----------------------------+-----------------------------+

                If certain symmetries are chosen, only one set of base or
                transmission probabilities might have to be provided.

            t_stages: List of T-stages that should be included in the learning
                process.

            time_dists: Distribution over the probability of diagnosis at
                different times :math:`t` given T-stage.

        Returns:
            The log-likelihood of a parameter sample.

        See Also:
            :meth:`log_likelihood`: Simply calls the actual likelihood function
            where it sets the `diag_times` to `None`.
        """
        return self.log_likelihood(
            theta, t_stages,
            diag_times=None, time_dists=time_dists
        )

    def time_log_likelihood(
        self,
        theta: np.ndarray,
        t_stages: List[Any],
        max_t: int = 10
    ) -> float:
        """
        Compute likelihood given the spread parameters and the time of diagnosis
        for each T-stage. Wraps the :math:`log_likelihood` method.

        Args:
            theta: Set of parameters, consisting of the spread probabilities
                and the time of diagnosis for all T-stages. It is therefore
                made up these parts and in that order:

                +------------+-------------+----------------+
                | base probs | trans probs | diagnose times |
                +------------+-------------+----------------+

            t_stages: keywords of T-stages that are present in the dictionary of
                C matrices and the previously loaded dataset.

            max_t: Latest accepted time-point.

        Returns:
            The likelihood of the data, given the spread parameters as well as
            the diagnose time for each T-stage.

        See Also:
            :meth:`log_likelihood`: The `theta` argument of this function is
            split into `spread_probs` and `diag_times`, which are then passed
            to the actual likelihood function.
        """
        # splitting theta into spread parameters and...
        len_spread_probs = len(theta) - len(t_stages)
        spread_probs = theta[:len_spread_probs]
        # ...diagnose times for each T-stage
        tmp = theta[len_spread_probs:]
        diag_times = {t_stages[t]: tmp[t] for t in range(len(t_stages))}

        return self.log_likelihood(
            spread_probs, t_stages,
            diag_times=diag_times, max_t=max_t, time_dists=None
        )

    def binom_marg_log_likelihood(
        self,
        theta: np.ndarray,
        t_stages: List[Any],
        max_t: int = 10
    ) -> float:
        """
        Compute marginal log-likelihood using binomial distributions to sum
        over the diagnose times.

        Args:
            theta: Set of parameters, consisting of the spread probabilities
                and the binomial distribution's :math:`p` parameters for each
                T-category. One has to provide a concatenated array of these
                numbers like this:

                +------------+-------------+--------------------------+
                | base probs | trans probs | binomial :math`p` params |
                +------------+-------------+--------------------------+

            t_stages: keywords of T-stages that are present in the dictionary of
                C matrices and the previously loaded dataset.

            max_t: Latest accepted time-point.

        Returns:
            The log-likelihood of the (already stored) data, given the spread
            prbabilities as well as the parameters for binomial distribtions
            used to marginalize over diagnose times.
        """
        # splitting theta into spread parameters and...
        len_spread_probs = len(theta) - len(t_stages)
        spread_probs = theta[:len_spread_probs]
        # ...p-values for the binomial distribution
        p = theta[len_spread_probs:]

        if np.any(np.greater(p, 1.)) or np.any(np.less(p, 0.)):
            return -np.inf

        t = np.arange(max_t + 1)
        time_dists = {}
        for i,stage in enumerate(t_stages):
            time_dists[stage] = fast_binomial_pmf(t, max_t, p[i])

        return self.marginal_log_likelihood(
            spread_probs, t_stages,
            time_dists=time_dists
        )

    def risk(
        self,
        spread_probs: Optional[np.ndarray] = None,
        inv: Dict[str, Optional[np.ndarray]] = {"ipsi": None, "contra": None},
        diagnoses: Dict[str, Dict] = {"ipsi": {}, "contra": {}},
        diag_time: Optional[int] = None,
        time_dist: Optional[np.ndarray] = None,
        mode: str = "HMM"
    ) -> float:
        """Compute risk of ipsi- & contralateral involvement given specific (but
        potentially incomplete) diagnoses for each side of the neck.

        Args:
            spread_probs: Set of new spread parameters. If not given (``None``),
                the currently set parameters will be used.

            inv: Dictionary that can have the keys ``"ipsi"`` and ``"contra"``
                with the respective values being the involvements of interest.
                If (for one side or both) no involvement of interest is given,
                it'll be marginalized.
                The array themselves may contain ``True``, ``False`` or ``None``
                for each LNL corresponding to the risk for involvement, no
                involvement and "not interested".

            diagnoses: Dictionary that itself may contain two dictionaries. One
                with key "ipsi" and one with key "contra". The respective value
                is then a dictionary that can hold a potentially incomplete
                (mask with ``None``) diagnose for every available modality.
                Leaving out available modalities will assume a completely
                missing diagnosis.

            diag_time: Time of diagnosis. Either this or the `time_dist` to
                marginalize over diagnose times must be given.

            time_dist: Distribution to marginalize over diagnose times. Either
                this, or the `diag_time` must be given.

            mode: Set to ``"HMM"`` for the hidden Markov model risk (requires
                the ``time_dist``) or to ``"BN"`` for the Bayesian network
                version.
        """
        if spread_probs is not None:
            self.spread_probs = spread_probs

        cX = {}   # marginalize over matching complete involvements.
        cZ = {}   # marginalize over Z for incomplete diagnoses.
        pXt = {}  # probability p(X|t) of state X at time t as 2D matrices
        pD = {}   # probability p(D|X) of a (potentially incomplete) diagnose,
                  # given an involvement. Should be a 1D vector

        for side in ["ipsi", "contra"]:
            involvement = np.array(inv[side])
            # build vector to marginalize over involvements
            cX[side] = np.zeros(shape=(len(self.system[side].state_list)),
                                dtype=bool)
            for i,state in enumerate(self.system[side].state_list):
                cX[side][i] = np.all(
                    np.equal(
                        involvement, state,
                        where=(involvement!=None),
                        out=np.ones_like(involvement, dtype=bool)
                    )
                )

            # create one large diagnose vector from the individual modalitie's
            # diagnoses
            obs = np.array([])
            for mod in self.system[side]._spsn_tables:
                if mod in diagnoses[side]:
                    obs = np.append(obs, diagnoses[side][mod])
                else:
                    obs = np.append(obs, np.array([None] * len(self.system[side].lnls)))

            # build vector to marginalize over diagnoses
            cZ[side] = np.zeros(shape=(len(self.system[side].obs_list)),
                                dtype=bool)
            for i,complete_obs in enumerate(self.system[side].obs_list):
                cZ[side][i] = np.all(
                    np.equal(
                        obs, complete_obs,
                        where=(obs!=None),
                        out=np.ones_like(obs, dtype=bool)
                    )
                )

            if diag_time is not None:
                pXt[side] = self.system[side]._evolve(diag_time)

            elif time_dist is not None:
                max_t = len(time_dist)
                pXt[side] = self.system[side]._evolve(t_last=max_t-1)

            else:
                msg = ("Either diagnose time or distribution to marginalize "
                       "over it must be given.")
                raise ValueError(msg)

            pD[side] = self.system[side].observation_matrix @ cZ[side]

        # joint probability of Xi & Xc (marginalized over time). Acts as prior
        # for p( Di,Dc | Xi,Xc ) and should be a 2D matrix
        if diag_time is not None:
            pXX = np.outer(pXt["ipsi"], pXt["contra"])

        elif time_dist is not None:
            # time-prior in diagnoal matrix form
            PT = np.diag(time_dist)
            pXX = pXt["ipsi"].T @ PT @ pXt["contra"]

        # joint probability of all hidden states and the requested diagnosis
        pDDXX = np.einsum("i,ij,j->ij", pD["ipsi"], pXX, pD["contra"])
        # joint probability of the requested involvement and diagnosis
        pDDII = cX["ipsi"].T @ pDDXX @ cX["contra"]

        # denominator p(Di, Dc). Joint probability for ipsi- & contralateral
        # diagnoses. Marginalized over all hidden involvements and over all
        # matching complete observations that give rise to the specific
        # diagnose. The result should be just a number
        pDD = (cZ["ipsi"].T
               @ self.ipsi.observation_matrix.T
               @ pXX
               @ self.contra.observation_matrix
               @ cZ["contra"])

        return pDDII / pDD

    def generate_dataset(
        self,
        num_patients: int,
        stage_dist: List[float],
        diag_times: Optional[Dict[Any, int]] = None,
        time_dists: Optional[Dict[Any, np.ndarray]] = None,
    ) -> pd.DataFrame:
        """Generate/sample a pandas :class:`DataFrame` from the defined network.

        Args:
            num_patients: Number of patients to generate.
            stage_dist: Probability to find a patient in a certain T-stage.
            diag_times: For each T-stage, one can specify until which time step
                the corresponding patients should be evolved. If this is set to
                ``None``, and a distribution over diagnose times ``time_dists``
                is provided, the diagnose time is drawn from the ``time_dist``.
            time_dists: Distributions over diagnose times that can be used to
                draw a diagnose time for the respective T-stage.
        """
        drawn_t_stages, drawn_diag_times = draw_diagnose_times(
            num_patients=num_patients,
            stage_dist=stage_dist,
            diag_times=diag_times,
            time_dists=time_dists
        )

        drawn_obs_ipsi = self.ipsi._draw_patient_diagnoses(drawn_diag_times)
        drawn_obs_contra = self.contra._draw_patient_diagnoses(drawn_diag_times)
        drawn_obs = np.concatenate([drawn_obs_ipsi, drawn_obs_contra], axis=1)

        # construct MultiIndex for dataset from stored modalities
        sides = ["ipsi", "contra"]
        modalities = list(self.modalities.keys())
        lnl_names = [lnl.name for lnl in self.ipsi.lnls]
        multi_cols = pd.MultiIndex.from_product([sides, modalities, lnl_names])

        # create DataFrame
        dataset = pd.DataFrame(drawn_obs, columns=multi_cols)
        dataset = dataset.reorder_levels(order=[1, 0, 2], axis="columns")
        dataset = dataset.sort_index(axis="columns", level=0)
        dataset[('info', 'tumor', 't_stage')] = drawn_t_stages

        return dataset

class BilateralSystem(Bilateral):
    """Class kept for compatibility after renaming to :class:`Bilateral`.

    See Also:
        :class:`Bilateral`
    """
    def __init__(self, *args, **kwargs):
        msg = ("This class has been renamed to `Bilateral`.")
        warnings.warn(msg, DeprecationWarning)

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


#!/usr/bin/env python
#
#Copyright 2019 Allan Haldane.

#This file is part of Mi3-GPU.

#Mi3-GPU is free software: you can redistribute it and/or modify
#it under the terms of the GNU General Public License as published by
#the Free Software Foundation, version 3 of the License.

#Mi3-GPU is distributed in the hope that it will be useful,
#but WITHOUT ANY WARRANTY; without even the implied warranty of
#MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
#GNU General Public License for more details.

#You should have received a copy of the GNU General Public License
#along with Mi3-GPU.  If not, see <http://www.gnu.org/licenses/>.

#Contact: allan.haldane _AT_ gmail.com
import sys, os, argparse
from Bio.Alphabet import IUPAC
import seqload
import numpy as np
from alphabet_reduction import getLq

def indmap(oldalpha, amap):
    def ind(x):
        for i in range(len(amap)):
            if x in amap[i]:
                return i
        return 0

    return  np.array([ind(let) for let in oldalpha])

def reduceSeqAlphaPerpos(seqs, newalphas, oldalpha, out=None):
    rseqs = np.empty(seqs.shape, dtype=int)
    for n,a in enumerate(newalphas):
        conv = indmap(oldalpha, a)
        rseqs[:,n] = conv[seqs[:,n]]

    if out is None:
        out = sys.stdout
    seqload.writeSeqs(out, rseqs, "ABCDEFGHIJKLMNOPQRSTUVWXYZ", noheader=True)

def reduceBimAlphaPerpos(bimarg, newalphas, oldalpha, out):
    if out is None:
        raise ValueError('out argument required for bimarg reduction')

    L, q = getLq(bimarg)
    qout = len(newalphas[0])
    nPairs = L*(L-1)//2

    # strategy: compute index mapping array to use in add.at
    maps = np.array([indmap(oldalpha, a) for a in newalphas])
    qqmap = np.array([np.add.outer(qout*maps[i], maps[j]).ravel()
                      for i in range(L-1) for j in range(i+1,L)])
    nmap = np.broadcast_to(np.arange(nPairs, dtype=int)[:,None], (nPairs, q*q))

    # allocate new bimarg
    newbim = np.zeros((nPairs, qout*qout), dtype='f8')

    # do the accumulation with `at`
    np.add.at(newbim, (nmap, qqmap), bimarg)

    # renormalize
    newbim /= np.sum(newbim, axis=1, keepdims=True)

    np.save(out, newbim.astype('f4'))

def main():
    parser = argparse.ArgumentParser(
                                description='Apply alphabet reduction to MSA')
    parser.add_argument('file', help='either seq file or bimarg file')
    parser.add_argument('alphamap')
    parser.add_argument('--alpha', default='protgap')
    parser.add_argument('--out')

    args = parser.parse_args(sys.argv[1:])
    alphabets = {'protein': IUPAC.protein.letters,
                 'protgap': '-' + IUPAC.protein.letters,
                 'charge': '0+-',
                 'nuc': "ACGT"}
    alpha = alphabets.get(args.alpha, args.alpha)

    with open(args.alphamap) as f:
        # assumed to be a file containing the output of alphabet reduction, but
        # only for one reduction level.  Each line should look like:
        # ALPHA8 -DNAGSQFMYCI E HWP K L R T V
        newalphas = [a.split()[1:] for a in f.readlines()]

    try:
        bimarg = np.load(args.file)
    except ValueError:
        seqs = seqload.loadSeqs(args.file, alpha)[0]
        reduceSeqAlphaPerpos(seqs, newalphas, alpha, args.out)
    else:
        reduceBimAlphaPerpos(bimarg, newalphas, alpha, args.out)

if __name__ == '__main__':
    main()


# Script for running MCMC sampling of EDF
# Constructs mock sample, then re-samples the velocities

import sys, os
dir_path = os.path.dirname(os.path.realpath(__file__))
sys.path.append(dir_path+'/py/')
import edf_sampling
import json
import numpy as np
import emcee
from itertools import izip
import pandas as pd
import subprocess
import pickle

# default parameters
names = ['age', 'RcP', 'mass', 'vR', 'vphi', 'vz', 'I', 'l', 'b', 's',
         'Z', 'Teff', 'logg', 'R', 'phi', 'z', 'vlos', 'pm_l', 'pm_b',
         'RA', 'DEC', 'pm_ra', 'pm_dec', 'I', 'JR', 'Lz', 'Jz', 'Rc', 'logl']

def LogL_Zsample(*args):
    return edf_sampling.LogL_sample(*args)

def deccut_fn(l,b,deccut):
    ad = edf_sampling.check_dec(l,b)
    if(deccut<np.pi):
        if(deccut>0.):
            if(ad<deccut-.5*np.pi):
                return 0
        if(deccut<0.):
            if(ad>deccut+.5*np.pi):
                return 0
    return 0

def run_mcmc_sample(
                 Nsamples, mag_band, maglimits,
                 # Isochrone set -- currently probably only Padova will work in all cases
                 which_iso='Padova',
                 # Extra magnitudes to calculate
                 extra_mags=None,
                 # Introduce a colour cut
                 color=None, colorlimits=None,
                 # Coordinate cuts (in radians) -- remove everything |b|<modbcut
                 # If deccut>0 cut all dec<deccut-pi/2
                 # If deccut<0 cut all dec>deccut+pi/2
                 modbcut=0.,deccut=2.*np.pi,
                 # Spectroscopic parameter cut -- log10Teff
                 logglimits=None,Tefflimits=None,fehlimits=None,
                 # Flags for different options
                 extinct=True, with_halo=True, interp=False,dered=False,
                 # Debugging and outputting
                 output_file=None,messages=True, extra_magnitudes=None,
                 # Sampler parameters
                 nwalkers=1000,Nburn=None,threads=1,thin=20,pt=False,ntemp=1,asampler=2.,
                 # Adjust magnitude selection with a taper (1-break_grad(mag-break_pos))
                 #break_pos=1000.,break_grad=1000.,
                 # If cutting on error-convolved teff, logg and color -- currently won't do anything
                 tgcolor_errs=np.array([0.,0.,0.])):

    '''
	Draws a Monte Carlo sample from an EDF model

	Draws Nsamples across the sky between magnitudes maglimits for band mag_band
    '''
    nsamples = int(Nsamples/nwalkers)+1

    # Turn on/off halo
    if(with_halo):
        edf_sampling.turn_on_halo()
    else:
        edf_sampling.turn_off_halo()

    # -- Set arguments
    check = lambda x: np.array([-10000.,10000.]) if x is None else x
    checkcolor = lambda x: np.array(['J','K']) if x is None else color # as a default
    args =[which_iso,mag_band,maglimits,extinct,interp,modbcut,deccut,check(color),
           check(colorlimits),check(Tefflimits),check(logglimits),check(fehlimits)]

    # -- Set sampler
    ndim, nwalker = 9, nwalkers
    sampler = emcee.EnsembleSampler(nwalker, ndim,
                                    LogL_Zsample, threads=threads,
                                    args=args)
    if(pt):
        # Parallel tempering
        sampler = emcee.PTSampler(ntemp, nwalker, ndim,los_magbox_LogL_sample_py,
                                  logp, loglargs=llos_args, threads=threads)
    # Initialize walkers
    # Sample uniform in age, mass
    # Sample gaussian in Z and velocities
    # Sample exp in (unextincted) magnitude
    lomag,himag=maglimits[0],maglimits[1]
    Zlim=[-3.,1.]
    if fehlimits is not None:
        Zlim=fehcut
    if not with_halo:
        Zlim[0]=edf_sampling.minZ()
        Zlim[1]=edf_sampling.maxZ()
    p0 = np.random.uniform(
        low=[0.2, Zlim[0], 0.5, -150., 50., -100., lomag, 0., modbcut],
        high=[12., Zlim[1], 3.5, 150., 350., 100., himag, 2.*np.pi, np.pi /2.],
        size=(nwalker, ndim))
    p0[:,-1]*=(1.-2.*(np.random.uniform(low=0,high=1,size=nwalker)>0.5))
    if fehlimits is not None:
        meanZ,sigZ = 0.,.4
        p0.T[1]=np.random.normal(size=nwalker)*sigZ+meanZ
    meanV,sigV = 0.,50.
    for i in range(3,6):
        p0.T[i]=np.random.normal(size=nwalker)*sigV+meanV
    p0.T[4]+=200.
    p0.T[6]=trunc_exp_rv(lomag,himag,1.,size=len(p0))

    # -- Only choose physically allowed masses for selected age, Z
    for i in np.arange(len(p0)):
        maxmass,minmass=0.,0.
        maxmass = edf_sampling.check_highmass_Z(p0[i][0], p0[i][1],which_iso)
        minmass = edf_sampling.check_lowmass_Z(p0[i][0], p0[i][1],which_iso)
        while(maxmass < 0. or minmass<0. or
	      edf_sampling.check_radius_positive(p0[i][0], p0[i][1]) == 0 or
              deccut_fn(p0[i][-2],p0[i][-1],deccut)):
            p0[i][-3] = np.random.uniform(low=lomag, high=himag)
            p0[i][1] = np.random.uniform(low=Zlim[0], high=Zlim[1])
            p0[i][0] = np.random.uniform(low=0., high=12.)
            maxmass = edf_sampling.check_highmass_Z(p0[i][0], p0[i][1],which_iso)
            minmass = edf_sampling.check_lowmass_Z(p0[i][0], p0[i][1],which_iso)
        MAXMASS=3.
        if(maxmass>MAXMASS):
            maxmass=MAXMASS
        p0[i][2] = np.random.uniform(low=minmass, high=maxmass)
        # Check other cuts satisfied
        if(check(colorlimits)[0]>-10. or check(logglimits)[0]>-10. or check(Tefflimits)[0]>-10.):
            while(edf_sampling.check_color_logg_cut(p0[i][0], p0[i][1], p0[i][2],
                                                 0.,180.,checkcolor(color),
						 check(colorlimits),check(logglimits),check(Tefflimits),
						 which_iso,False)==0):
                    p0[i][2] = np.random.uniform(low=minmass, high=maxmass)

    # Now wiggle the walkers a bit
    err = [0.5,0.1,0.1,10.,10.,10.,0.1,0.01,0.01]
    for i in np.arange(len(p0)):
        extinctMag = edf_sampling.get_extinct(p0[i][-2],p0[i][-1],3.,mag_band)
        if(extinctMag>err[-3]):
            err[-3]=extinctMag
        pp = edf_sampling.LogL_sample(p0[i],*args)
        n=0
        while(np.isinf(pp) and n<10000):
            p0[i]=p0[np.random.randint(len(p0))]+np.random.normal(size=ndim)*err
            pp = edf_sampling.LogL_sample(p0[i],*args)
            n+=1
        if(n==10000):
            print "Can't find start point:",p0[i],deccut_fn(p0[i][-2],p0[i][-1],deccut)

    if(messages): print 'Initial points sampled'
    if(pt):
        p0=np.reshape(p0,(ntemp,nwalker,ndim))

    if(Nburn==None):
        Nburn=2*thin*nsamples
    # -- Run a burn-in
    pos, prob, state = sampler.run_mcmc(p0, Nburn, storechain=False)
    if messages: print 'Number of logl=-inf = '+str(len(prob[np.isinf(prob)]))
    if messages: print("Burnt")
    sampler.reset()

    # -- Sample with thinning and calculating dependent variables
    pos, prob, state = sampler.run_mcmc(pos, nsamples*thin, thin=thin)
    if messages: print ("Sampled")
    flatchain = sampler.flatchain
    lnprob = sampler.lnprobability
    if(pt):
        flatchain=flatchain[0]
        lnprob=lnprob[0]
    extras = np.array(
        map(lambda i: edf_sampling.get_extra_data(i,mag_band,which_iso,False,extinct,interp), sampler.flatchain))
    actions = np.array([edf_sampling.get_actions(np.concatenate((b[4:7], a[3:6])))
                        for a, b in izip(flatchain, extras)])

    nameslist = np.copy(names)
    nameslist[1]="Z"
    nameslist[10]="RcP"
    nameslist[23]=mag_band
    nameslist[6]=mag_band+"0"

    if messages: print("Mean acceptance fraction:", np.mean(sampler.acceptance_fraction))
    everything = np.vstack((flatchain.T, extras.T,
                     actions.T,lnprob.flatten())).T
    if color is not None:
        if extra_magnitudes is not None:
            extra_magnitudes=np.unique(np.concatenate((color,extra_magnitudes,np.array([mag_band]))))
        else:
            extra_magnitudes=np.unique(np.concatenate((color,np.array([mag_band]))))
        extra_magnitudes = extra_magnitudes[extra_magnitudes!=mag_band]

    if extra_magnitudes is not None:
        extra_mags = np.array(
        map(lambda i: edf_sampling.get_extra_magnitudes(i,
                       mag_band, extra_magnitudes,which_iso, False, extinct,interp), flatchain))
        nameslist=np.concatenate((nameslist,extra_magnitudes))
        print extra_magnitudes
        nameslist=np.concatenate((nameslist,np.array([e+'0' for e in list(extra_magnitudes)])))
        everything = np.vstack((everything.T, extra_mags.T)).T

    df = pd.DataFrame(everything,columns=nameslist)
    df = df.sample(Nsamples,replace=False).reset_index(drop=True)
    if(output_file):
        df.to_csv(output_file)
    return df

def los_magbox_LogL_sample_py(*args):
    return edf_sampling.los_magbox_LogL_sample(*args)

import scipy.stats

def trunc_exp_rv(lo, hi, scale, size):
    low = 0.
    high=hi-lo
    rnd_cdf = np.random.uniform(scipy.stats.expon.cdf(x=low, scale=scale),
                                scipy.stats.expon.cdf(x=high, scale=scale),
                                size=size)
    return hi-scipy.stats.expon.ppf(q=rnd_cdf, scale=scale)

def logp(p):
    return 1.

def run_mcmc_los(Nsamples, l, b, mag_band, maglimits,
                 # Field radius (in radians)
                 fieldradius=-10.,
                 # Isochrone set -- currently probably only Padova will work in all cases
                 which_iso='Padova',
                 # Extra magnitudes to calculate
                 extra_mags=None,
                 # Introduce a colour cut
                 color=None, colorlimits=None,
                 # Spectroscopic parameter cut
                 logglimits=None,Tefflimits=None,
                 # Flags for different options
                 extinct=True, RcPorZ=False, with_halo=True, interp=False,dered=False,
                 # Debugging and outputting
                 output_file=None,messages=True, extra_magnitudes=None,
                 # Sampler parameters
                 nwalkers=1000,Nburn=None,threads=1,thin=20,pt=False,ntemp=1,asampler=2.,
                 # Adjust magnitude selection with a taper (1-break_grad(mag-break_pos))
                 break_pos=1000.,break_grad=1000.,
                 # If cutting on error-convolved teff, logg and color
                 tgcolor_errs=np.array([0.,0.,0.])):

    '''
	Draws a Monte Carlo sample from an EDF model

	Draws Nsamples in a field of radius
        fieldradius (in radians) at (l,b) (in radians) between magnitudes maglimits for
        band mag_band
    '''
    nsamples = int(Nsamples/nwalkers)+1

    if(with_halo):
        edf_sampling.turn_on_halo()
    else:
        edf_sampling.turn_off_halo()

    if(np.fabs(b)<fieldradius/4.):
        b+=fieldradius/2.

    # -- Initialise walkers
    ndim, nwalker = 7 + 2*(fieldradius>0), nwalkers
    # Parameters for log-likelihood
    check = lambda x: np.array([-10000.,10000.]) if x is None else x
    checkcolor = lambda x: np.array(['J','K']) if x is None else color
    llos_args =[np.array([l,b]), mag_band, maglimits, which_iso, RcPorZ,
                checkcolor(color), check(colorlimits), check(logglimits), check(Tefflimits),
                np.array([1*extinct,1*dered]), fieldradius,interp,
                np.array([break_pos,break_grad]),tgcolor_errs]
    ## Samplers
    sampler = emcee.EnsembleSampler(nwalker, ndim,los_magbox_LogL_sample_py,
                                    args=llos_args, threads=threads,a=asampler)
    if(pt):
        # Parallel tempering
        sampler = emcee.PTSampler(ntemp, nwalker, ndim,los_magbox_LogL_sample_py,
                                  logp, loglargs=llos_args, threads=threads)
    # Either sample in birth radius or metallicity -- birth radius sampling obsolete but good for checks
    RorZlim=[0.,18.]
    if(RcPorZ==False):
        RorZlim=[-1.,0.6]
        if not with_halo:
            RorZlim[0]=edf_sampling.minZ()
            RorZlim[1]=edf_sampling.maxZ()

    # Sample uniform in age, mass
    # Sample gaussian in Z and velocities
    # Sample exp in (unextincted) magnitude
    meanZ,sigZ = 0.,.4
    meanV,sigV = 0.,50.
    #     age  Z           mass  vr     vphi    vz   mag
    lo = [0.5, RorZlim[0], 0.5, -150., -200., -150., maglimits[0]]
    hi = [12., RorZlim[1], 2.,   150.,  200.,  150., maglimits[1]]
    if(fieldradius>0.):
        # also sample in l,b
        hi=np.concatenate((hi,[l,b]))
        lo=np.concatenate((lo,[l,b]))
    # Initial sample
    p0 = np.random.uniform(low=lo,high=hi,size=(nwalker*ntemp, ndim))
    p0.T[1]=np.random.normal(size=(nwalker*ntemp))*sigZ+meanZ # metallicity
    for i in range(3,6):
        # velocities
        p0.T[i]=np.random.normal(size=(nwalker*ntemp))*sigV+meanV
    p0.T[4]+=200. # shift vphi
    # magnitudes
    p0.T[6]=trunc_exp_rv(maglimits[0],maglimits[1],1.,size=len(p0))

    # -- Only choose physically allowed masses for selected age, Z
    for i in np.arange(len(p0)):
        maxmass,minmass=0.,0.
        check_highmass_fn = edf_sampling.check_highmass_Z
        check_lowmass_fn = edf_sampling.check_lowmass_Z
        if(RcPorZ==True):
            check_highmass_fn = edf_sampling.check_highmass
            check_lowmass_fn = edf_sampling.check_lowmass
        maxmass = check_highmass_fn(p0[i][0], p0[i][1],which_iso)
        minmass = check_lowmass_fn(p0[i][0], p0[i][1],which_iso)

        while(maxmass < 0. or minmass<0. or
              edf_sampling.check_radius_positive(p0[i][0], p0[i][1]) == 0):
            p0[i][1] = np.random.uniform(low=RorZlim[0], high=RorZlim[1])
            maxmass = check_highmass_fn(p0[i][0], p0[i][1],which_iso)
            minmass = check_lowmass_fn(p0[i][0], p0[i][1],which_iso)
        MAXMASS=3.
        if(maxmass>MAXMASS):
            maxmass=MAXMASS
        p0[i][2] = np.random.uniform(low=minmass, high=maxmass)
        # Check other cuts satisfied
        if(check(colorlimits)[0]>-10. or check(logglimits)[0]>-10. or check(Tefflimits)[0]>-10.):
            while(edf_sampling.check_color_logg_cut(p0[i][0], p0[i][1], p0[i][2],
                                                 0.,180.,checkcolor(color),
						 check(colorlimits),check(logglimits),check(Tefflimits),
						 which_iso,False)==0):
                p0[i][2] = np.random.uniform(low=minmass, high=maxmass)

    # Got reasonable starting points but now want to shuffle slightly to ensure
    err = [0.5,0.1,0.1,10.,10.,10.,0.1]
    # Shuffle by extinction if large
    max_distance=3.
    extinctBand = edf_sampling.get_extinct(l,b,max_distance,mag_band)
    if(extinct>err[-1]):
        err[-1]=extinctBand
    if(fieldradius>0.):
        err=np.concatenate((err,[0.,0.]))
    # Shuffle
    for i in np.arange(len(p0)):
        pp = los_magbox_LogL_sample_py(p0[i],*llos_args)
        n,maxn=0,10000
        while(np.isinf(pp) and n<maxn):
            p0[i]=p0[np.random.randint(len(p0))]+np.random.normal(size=ndim)*err
            pp = los_magbox_LogL_sample_py(p0[i],*llos_args)
            n+=1
        if(n==maxn):
            print 'No finite LogL found:', p0[i]
    ## Now sample l and b if needed
    if(fieldradius>0.):
        rad = np.random.uniform(size=len(p0))*fieldradius
        thet = np.random.uniform(size=len(p0))*2.*np.pi
        p0.T[7]=l+np.cos(thet)*rad/np.cos(b)
        p0.T[8]=b+np.sin(thet)*rad

    if(messages): print 'Initial points sampled'
    if(pt):
        p0=np.reshape(p0,(ntemp,nwalker,ndim))
    # -- Run a burn-in
    Nb=16000
    if(Nburn):
        Nb=Nburn
    else:
        Nb=nsamples*thin*2
    pos, prob, state = sampler.run_mcmc(p0, Nb, storechain=False)
    if(messages):
        print 'Number of logl=-inf = '+str(len(prob[np.isinf(prob)]))
        print("Burnt")
    sampler.reset()

    # -- Sample with thinning and calculating dependent variables
    pos, prob, state = sampler.run_mcmc(pos, nsamples*thin, thin=thin)
    flatchain = sampler.flatchain
    lnprob = sampler.lnprobability
    if(pt):
        flatchain=flatchain[0]
        lnprob=lnprob[0]

    extras = np.array(
        map(lambda i: edf_sampling.get_extra_data(i if fieldradius>0. else np.append(i,[l,b]),
                       mag_band, which_iso, RcPorZ, extinct,interp), flatchain))
    actions = np.array([edf_sampling.get_actions(np.concatenate((j[4:7], i[3:6])))
                        for i,j in izip(flatchain, extras)])

    nameslist = np.copy(names)
    if(RcPorZ==False):
        nameslist[1]="Z"
        nameslist[10]="RcP"
    nameslist[23]=mag_band
    nameslist[6]=mag_band+"0"

    LB = np.reshape(np.tile([l,b],len(flatchain)),(len(flatchain),2))
    everything = np.vstack((flatchain.T, LB.T, extras.T,
                 actions.T,lnprob.flatten())).T
    if fieldradius>0.:
        everything = np.vstack((flatchain.T, extras.T,
                 actions.T,lnprob.flatten())).T

    if color is not None:
        if extra_magnitudes is not None:
            extra_magnitudes=np.unique(np.concatenate((color,extra_magnitudes,np.array([mag_band]))))
        else:
            extra_magnitudes=np.unique(np.concatenate((color,np.array([mag_band]))))
        extra_magnitudes = extra_magnitudes[extra_magnitudes!=mag_band]

    if extra_magnitudes is not None:
        extra_mags = np.array(
        map(lambda i: edf_sampling.get_extra_magnitudes(i if fieldradius>0. else np.append(i,[l,b]),
                       mag_band, extra_magnitudes,which_iso, RcPorZ, extinct,interp), flatchain))
        nameslist=np.concatenate((nameslist,extra_magnitudes))
        print extra_magnitudes
        nameslist=np.concatenate((nameslist,np.array([e+'0' for e in list(extra_magnitudes)])))
        everything = np.vstack((everything.T, extra_mags.T)).T

    if(messages): print("Mean acceptance fraction:", np.mean(sampler.acceptance_fraction))
    df=pd.DataFrame(everything,columns=nameslist)
    df = df.sample(Nsamples,replace=False).reset_index(drop=True)
    if(output_file):
        df.to_csv(output_file)
    return df

from sklearn.neighbors import KDTree

def add_errors_to_sample(sample,data,
                         data_tree_fields = ['Teff_K','logg_K','Met_N_K_DR5'],
                         sample_tree_fields = ['Teff','logg','Z'],
                         err_fields = {'Teff':'eTeff_K',
                                       'logg':'elogg_K',
                                       'Z':'eMet_N_K',
                                       'vlos':'eHRV'},
                         leaf_size=10):
    '''
        Add errors to a mock sample
        sample -- sample catalogue
        data -- data
        data_tree_fields -- the data entries used to construct the tree
        sample_tree_fields -- the sample entries corresponding to the data used
                              to generate the tree
        err_fields -- a dictionary of the fields to add errors to in the keys
                      and the corresponding error fields in the data
    '''
    # Build tree
    data_tr = data[data_tree_fields]
    m,s = data_tr.mean(), data_tr.std()
    data_tr = (data_tr-m)/s
    kdt = KDTree(data_tr, leaf_size=leaf_size, metric='euclidean')
    RAVE_dd = (sample[sample_tree_fields].values-m.values)/s.values

    index = kdt.query(RAVE_dd, k=1, return_distance=False)
    for k,e in err_fields.items():
        sample['e'+k]=data[e].take(index.flatten()).values
        sample[k+'_e']=np.random.normal(size=len(sample))*sample['e'+k]+sample[k]

    return sample

def add_errors_to_sample_astrometry_covariance(sample,data,
                         data_tree_fields = ['BVmag','Vmag','l','b'],
                         sample_tree_fields = ['BV','V','l','b'],
                         err_fields = {'B':'Bmag','V':'Vmag'},
                         leaf_size=10):
    '''
        Add errors to a mock sample
        sample -- sample catalogue
        data -- data
        data_tree_fields -- the data entries used to construct the tree
        sample_tree_fields -- the sample entries corresponding to the data used
                              to generate the tree
        err_fields -- a dictionary of the fields to add errors to in the keys
                      and the corresponding error fields in the data
        Will also use tree fields to assign covariant astrometry errors
    '''
    # Build tree
    data_tr = data[data_tree_fields]
    m,s = data_tr.mean(), data_tr.std()
    data_tr = (data_tr-m)/s
    kdt = KDTree(data_tr, leaf_size=leaf_size, metric='euclidean')
    sample_dd = (sample[sample_tree_fields].values-m.values)/s.values

    sample['parallax'] = 1./sample['s']

    index = kdt.query(sample_dd, k=1, return_distance=False)
    ## First photometry
    for k,e in err_fields.items():
        sample['e'+k]=data[e].take(index.flatten()).values
        sample[k+'_e']=np.random.normal(size=len(sample))*sample['e'+k]+sample[k]
    ## Now astrometry
    cov = np.array([[data.parallax_error.values**2,
        data.parallax_pmra_corr.values*data.parallax_error.values*data.pmra_error.values,
        data.parallax_pmdec_corr.values*data.parallax_error.values*data.pmdec_error.values],
       [data.parallax_pmra_corr.values*data.parallax_error.values*data.pmra_error.values,
        data.pmra_error.values**2,
        data.pmra_pmdec_corr.values*data.pmra_error.values*data.pmdec_error.values],
       [data.parallax_pmdec_corr.values*data.parallax_error.values*data.pmdec_error.values,
        data.pmra_pmdec_corr.values*data.pmra_error.values*data.pmdec_error.values,
        data.pmdec_error.values**2]]).T
    #DD = np.array(map(lambda x: np.random.multivariate_normal(np.zeros(3),x),
    #        cov[index.flatten()]))
    DD = np.einsum('jik,jk->ji',np.linalg.cholesky(cov[index.flatten()]),
                                np.random.normal(size=(len(sample),3)))
    for i,p in enumerate(['parallax','pm_ra','pm_dec']):
        sample[p+'_e']=DD[:,i]+sample[p]
        sample['e'+p]=np.sqrt(cov[index.flatten()][:,i,i])
    sample['parallax_pmra_corr']=cov[index.flatten()][:,1,0]/np.sqrt(cov[index.flatten()][:,0,0]*cov[index.flatten()][:,1,1])
    sample['parallax_pmdec_corr']=cov[index.flatten()][:,2,0]/np.sqrt(cov[index.flatten()][:,0,0]*cov[index.flatten()][:,2,2])
    sample['pmra_pmdec_corr']=cov[index.flatten()][:,1,2]/np.sqrt(cov[index.flatten()][:,2,2]*cov[index.flatten()][:,1,1])
    return sample

def compute_velocities_actions_with_errors(sample):
    '''
        compute the velocities and actions of a sample with erroneous
        quantities.
        sample must have fields RA, DEC, s_e, vlos_e, pm_ra_e, pm_dec_e

        NEED TO RUN edf_sampling.setup first

    '''

    fields =['PMl_e','PMb_e','R_e','phi_e','z_e','vR_e','vphi_e','vz_e','JR_e','Lz_e','Jz_e','Rc_e']
    for f in fields:
        sample[f]=np.ones(len(sample))*-9999.

    for i in range(len(sample)):
        if(sample['s_e'][i]>0. and sample['s_e'][i]==sample['s_e'][i]):
            X = edf_sampling.process_data(
                np.array([sample['RA'][i],     sample['DEC'][i],
                          sample['s_e'][i],    sample['vlos_e'][i],
                          sample['pm_ra_e'][i],sample['pm_dec_e'][i]]))
            sample['PMl_e'][i],sample['PMb_e'][i]=X[4],X[5]
            sample['R_e'][i],sample['phi_e'][i],sample['z_e'][i]=X[6],X[7],X[8]
            sample['vR_e'][i],sample['vphi_e'][i],sample['vz_e'][i]=X[9],X[10],X[11]
            sample['JR_e'][i],sample['Lz_e'][i],sample['Jz_e'][i],sample['Rc'][i]=X[12],X[13],X[14],X[15]

    return sample

def compute_velocities_with_errors(sample):
    '''
        compute the velocities and actions of a sample with erroneous
        quantities.
        sample must have fields RA, DEC, s_e, vlos_e, pm_ra_e, pm_dec_e
    '''

    fields =['PMl_e','PMb_e','R_e','phi_e','z_e','vR_e','vphi_e','vz_e','JR_e','Lz_e','Jz_e','Rc_e']
    for f in fields:
        sample[f]=np.ones(len(sample))*-9999.

    for i in range(len(sample)):
        if(sample['s_e'][i]>0. and sample['s_e'][i]==sample['s_e'][i]):
            X = edf_sampling.process_data(
                np.array([sample['RA'][i],     sample['DEC'][i],
                          sample['s_e'][i],    sample['vlos_e'][i],
                          sample['pm_ra_e'][i],sample['pm_dec_e'][i]]))
            sample['PMl_e'][i],sample['PMb_e'][i]=X[4],X[5]
            sample['R_e'][i],sample['phi_e'][i],sample['z_e'][i]=X[6],X[7],X[8]
            sample['vR_e'][i],sample['vphi_e'][i],sample['vz_e'][i]=X[9],X[10],X[11]
            sample['JR_e'][i],sample['Lz_e'][i],sample['Jz_e'][i],sample['Rc'][i]=X[12],X[13],X[14],X[15]

    return sample


import sys, os
import time
import numpy as np
import torch
import torch.nn as nn
from torch.utils import data
from parsers import parse_a3m, read_templates
from TrunkModel  import TrunkModule
import util
from collections import namedtuple
from ffindex import *
from kinematics import xyz_to_c6d, c6d_to_bins2, xyz_to_t2d
from trFold import TRFold

NBIN = [37, 37, 37, 19]

MODEL_PARAM ={
        "n_module"     : 8,
        "n_module_str" : 4,
        "n_layer"      : 1,
        "d_msa"        : 384 ,
        "d_pair"       : 288,
        "d_templ"      : 64,
        "n_head_msa"   : 12,
        "n_head_pair"  : 8,
        "n_head_templ" : 4,
        "d_hidden"     : 64,
        "r_ff"         : 4,
        "n_resblock"   : 1,
        "p_drop"       : 0.0,
        "use_templ"    : True,
        "performer_N_opts": {"nb_features": 64},
        "performer_L_opts": {"nb_features": 64}
        }

SE3_param = {
        "num_layers"    : 3,
        "num_channels"  : 16,
        "num_degrees"   : 2,
        "l0_in_features": 32,
        "l0_out_features": 8,
        "l1_in_features": 3,
        "l1_out_features": 2,
        "num_edge_features": 32,
        "div": 2,
        "n_heads": 4
        }
MODEL_PARAM['SE3_param'] = SE3_param

# params for the folding protocol
fold_params = {
    "SG7"     : np.array([[[-2,3,6,7,6,3,-2]]])/21,
    "SG9"     : np.array([[[-21,14,39,54,59,54,39,14,-21]]])/231,
    "DCUT"    : 19.5,
    "ALPHA"   : 1.57,

    # TODO: add Cb to the motif
    "NCAC"    : np.array([[-0.676, -1.294,  0.   ],
                          [ 0.   ,  0.   ,  0.   ],
                          [ 1.5  , -0.174,  0.   ]], dtype=np.float32),
    "CLASH"   : 2.0,
    "PCUT"    : 0.5,
    "DSTEP"   : 0.5,
    "ASTEP"   : np.deg2rad(10.0),
    "XYZRAD"  : 7.5,
    "WANG"    : 0.1,
    "WCST"    : 0.1
}

fold_params["SG"] = fold_params["SG9"]

class Predictor():
    def __init__(self, model_dir=None, device="cuda:0"):
        if model_dir == None:
            self.model_dir = "%s/models"%(os.path.dirname(os.path.realpath(__file__)))
        else:
            self.model_dir = model_dir
        #
        # define model name
        self.model_name = "BFF"
        self.device = device
        self.active_fn = nn.Softmax(dim=1)

        # define model & load model
        self.model = TrunkModule(**MODEL_PARAM).to(self.device)
        could_load = self.load_model(self.model_name)
        if not could_load:
            print ("ERROR: failed to load model")
            sys.exit()

    def load_model(self, model_name, suffix='last'):
        chk_fn = "%s/%s_%s.pt"%(self.model_dir, model_name, suffix)
        print (chk_fn)
        if not os.path.exists(chk_fn):
            return False
        checkpoint = torch.load(chk_fn, map_location=self.device)
        self.model.load_state_dict(checkpoint['model_state_dict'])
        return True

    def predict(self, a3m_fn, out_prefix, hhr_fn=None, atab_fn=None, window=150, shift=50, n_latent=128):
        msa_orig = parse_a3m(a3m_fn)
        N, L = msa_orig.shape
        #
        if os.path.exists(hhr_fn):
            xyz_t, t1d, t0d = read_templates(L, ffdb, hhr_fn, atab_fn, n_templ=10)
        else:
            xyz_t = torch.full((1, L, 3, 3), np.nan).float()
            t1d = torch.zeros((1, L, 1)).float()
            t0d = torch.zeros((1,3)).float()
        #
        # template features
        xyz_t = xyz_t.float().unsqueeze(0)
        t1d = t1d.float().unsqueeze(0)
        t0d = t0d.float().unsqueeze(0)
        t2d = xyz_to_t2d(xyz_t, t0d)

        self.model.eval()
        for i_trial in range(10):
            self.run_prediction(msa_orig, t1d, t2d, "%s_%02d"%(out_prefix, i_trial), n_latent=n_latent)
            torch.cuda.empty_cache()
    def run_prediction(self, msa_orig, t1d, t2d, out_prefix, n_latent=128):
        N, L = msa_orig.shape
        with torch.no_grad():
            #
            msa = torch.tensor(msa_orig).long() # (N, L)
            random_idx = torch.randperm(N-1)
            msa = torch.cat([msa[:1,:], msa[1:,:][random_idx]], dim=0).view(1, N, L)
            #
            if n_latent < N:
                msa_latent = msa[:,:n_latent]
                msa_extra = msa[:,n_latent:,:]
            else:
                msa_latent = msa
                msa_extra = msa[:,:1,:]
            msa_extra = msa_extra[:,:10000]
            print (msa_latent.shape, msa_extra.shape, N, L)
            #
            idx_pdb = torch.arange(L).long().view(1, L)
            #
            msa_latent = msa_latent.to(self.device)
            msa_extra = msa_extra.to(self.device)
            seq = msa_latent[:,0]
            idx_pdb = idx_pdb.to(self.device)
            t1d = t1d.to(self.device)
            t2d = t2d.to(self.device)
            with torch.cuda.amp.autocast(enabled=True):
                logit_s, init_crds, pred_lddt = self.model(msa_latent, msa_extra, seq, idx_pdb, t1d=t1d, t2d=t2d)
            prob_s = list()
            for logit in logit_s:
                prob = self.active_fn(logit.float()) # distogram
                prob = prob.reshape(-1, L, L) #.permute(1,2,0).cpu().numpy()
                prob_s.append(prob)

        for prob in prob_s:
            prob += 1e-8
            prob = prob / torch.sum(prob, dim=0)[None]
        self.write_pdb(seq[0], init_crds[0], Bfacts=pred_lddt[0], prefix="%s_init"%(out_prefix))
        xyz = init_crds[0, :, 1] # initial ca coordinates
        TRF = TRFold(prob_s, fold_params)
        xyz = TRF.fold(xyz, batch=45, lr=0.1, nsteps=200)
        self.write_pdb(seq[0], xyz, prefix="%s"%(out_prefix), Bfacts=pred_lddt[0])

        prob_s = [prob.permute(1,2,0).detach().cpu().numpy().astype(np.float16) for prob in prob_s]
        np.savez_compressed("%s.npz"%(out_prefix), dist=prob_s[0].astype(np.float16), \
                            omega=prob_s[1].astype(np.float16),\
                            theta=prob_s[2].astype(np.float16),\
                            phi=prob_s[3].astype(np.float16))

    def write_pdb(self, seq, atoms, Bfacts=None, prefix=None):
        L = len(seq)
        filename = "%s.pdb"%prefix
        ctr = 1
        with open(filename, 'wt') as f:
            if Bfacts == None:
                Bfacts = np.zeros(L)
            else:
                Bfacts = torch.clamp( Bfacts, 0, 1)

            for i,s in enumerate(seq):
                if (len(atoms.shape)==2):
                    f.write ("%-6s%5s %4s %3s %s%4d    %8.3f%8.3f%8.3f%6.2f%6.2f\n"%(
                            "ATOM", ctr, " CA ", util.num2aa[s],
                            "A", i+1, atoms[i,0], atoms[i,1], atoms[i,2],
                            1.0, Bfacts[i] ) )
                    ctr += 1

                elif atoms.shape[1]==3:
                    for j,atm_j in enumerate((" N  "," CA "," C  ")):
                        f.write ("%-6s%5s %4s %3s %s%4d    %8.3f%8.3f%8.3f%6.2f%6.2f\n"%(
                                "ATOM", ctr, atm_j, util.num2aa[s],
                                "A", i+1, atoms[i,j,0], atoms[i,j,1], atoms[i,j,2],
                                1.0, Bfacts[i] ) )
                        ctr += 1

def get_args():
    #DB="/home/robetta/rosetta_server_beta/external/databases/trRosetta/pdb100_2021Mar03/pdb100_2021Mar03"
    DB = "/projects/ml/TrRosetta/pdb100_2020Mar11/pdb100_2020Mar11"
    import argparse
    parser = argparse.ArgumentParser(description="RoseTTAFold: Protein structure prediction with 3-track attentions on 1D, 2D, and 3D features")
    parser.add_argument("a3m_fn", help="input MSA in a3m format")
    parser.add_argument("out_prefix", help="prefix for output file. [out_prefix].npz file will be generated")
    parser.add_argument("hhr_fn",
                        help="HHsearch result file in hhr format")
    parser.add_argument("atab_fn",
                        help="HHsearch result file in atab format")
    parser.add_argument("-db", default=DB, required=False,
                        help="HHsearch database [%s]"%DB)
    args = parser.parse_args()
    return args

if __name__ == "__main__":
    args = get_args()
    if not os.path.exists("%s.npz"%args.out_prefix):
        FFDB = args.db
        FFindexDB = namedtuple("FFindexDB", "index, data")
        ffdb = FFindexDB(read_index(FFDB+'_pdb.ffindex'),
                         read_data(FFDB+'_pdb.ffdata'))
        pred = Predictor()
        pred.predict(args.a3m_fn, args.out_prefix, args.hhr_fn, args.atab_fn)