<reponame>berkalpay/EVcouplings
"""
Functions for handling evolutionary couplings data.

.. todo::

    1. clean up
    2. add Pompom score
    3. add mapping tools (multidomain, complexes)
    4. ECs to matrix
    5. APC on subsets of positions (e.g. for complexes)

Authors:
  <NAME>
  <NAME> (original mixture model code)
  <NAME> (skew normal mixture model)
  <NAME> (EVComplex Score code)
"""

from math import ceil
from copy import deepcopy
from pkg_resources import resource_filename

import numpy as np
import pandas as pd
import matplotlib.pyplot as plt
import scipy.optimize as op
from scipy import stats
from sklearn.linear_model import LogisticRegression

from evcouplings.utils.calculations import median_absolute_deviation
from evcouplings.utils.config import read_config_file


def read_raw_ec_file(filename, sort=True, score="cn"):
    """
    Read a raw EC file (e.g. from plmc) and sort
    by scores

    Parameters
    ----------
    filename : str
        File containing evolutionary couplings
    sort : bool, optional (default: True)
        If True, sort pairs by coupling score in
        descending order
    score_column : str, optional (default: True)
        Score column to be used for sorting

    Returns
    -------
    ecs : pd.DataFrame
        Table of evolutionary couplings

    """
    ecs = pd.read_csv(
        filename, sep=" ",
        names=["i", "A_i", "j", "A_j", "fn", "cn"]
    )

    if sort:
        ecs = ecs.sort_values(
            by=score, ascending=False
        )

    return ecs


def enrichment(ecs, num_pairs=1.0, score="cn", min_seqdist=6):
    """
    Calculate EC "enrichment" as first described in
    Hopf et al., Cell, 2012.

    .. todo::

        Make this handle segments if they are in EC table

    Parameters
    ----------
    ecs : pd.DataFrame
        Dataframe containing couplings
    num_pairs : int or float, optional (default: 1.0)
        Number of ECs to use for enrichment calculation.
        - If float, will be interpreted as fraction of the
        length of the sequence (e.g. 1.0*L)
        - If int, will be interpreted as
        absolute number of pairs
    score : str, optional (default: cn)
        Pair coupling score used for calculation
    min_seqdist : int, optional (default: 6)
        Minimum sequence distance of couplings that will
        be included in the calculation

    Returns
    -------
    enrichment_table : pd.DataFrame
        Sorted table with enrichment values for each
        position in the sequence
    """
    # determine how many positions ECs are over
    pos = set(ecs.i.unique()) | set(ecs.j.unique())
    num_pos = len(pos)

    # calculate absolute number of pairs if
    # fraction of length is given
    if isinstance(num_pairs, float):
        num_pairs = int(ceil(num_pairs * num_pos))

    # get longrange ECs and sort by score
    sorted_ecs = ecs.query(
        "abs(i-j) >= {}".format(min_seqdist)
    ).sort_values(
        by=score, ascending=False
    )

    # select top EC pairs
    top_ecs = sorted_ecs.iloc[0:num_pairs]

    # stack dataframe so it contains each
    # EC twice as forward and backward pairs
    # (i, j) and (j, i)
    flipped = top_ecs.rename(
        columns={"i": "j", "j": "i", "A_i": "A_j", "A_j": "A_i"}
    )

    stacked_ecs = top_ecs.append(flipped)

    # now sum cumulative strength of EC for each position
    ec_sums = pd.DataFrame(
        stacked_ecs.groupby(["i", "A_i"]).sum()
    )

    # average EC strength for top ECs
    avg_degree = top_ecs.loc[:, score].sum() / len(top_ecs)

    # "enrichment" is ratio how much EC strength on
    # an individual position exceeds average strength in top
    ec_sums.loc[:, "enrichment"] = ec_sums.loc[:, score] / avg_degree

    e = ec_sums.reset_index().loc[:, ["i", "A_i", "enrichment"]]
    return e.sort_values(by="enrichment", ascending=False)


class LegacyScoreMixtureModel:
    """
    Assign to each EC score the probability of being in the
    lognormal tail of a normal-lognormal mixture model.

    .. note::
        this is the original version of the score mixture model
        with a normal distribution noise component, this has been
        superseded by a model using a skew normal distribution
    """
    def __init__(self, x, clamp_mu=False, max_fun=10000, max_iter=1000):
        """
        Mixture model of evolutionary coupling scores to
        determine significant scores that are in high-scoring,
        positive tail of distribution.

        Parameters
        ----------
        x : np.array (or list-like)
            EC scores from which to infer the mixture model
        clamp_mu : bool, optional (default: False)
            Fix mean of Gaussian component to 0 instead of
            fitting it based on data
        max_fun : int
            Maximum number of function evaluations
        max_iter : int
            Maximum number of iterations
        """
        x = np.array(x)

        # Infer parameters of mixture model
        self.params = self._learn_params(x, clamp_mu, max_fun, max_iter)

    @classmethod
    def _learn_params(cls, x, clamp_mu, max_fun, max_iter):
        """
        Infer parameters of mixture model.

        Parameters
        ----------
        x : np.array (or list-like)
            EC scores from which to infer the mixture model
        clamp_mu : bool, optional (default: False)
            Fix mean of Gaussian component to 0 instead of
            fitting it based on data
        max_fun : int
            Maximum number of function evaluations
        max_iter : int
            Maximum number of iterations

        Returns
        -------
        mu : float
            Mean of normal distribution
        sigma : float
            Standard deviation of normal distribution
        q : float
            Relative weights of each distribution
        logmu : float
            Mean of lognormal distribution
        logsigma : float
            Standard deviation of lognormal distribution
        """
        # Initial starting parameters for normal
        # and lognormal distributions
        # q: relative contribution of each distribtuion
        mu = 0
        sigma = np.std(x)
        q = 1
        logsigma = 0.4
        logmu = np.percentile(x, 75) - logsigma**2 / 2
        param = np.array([mu, sigma, q, logmu, logsigma])

        # Target function for minimization
        def target_func(params):
            if clamp_mu:
                params[0] = 0

            return -np.sum(np.log(cls._gaussian_lognormal(x, params)))

        # Minimize function
        coeff = op.fmin(
            target_func, param, maxfun=max_fun, maxiter=max_iter, disp=False
        )

        # If clamping mu, also set to 0 in the end, so this value
        # is used for probability calculations
        if clamp_mu:
            coeff[0] = 0

        q = coeff[2]
        # Check if fit worked
        if q >= 1 or np.isinf(q) or np.isneginf(q):
            raise ValueError(
                "No tail, fit failed. q={}".format(q)
            )

        return coeff

    @classmethod
    def _gaussian_lognormal(cls, x, params):
        """
        Gaussian-lognormal mixture probability
        density function.

        Parameters
        ----------
        x : np.array
            Scores for which PDF is calculated
        params : tuple
            Parameters of lognormal-Gaussian mixture
            (mu, sigma, class weight q, loglog, logsigma)

        Returns
        -------
        np.array
            Probabilities
        """
        return cls._gaussian(x, params) + cls._lognormal(x, params)

    @classmethod
    def _gaussian(cls, x, params):
        """
        Normal probability density (multiplied
        by class weight).

        Parameters
        ----------
        x : np.array
            Scores for which PDF is calculated
        params : tuple
            Parameters of lognormal-Gaussian mixture
            (mu, sigma, class weight q, loglog, logsigma)

        Returns
        -------
        np.array
            Probabilities
        """
        mu, sigma, q, logmu, logsigma = params
        return q * stats.norm.pdf(x, loc=mu, scale=sigma)

    @classmethod
    def _lognormal(cls, x, params):
        """
        Log normal probability density (multiplied
        by class weight).

        Parameters
        ----------
        x : np.array
            Scores for which PDF is calculated
        params : tuple
            Parameters of lognormal-Gaussian mixture
            (mu, sigma, class weight q, loglog, logsigma)

        Returns
        -------
        np.array
            Probabilities
        """
        mu, sigma, q, logmu, logsigma = params

        # only assign probability to defined (i.e. positive) values,
        # set all others to zero
        prob = np.zeros(len(x))
        xpos = x > 0
        prob[xpos] = (1 - q) * stats.norm.pdf(
            np.log(x[xpos]), loc=logmu, scale=logsigma
        ) / x[xpos]

        return prob

    def probability(self, x, plot=False):
        """
        Calculate posterior probability of EC pair to
        be located in positive (lognormal) tail of the
        distribution.

        Parameters
        ----------
        x : np.array (or list-like)
            List of scores
        plot : bool, optional (default: False)
            Plot score distribution and probabilities

        Returns
        -------
        posterior : np.array(float)
            Posterior probability of being in signal
            component of mixture model
        """
        x = np.array(x)
        p_lognormal = self._lognormal(x, self.params)
        p_gaussian = self._gaussian(x, self.params)

        posterior = p_lognormal / (p_lognormal + p_gaussian)

        if plot:
            plt.figure(figsize=(12, 8))
            # fig = plt.figure(figsize=(4,4))
            c = "#fdc832"
            n_ECs, edges = np.histogram(x, 1000, density=True)
            mid = []
            for l, r in zip(edges[:-1], edges[1:]):
                mid.append((l + r) / 2)

            plt.plot(
                mid, n_ECs, '-', color=c, markerfacecolor=c,
                markeredgecolor='None', linewidth=1
            )
            plt.plot(x, posterior, '-k', linewidth=2)
            plt.plot(x, p_lognormal, 'r', linewidth=1)
            plt.plot(x, p_gaussian, 'b', linewidth=1)

            takeover = x[p_lognormal > p_gaussian].min()
            plt.axvline(takeover, color="grey")
            plt.axvline(x[posterior > 0.99].min(), color="grey", lw=1)

            pompom = abs(x.min())
            plt.axvline(-pompom, color="grey", ls="--")
            plt.axvline(pompom, color="grey", ls="--")

            plt.xlabel("EC scores")
            plt.ylabel("PDF")

        return posterior


class ScoreMixtureModel:
    """
    Assign to each EC score the probability of being in the
    lognormal tail of a skew normal-lognormal mixture model.
    """
    def __init__(self, x):
        """
        Mixture model of evolutionary coupling scores to
        determine signifcant scores that are in high-scoring,
        positive tail of distribution.

        Parameters
        ----------
        x : np.array (or list-like)
            EC scores from which to infer the mixture model
        """
        x = np.array(x)

        # Infer parameters of mixture model
        self.params = self._learn_params(x)

    @classmethod
    def skewnorm_pdf(cls, x, location, scale, skew):
        """
        Probability density of skew normal distribution
        (noise component)

        Parameters
        ---------
        x : np.array(float)
            Data for which probability density should be calculated
        location : float
            Location parameter of skew normal distribution
        scale : float
            Scale parameter of skew normal distribution
        skew : float
            Skew parameter of skew normal distribution

        Returns
        ------
        density : np.array(float)
            Probability density for input array x
        """
        x_transform = (x - location) / scale
        density = (
            2 / scale *
            stats.norm.pdf(x_transform) *
            stats.norm.cdf(skew * x_transform)
        )

        return density

    @classmethod
    def lognorm_pdf(cls, x, logmu, logsig):
        """
        Probability density of lognormal distribution
        (signal component)

        Parameters
        ---------
        x : np.array(float)
            Data for which probability density should be calculated
        logmu : float
            Location of lognormal distribution (signal component)
        logsig : float
            Scale parameter of lognormal distribution (signal component)

        Returns
        ------
        density : np.array(float)
            Probability density for input array x
        """
        density = np.zeros(len(x))
        xpos = x > 0
        density[xpos] = stats.norm.pdf(
            np.log(x[xpos]), loc=logmu, scale=logsig
        ) / x[xpos]

        return density

    @classmethod
    def skewnorm_constraint(cls, scale, skew):
        """
        Given scale and skew, returns location parameter to yield mean zero

        Parameters
        ----------
        scale : float
            Scale parameter of skew normal distribution
        skew : float
            Skew parameter of skew normal distribution

        Returns
        -------
        location : float
            Location parameter of skew normal distribution s.t.
            mean of distribution is equal to 0
        """
        location = -scale * skew / np.sqrt(1 + skew**2) * np.sqrt(2 / np.pi)
        return location

    @classmethod
    def mixture_pdf(cls, x, p, scale, skew, logmu, logsig):
        """
        Compute mixture probability

        Parameters
        ----------
        x : np.array(float)
            Data for which probability density should be calculated
        p : float
            Mixing fraction between components for noise component
            (signal component will be 1-p)
        scale : float
            Scale parameter of skew normal distribution (noise component)
        skew : float
            Skew parameter of skew normal distribution (noise component)
        logmu : float
            Location of lognormal distribution (signal component)
        logsig : float
            Scale parameter of lognormal distribution (signal component)

        Returns
        -------
        density : np.array(float)
            Probability density for input array x
        """
        location = cls.skewnorm_constraint(scale, skew)
        density = (
            p * cls.skewnorm_pdf(x, location, scale, skew) +
            (1 - p) * cls.lognorm_pdf(x, logmu, logsig)
        )

        return density

    @classmethod
    def posterior_signal(cls, x, p, scale, skew, logmu, logsig):
        """
        Compute posterior probability of being in signal component

        Parameters
        ----------
        x : np.array(float)
            Data for which probability density should be calculated
        p : float
            Mixing fraction between components for noise component
            (signal component will be 1-p)
        scale : float
            Scale parameter of skew normal distribution (noise component)
        skew : float
            Skew parameter of skew normal distribution (noise component)
        logmu : float
            Location of lognormal distribution (signal component)
        logsig : float
            Scale parameter of lognormal distribution (signal component)

        Returns
        -------
        posterior : np.array(float)
            Posterior probability of being in signal component
            for input array x
        """
        P = cls.mixture_pdf(x, p, scale, skew, logmu, logsig)
        posterior = np.zeros(P.shape)
        f2 = cls.lognorm_pdf(x, logmu, logsig)
        posterior[x > 0] = (1 - p) * f2[x > 0] / P[x > 0]

        return posterior

    @classmethod
    def _learn_params(cls, x):
        """
        Infer parameters of mixture model.

        Parameters
        ----------
        x : np.array(float)
            EC scores from which to infer the mixture model

        Returns
        -------
        theta : np.array(float)
            Array with inferred parameters of model
            (mixing fraction, skew-normal scale, skew-normal skew,
            log-normal mean, log-normal stddev)
        """
        # mixing fraction, skew-normal scale, skew-normal skew,
        # log-normal mean, log-normal stddev
        theta = np.array(
            [0.5, np.std(x), 0, np.log(np.max(x)), 0.1]
        )

        def loglk_fun(params):
            return np.sum(np.log(cls.mixture_pdf(x, *params)))

        loglk = loglk_fun(theta)
        delta_loglk = 100
        max_iter = 200
        cur_iter = 0
        tolerance = 0.0001

        while delta_loglk > tolerance and cur_iter < max_iter:
            # E step
            z = 1 - cls.posterior_signal(x, *theta)

            # M step
            # MLE of the mixing fraction is the mean z
            theta[0] = np.mean(z)

            # Log-normal component
            # MLE is the z-weighted mean and std deviation of the log-scores
            pos_ix = x > 0
            z_complement = 1 - z[pos_ix]
            log_score = np.log(x[pos_ix])
            theta[3] = np.sum(z_complement * log_score) / np.sum(z_complement)
            theta[4] = np.sqrt(
                np.sum(z_complement * (log_score - theta[3])**2) / z_complement.sum()
            )

            # Skew-normal distribution
            # MLE requires numerical optimization
            def objfun(params):
                return -np.sum(
                    z * np.log(
                        cls.skewnorm_pdf(
                            x, cls.skewnorm_constraint(params[0], params[1]),
                            params[0], params[1]
                        )
                    )
                )

            theta[1:3] = op.fmin(objfun, theta[1:3], disp=False)

            # Test for EM convergence
            loglk_new = loglk_fun(theta)
            delta_loglk = loglk_new - loglk
            loglk = loglk_new

            # status update
            cur_iter += 1

        return theta

    def probability(self, x, plot=False):
        """
        Calculate posterior probability of EC pair to
        be located in positive (lognormal) tail of the
        distribution.

        Parameters
        ----------
        x : np.array (or list-like)
            List of scores

        Returns
        -------
        posterior : np.array(float)
            Posterior probability of being in signal
            component of mixture model
        """
        posterior = self.posterior_signal(x, *self.params)

        if plot:
            plt.hist(x, normed=True, bins=50, color="k")
            plt.plot(x, self.mixture_pdf(x, *self.params), color="r", lw=3)
            plt.plot(x, posterior, color="gold", lw=3)

        return posterior


class EVComplexScoreModel:
    """
    Assign to each EC score a (unnormalized) EVcomplex score as
    described in Hopf, Schärfe et al. (2014).

    TODO: this implementation currently does not take into account
    score normalization for the number of sequences and length of
    the model
    """
    def __init__(self, x):
        """
        Initialize EVcomplex score model
        
        Parameters
        ----------
        x : np.array (or list-like)
            EC scores from which to infer the mixture model
        """
        self.x = np.array(x)

    def probability(self, x, plot=False):
        """
        Calculates evcomplex score as cn_score / min_cn_score.
        TODO: plotting functionality not yet implemented

        Parameters
        ----------
        x : np.array (or list-like)
            List of scores
        plot: bool, optional (default: False)
            Plot score distribution

        Returns
        -------
        probability: np.array(float)
            EVcomplex score
        """
        # Calculate the minimum score
        min_score = abs(np.min(self.x))

        return x / min_score


def add_mixture_probability(ecs, model="skewnormal", score="cn",
                            clamp_mu=False, plot=False):
    """
    Add lognormal mixture model probability to EC table.

    Parameters
    ----------
    ecs : pd.DataFrame
        EC table with scores
    model : {"skewnormal", "normal"}, optional (default: skewnormal)
        Use model with skew-normal or normal distribution
        for the noise component of mixture model
    score : str, optional (default: "cn")
        Score on which mixture model will be based
    clamp_mu : bool, optional (default: False)
        Fix mean of Gaussian component to 0 instead of
        fitting it based on data
    plot : bool, optional (default: False)
        Plot score distribution and probabilities

    Returns
    -------
    ec_prob : pd.DataFrame
        EC table with additional column "probability"
        that for each EC contains the posterior
        probability of belonging to the lognormal
        tail of the distribution.
    """
    ec_prob = deepcopy(ecs)

    # learn mixture model
    if model == "skewnormal":
        mm = ScoreMixtureModel(ecs.loc[:, score].values)
    elif model == "normal":
        mm = LegacyScoreMixtureModel(ecs.loc[:, score].values, clamp_mu)
    elif model == "evcomplex":
        mm = EVComplexScoreModel(ecs.loc[:, score].values)
    else:
        raise ValueError(
            "Invalid model selection, valid options are: "
            "skewnormal, normal, evcomplex"
        )

    # assign probability
    ec_prob.loc[:, "probability"] = mm.probability(
        ec_prob.loc[:, score].values, plot=plot
    )

    return ec_prob


def logreg_classifier_to_dict(classifier, feature_names=None):
    """
    Serialize sklearn logistic regression classifier

    Inspired by https://stackoverflow.com/questions/48328012/python-scikit-learn-to-json

    Parameters
    ----------
    classifier : sklearn.linear_model.LogisticRegression
        Logistic regression classifier to be serialized
    feature_names : list(str)
        Feature names of coefficients in classifier, must
        be in the same order as classifier.coef_

    Returns
    -------
    model : dict
        Serialized classifier
    """
    params = {
        "classifier_settings": classifier.get_params(),
        "model_settings": {}
    }

    for attrib in ["classes_", "intercept_", "coef_", "n_iter_"]:
        params["model_settings"][attrib] = getattr(classifier, attrib).tolist()

    if feature_names is not None:
        params["feature_names"] = feature_names

    return params


def logreg_classifier_from_dict(params):
    """
    Deserialize model parameters into sklearn LogisticRegression classifier

    Inspired by https://stackoverflow.com/questions/48328012/python-scikit-learn-to-json

    Parameters
    -------
    model : dict
        Serialized classifier

    Parameters
    ----------
    classifier : sklearn.linear_model.LogisticRegression
        Deserialized logistic regression classifier
    feature_names : list(str)
        Feature names of coefficients in classifier
        (in same order as classifier.coef_)
    """
    classifier = LogisticRegression(
        **params["classifier_settings"]
    )

    for attrib, values in params["model_settings"].items():
        setattr(classifier, attrib, np.array(values))

    feature_names = params.get("feature_names")

    return classifier, feature_names


def add_freqs_to_ec_table(ecs, freqs):
    """
    Add residue and gap frequency as well as conservation info
    to data table by merging EC table and frequencies table

    Parameters
    ----------
    ecs : pd.DataFrame
        EC table generated by pipeline (e.g. _ECs.txt, _CouplingsScores.csv)
    freqs : pd.DataFrame
        Frequency table generated by pipeline (_frequencies.csv) using
        evcouplings.align.protocol.describe_frequencies()
    """
    # rename column names to be more descriptive
    freqs = freqs.rename(
        columns={
            "-": "gap_i",
            "conservation": "cons_i",
        }
    )

    # extract frequency of target residue for each position
    freqs.loc[:, "freq_i"] = [r[r["A_i"]] for idx, r in freqs.iterrows()]

    # limit to relevant columns and create second version
    # for mapping second position suffixed with _j
    freqs_sel_i = freqs[["i", "A_i", "freq_i", "gap_i", "cons_i"]]

    freqs_sel_j = freqs_sel_i.rename(columns={
        c: c.replace("i", "j") for c in freqs_sel_i.columns
    })

    # add frequency/conservation information to table
    ecs_with_freqs = ecs.merge(
        freqs_sel_i, on=["i", "A_i"]
    ).merge(
        freqs_sel_j, on=["j", "A_j"]
    )

    # check nothing went wrong and we kept full table
    assert len(ecs_with_freqs) == len(ecs)

    return ecs_with_freqs


def mad_outlier_score(x):
    """
    Median absolute deviation outlier scoring of
    ECs in tail vs. background noise distribution
    (robust z-scoring using median and median
    absolute deviation instead of mean and
    sample standard deviation)

    Parameters
    ----------
    x : np.array
        List of scores (note this must contain *all*
        scores for estimation of median and median
        absolute deviation, i.e. a list of top X
        ECs only will give wrong results)

    Returns
    -------
    np.array
        Outlier scores
    """
    med = np.median(x)
    mad = median_absolute_deviation(x)
    return (x - med) / mad


class LogisticRegressionScorer:
    """
    Rescore ECs based on logistic regression model
    fitted to large set of runs
    """
    def __init__(self, logreg_model_file=None, min_n_eff_over_l=0.375):
        """
        Create new logistic regression-based
        EC rescorer

        Parameters
        ----------
        logreg_model_file : str, optional (default: None)
            Specify path to yml file with logistic regression
            model parameters; if None, will use default
            model included with package
            (evcouplings/couplings/scoring_models/logistic_regression_all.yml)
        min_n_eff_over_l : float, optional (default: 0.3)
            Minimum number of effective sequences per model site required
            for rescoring to be applied; otherwise standard score will
            be returned and all probabilities will be set to 0. The
            default value will be divided by theta for the rescored run,
            the default of 0.375 derives from N_eff/L = 0.3 at theta = 0.8
        """
        # by default load internal classifier included with package
        if logreg_model_file is None:
            logreg_model_file = resource_filename(
                __name__, "scoring_models/logistic_regression_all.yml"
            )

        # load classifier from param file
        logreg_model_serialized = read_config_file(logreg_model_file)

        # deserialize and store classifier
        self.classifier, self.feature_names = logreg_classifier_from_dict(
            logreg_model_serialized
        )

        # store min N_eff/L requirement
        self.min_n_eff_over_l = min_n_eff_over_l

    @classmethod
    def _create_full_data_table(cls, ecs, freqs, theta, effective_sequences, num_sites):
        """
        Create full EC feature table

        Parameters
        ----------
        ecs : pd.DataFrame
            Full, unfiltered EC table
        freqs : pd.DataFrame
            Frequencies table
        theta : float
            Theta parameter used for sequence reweighting
            (used to adjust influence of different thetas in scoring model)
        effective_sequences : float
            Number of effective sequences after clustering and downweighting
        num_sites : float, optional (default: None)
            Number of sites in model / alignment; if None, will
            be automatically inferred from Ec table

        Returns
        -------
        pd.DataFrame
            Full EC table with annotated input features for model
        """
        meff_over_l = effective_sequences / num_sites
        meff_over_l2 = effective_sequences / num_sites ** 2

        meff_over_l_norm = meff_over_l / theta
        meff_over_l2_norm = meff_over_l2 / theta

        # add frequency and conservation info to EC table
        ecs = add_freqs_to_ec_table(
            ecs, freqs
        )

        # assign additional features to EC table
        # needed by prediction model
        ecs = ecs.assign(
            num_sites_log=np.log10(num_sites),
            min_gap=np.minimum(ecs.gap_i, ecs.gap_j),
            max_gap=np.maximum(ecs.gap_i, ecs.gap_j),
            min_cons=np.minimum(ecs.cons_i, ecs.cons_j),
            max_cons=np.maximum(ecs.cons_i, ecs.cons_j),
            meff_over_l_norm_log=np.log10(meff_over_l_norm),
            meff_over_l2_norm_log=np.log10(meff_over_l2_norm)
        )

        return ecs

    def score(self, ecs, freqs, theta, effective_sequences, num_sites=None, score="cn"):
        """
        Rescore EC table (must be full, unfiltered table) using logistic regression
        based probabilistic scoring. If number of effective sequences is too low for reliable
        scoring, score will default to input score and probabilities will all be set to 0.

        Parameters
        ----------
        ecs : pd.DataFrame
            Full, unfiltered EC table
        freqs : pd.DataFrame
            Frequencies table
        theta : float
            Theta parameter used for sequence reweighting
            (used to adjust influence of different thetas in scoring model)
        effective_sequences : float
            Number of effective sequences after clustering and downweighting
        num_sites : float, optional (default: None)
            Number of sites in model / alignment; if None, will
            be automatically inferred from Ec table
        score : string, optional (default: "cn")
            Input score to be input into logistic regression model.
            Model has been trained on CN scores only, so is not applicable
            to other types of scoring like DI score

        Returns
        -------
        ecs_final : pd.DataFrame
            Rescored ECs, same table as ecs but with three
            additional columns (mad_score, probability, score) containing
            1) median absolute deviation outlier score, 2) logistic regression
            model probability for EC to be true EC, and 3) pair score
            derived from decision function of logistic regression model
            (this can reorder the original ECs based on positional conservation
            and gap frequencies).
            Final table is sorted by "score" in descending order.
        """
        # infer num_sites from EC table if not defined
        if num_sites is None:
            num_sites = len(
                set(ecs.i.unique()) | set(ecs.j.unique())
            )

        # if N_eff/L is too low for model to give reliable results,
        # keep standard score and set all probabilities to 0
        if effective_sequences / num_sites / theta < self.min_n_eff_over_l:
            return ecs.assign(
                score=ecs[score],
                probability=0
            )

        # add medium absolute deviation outlier score to copy of EC table
        ecs = ecs.assign(
            mad_score=mad_outlier_score(ecs[score])
        )

        # add additional information to table
        ecs_full = self._create_full_data_table(
            ecs, freqs, theta, effective_sequences, num_sites
        )

        # extract features for classifier ("X")
        feature_table = ecs_full.reindex(self.feature_names, axis=1)

        # predict probabilities of contact and decision function
        # (target class/true EC := 1)
        probs = self.classifier.predict_proba(feature_table)[:, 1]
        decision_func = self.classifier.decision_function(feature_table)

        # assign to EC table
        ecs_final = ecs_full.assign(
            score=decision_func,
            probability=probs
        ).sort_values(
            by="score", ascending=False
        )

        # remove temporary feature columns and return updated EC table
        return ecs_final[
            list(ecs.columns) + ["probability", "score"]
        ]

"""
Functions for handling evolutionary couplings data.

.. todo::

    1. clean up
    2. add Pompom score
    3. add mapping tools (multidomain, complexes)
    4. ECs to matrix
    5. APC on subsets of positions (e.g. for complexes)

Authors:
  <NAME>
  <NAME> (original mixture model code)
  <NAME> (skew normal mixture model)
  <NAME> (EVComplex Score code)
"""

from math import ceil
from copy import deepcopy
from pkg_resources import resource_filename

import numpy as np
import pandas as pd
import matplotlib.pyplot as plt
import scipy.optimize as op
from scipy import stats
from sklearn.linear_model import LogisticRegression

from evcouplings.utils.calculations import median_absolute_deviation
from evcouplings.utils.config import read_config_file


def read_raw_ec_file(filename, sort=True, score="cn"):
    """
    Read a raw EC file (e.g. from plmc) and sort
    by scores

    Parameters
    ----------
    filename : str
        File containing evolutionary couplings
    sort : bool, optional (default: True)
        If True, sort pairs by coupling score in
        descending order
    score_column : str, optional (default: True)
        Score column to be used for sorting

    Returns
    -------
    ecs : pd.DataFrame
        Table of evolutionary couplings

    """
    ecs = pd.read_csv(
        filename, sep=" ",
        names=["i", "A_i", "j", "A_j", "fn", "cn"]
    )

    if sort:
        ecs = ecs.sort_values(
            by=score, ascending=False
        )

    return ecs


def enrichment(ecs, num_pairs=1.0, score="cn", min_seqdist=6):
    """
    Calculate EC "enrichment" as first described in
    Hopf et al., Cell, 2012.

    .. todo::

        Make this handle segments if they are in EC table

    Parameters
    ----------
    ecs : pd.DataFrame
        Dataframe containing couplings
    num_pairs : int or float, optional (default: 1.0)
        Number of ECs to use for enrichment calculation.
        - If float, will be interpreted as fraction of the
        length of the sequence (e.g. 1.0*L)
        - If int, will be interpreted as
        absolute number of pairs
    score : str, optional (default: cn)
        Pair coupling score used for calculation
    min_seqdist : int, optional (default: 6)
        Minimum sequence distance of couplings that will
        be included in the calculation

    Returns
    -------
    enrichment_table : pd.DataFrame
        Sorted table with enrichment values for each
        position in the sequence
    """
    # determine how many positions ECs are over
    pos = set(ecs.i.unique()) | set(ecs.j.unique())
    num_pos = len(pos)

    # calculate absolute number of pairs if
    # fraction of length is given
    if isinstance(num_pairs, float):
        num_pairs = int(ceil(num_pairs * num_pos))

    # get longrange ECs and sort by score
    sorted_ecs = ecs.query(
        "abs(i-j) >= {}".format(min_seqdist)
    ).sort_values(
        by=score, ascending=False
    )

    # select top EC pairs
    top_ecs = sorted_ecs.iloc[0:num_pairs]

    # stack dataframe so it contains each
    # EC twice as forward and backward pairs
    # (i, j) and (j, i)
    flipped = top_ecs.rename(
        columns={"i": "j", "j": "i", "A_i": "A_j", "A_j": "A_i"}
    )

    stacked_ecs = top_ecs.append(flipped)

    # now sum cumulative strength of EC for each position
    ec_sums = pd.DataFrame(
        stacked_ecs.groupby(["i", "A_i"]).sum()
    )

    # average EC strength for top ECs
    avg_degree = top_ecs.loc[:, score].sum() / len(top_ecs)

    # "enrichment" is ratio how much EC strength on
    # an individual position exceeds average strength in top
    ec_sums.loc[:, "enrichment"] = ec_sums.loc[:, score] / avg_degree

    e = ec_sums.reset_index().loc[:, ["i", "A_i", "enrichment"]]
    return e.sort_values(by="enrichment", ascending=False)


class LegacyScoreMixtureModel:
    """
    Assign to each EC score the probability of being in the
    lognormal tail of a normal-lognormal mixture model.

    .. note::
        this is the original version of the score mixture model
        with a normal distribution noise component, this has been
        superseded by a model using a skew normal distribution
    """
    def __init__(self, x, clamp_mu=False, max_fun=10000, max_iter=1000):
        """
        Mixture model of evolutionary coupling scores to
        determine significant scores that are in high-scoring,
        positive tail of distribution.

        Parameters
        ----------
        x : np.array (or list-like)
            EC scores from which to infer the mixture model
        clamp_mu : bool, optional (default: False)
            Fix mean of Gaussian component to 0 instead of
            fitting it based on data
        max_fun : int
            Maximum number of function evaluations
        max_iter : int
            Maximum number of iterations
        """
        x = np.array(x)

        # Infer parameters of mixture model
        self.params = self._learn_params(x, clamp_mu, max_fun, max_iter)

    @classmethod
    def _learn_params(cls, x, clamp_mu, max_fun, max_iter):
        """
        Infer parameters of mixture model.

        Parameters
        ----------
        x : np.array (or list-like)
            EC scores from which to infer the mixture model
        clamp_mu : bool, optional (default: False)
            Fix mean of Gaussian component to 0 instead of
            fitting it based on data
        max_fun : int
            Maximum number of function evaluations
        max_iter : int
            Maximum number of iterations

        Returns
        -------
        mu : float
            Mean of normal distribution
        sigma : float
            Standard deviation of normal distribution
        q : float
            Relative weights of each distribution
        logmu : float
            Mean of lognormal distribution
        logsigma : float
            Standard deviation of lognormal distribution
        """
        # Initial starting parameters for normal
        # and lognormal distributions
        # q: relative contribution of each distribtuion
        mu = 0
        sigma = np.std(x)
        q = 1
        logsigma = 0.4
        logmu = np.percentile(x, 75) - logsigma**2 / 2
        param = np.array([mu, sigma, q, logmu, logsigma])

        # Target function for minimization
        def target_func(params):
            if clamp_mu:
                params[0] = 0

            return -np.sum(np.log(cls._gaussian_lognormal(x, params)))

        # Minimize function
        coeff = op.fmin(
            target_func, param, maxfun=max_fun, maxiter=max_iter, disp=False
        )

        # If clamping mu, also set to 0 in the end, so this value
        # is used for probability calculations
        if clamp_mu:
            coeff[0] = 0

        q = coeff[2]
        # Check if fit worked
        if q >= 1 or np.isinf(q) or np.isneginf(q):
            raise ValueError(
                "No tail, fit failed. q={}".format(q)
            )

        return coeff

    @classmethod
    def _gaussian_lognormal(cls, x, params):
        """
        Gaussian-lognormal mixture probability
        density function.

        Parameters
        ----------
        x : np.array
            Scores for which PDF is calculated
        params : tuple
            Parameters of lognormal-Gaussian mixture
            (mu, sigma, class weight q, loglog, logsigma)

        Returns
        -------
        np.array
            Probabilities
        """
        return cls._gaussian(x, params) + cls._lognormal(x, params)

    @classmethod
    def _gaussian(cls, x, params):
        """
        Normal probability density (multiplied
        by class weight).

        Parameters
        ----------
        x : np.array
            Scores for which PDF is calculated
        params : tuple
            Parameters of lognormal-Gaussian mixture
            (mu, sigma, class weight q, loglog, logsigma)

        Returns
        -------
        np.array
            Probabilities
        """
        mu, sigma, q, logmu, logsigma = params
        return q * stats.norm.pdf(x, loc=mu, scale=sigma)

    @classmethod
    def _lognormal(cls, x, params):
        """
        Log normal probability density (multiplied
        by class weight).

        Parameters
        ----------
        x : np.array
            Scores for which PDF is calculated
        params : tuple
            Parameters of lognormal-Gaussian mixture
            (mu, sigma, class weight q, loglog, logsigma)

        Returns
        -------
        np.array
            Probabilities
        """
        mu, sigma, q, logmu, logsigma = params

        # only assign probability to defined (i.e. positive) values,
        # set all others to zero
        prob = np.zeros(len(x))
        xpos = x > 0
        prob[xpos] = (1 - q) * stats.norm.pdf(
            np.log(x[xpos]), loc=logmu, scale=logsigma
        ) / x[xpos]

        return prob

    def probability(self, x, plot=False):
        """
        Calculate posterior probability of EC pair to
        be located in positive (lognormal) tail of the
        distribution.

        Parameters
        ----------
        x : np.array (or list-like)
            List of scores
        plot : bool, optional (default: False)
            Plot score distribution and probabilities

        Returns
        -------
        posterior : np.array(float)
            Posterior probability of being in signal
            component of mixture model
        """
        x = np.array(x)
        p_lognormal = self._lognormal(x, self.params)
        p_gaussian = self._gaussian(x, self.params)

        posterior = p_lognormal / (p_lognormal + p_gaussian)

        if plot:
            plt.figure(figsize=(12, 8))
            # fig = plt.figure(figsize=(4,4))
            c = "#fdc832"
            n_ECs, edges = np.histogram(x, 1000, density=True)
            mid = []
            for l, r in zip(edges[:-1], edges[1:]):
                mid.append((l + r) / 2)

            plt.plot(
                mid, n_ECs, '-', color=c, markerfacecolor=c,
                markeredgecolor='None', linewidth=1
            )
            plt.plot(x, posterior, '-k', linewidth=2)
            plt.plot(x, p_lognormal, 'r', linewidth=1)
            plt.plot(x, p_gaussian, 'b', linewidth=1)

            takeover = x[p_lognormal > p_gaussian].min()
            plt.axvline(takeover, color="grey")
            plt.axvline(x[posterior > 0.99].min(), color="grey", lw=1)

            pompom = abs(x.min())
            plt.axvline(-pompom, color="grey", ls="--")
            plt.axvline(pompom, color="grey", ls="--")

            plt.xlabel("EC scores")
            plt.ylabel("PDF")

        return posterior


class ScoreMixtureModel:
    """
    Assign to each EC score the probability of being in the
    lognormal tail of a skew normal-lognormal mixture model.
    """
    def __init__(self, x):
        """
        Mixture model of evolutionary coupling scores to
        determine signifcant scores that are in high-scoring,
        positive tail of distribution.

        Parameters
        ----------
        x : np.array (or list-like)
            EC scores from which to infer the mixture model
        """
        x = np.array(x)

        # Infer parameters of mixture model
        self.params = self._learn_params(x)

    @classmethod
    def skewnorm_pdf(cls, x, location, scale, skew):
        """
        Probability density of skew normal distribution
        (noise component)

        Parameters
        ---------
        x : np.array(float)
            Data for which probability density should be calculated
        location : float
            Location parameter of skew normal distribution
        scale : float
            Scale parameter of skew normal distribution
        skew : float
            Skew parameter of skew normal distribution

        Returns
        ------
        density : np.array(float)
            Probability density for input array x
        """
        x_transform = (x - location) / scale
        density = (
            2 / scale *
            stats.norm.pdf(x_transform) *
            stats.norm.cdf(skew * x_transform)
        )

        return density

    @classmethod
    def lognorm_pdf(cls, x, logmu, logsig):
        """
        Probability density of lognormal distribution
        (signal component)

        Parameters
        ---------
        x : np.array(float)
            Data for which probability density should be calculated
        logmu : float
            Location of lognormal distribution (signal component)
        logsig : float
            Scale parameter of lognormal distribution (signal component)

        Returns
        ------
        density : np.array(float)
            Probability density for input array x
        """
        density = np.zeros(len(x))
        xpos = x > 0
        density[xpos] = stats.norm.pdf(
            np.log(x[xpos]), loc=logmu, scale=logsig
        ) / x[xpos]

        return density

    @classmethod
    def skewnorm_constraint(cls, scale, skew):
        """
        Given scale and skew, returns location parameter to yield mean zero

        Parameters
        ----------
        scale : float
            Scale parameter of skew normal distribution
        skew : float
            Skew parameter of skew normal distribution

        Returns
        -------
        location : float
            Location parameter of skew normal distribution s.t.
            mean of distribution is equal to 0
        """
        location = -scale * skew / np.sqrt(1 + skew**2) * np.sqrt(2 / np.pi)
        return location

    @classmethod
    def mixture_pdf(cls, x, p, scale, skew, logmu, logsig):
        """
        Compute mixture probability

        Parameters
        ----------
        x : np.array(float)
            Data for which probability density should be calculated
        p : float
            Mixing fraction between components for noise component
            (signal component will be 1-p)
        scale : float
            Scale parameter of skew normal distribution (noise component)
        skew : float
            Skew parameter of skew normal distribution (noise component)
        logmu : float
            Location of lognormal distribution (signal component)
        logsig : float
            Scale parameter of lognormal distribution (signal component)

        Returns
        -------
        density : np.array(float)
            Probability density for input array x
        """
        location = cls.skewnorm_constraint(scale, skew)
        density = (
            p * cls.skewnorm_pdf(x, location, scale, skew) +
            (1 - p) * cls.lognorm_pdf(x, logmu, logsig)
        )

        return density

    @classmethod
    def posterior_signal(cls, x, p, scale, skew, logmu, logsig):
        """
        Compute posterior probability of being in signal component

        Parameters
        ----------
        x : np.array(float)
            Data for which probability density should be calculated
        p : float
            Mixing fraction between components for noise component
            (signal component will be 1-p)
        scale : float
            Scale parameter of skew normal distribution (noise component)
        skew : float
            Skew parameter of skew normal distribution (noise component)
        logmu : float
            Location of lognormal distribution (signal component)
        logsig : float
            Scale parameter of lognormal distribution (signal component)

        Returns
        -------
        posterior : np.array(float)
            Posterior probability of being in signal component
            for input array x
        """
        P = cls.mixture_pdf(x, p, scale, skew, logmu, logsig)
        posterior = np.zeros(P.shape)
        f2 = cls.lognorm_pdf(x, logmu, logsig)
        posterior[x > 0] = (1 - p) * f2[x > 0] / P[x > 0]

        return posterior

    @classmethod
    def _learn_params(cls, x):
        """
        Infer parameters of mixture model.

        Parameters
        ----------
        x : np.array(float)
            EC scores from which to infer the mixture model

        Returns
        -------
        theta : np.array(float)
            Array with inferred parameters of model
            (mixing fraction, skew-normal scale, skew-normal skew,
            log-normal mean, log-normal stddev)
        """
        # mixing fraction, skew-normal scale, skew-normal skew,
        # log-normal mean, log-normal stddev
        theta = np.array(
            [0.5, np.std(x), 0, np.log(np.max(x)), 0.1]
        )

        def loglk_fun(params):
            return np.sum(np.log(cls.mixture_pdf(x, *params)))

        loglk = loglk_fun(theta)
        delta_loglk = 100
        max_iter = 200
        cur_iter = 0
        tolerance = 0.0001

        while delta_loglk > tolerance and cur_iter < max_iter:
            # E step
            z = 1 - cls.posterior_signal(x, *theta)

            # M step
            # MLE of the mixing fraction is the mean z
            theta[0] = np.mean(z)

            # Log-normal component
            # MLE is the z-weighted mean and std deviation of the log-scores
            pos_ix = x > 0
            z_complement = 1 - z[pos_ix]
            log_score = np.log(x[pos_ix])
            theta[3] = np.sum(z_complement * log_score) / np.sum(z_complement)
            theta[4] = np.sqrt(
                np.sum(z_complement * (log_score - theta[3])**2) / z_complement.sum()
            )

            # Skew-normal distribution
            # MLE requires numerical optimization
            def objfun(params):
                return -np.sum(
                    z * np.log(
                        cls.skewnorm_pdf(
                            x, cls.skewnorm_constraint(params[0], params[1]),
                            params[0], params[1]
                        )
                    )
                )

            theta[1:3] = op.fmin(objfun, theta[1:3], disp=False)

            # Test for EM convergence
            loglk_new = loglk_fun(theta)
            delta_loglk = loglk_new - loglk
            loglk = loglk_new

            # status update
            cur_iter += 1

        return theta

    def probability(self, x, plot=False):
        """
        Calculate posterior probability of EC pair to
        be located in positive (lognormal) tail of the
        distribution.

        Parameters
        ----------
        x : np.array (or list-like)
            List of scores

        Returns
        -------
        posterior : np.array(float)
            Posterior probability of being in signal
            component of mixture model
        """
        posterior = self.posterior_signal(x, *self.params)

        if plot:
            plt.hist(x, normed=True, bins=50, color="k")
            plt.plot(x, self.mixture_pdf(x, *self.params), color="r", lw=3)
            plt.plot(x, posterior, color="gold", lw=3)

        return posterior


class EVComplexScoreModel:
    """
    Assign to each EC score a (unnormalized) EVcomplex score as
    described in Hopf, Schärfe et al. (2014).

    TODO: this implementation currently does not take into account
    score normalization for the number of sequences and length of
    the model
    """
    def __init__(self, x):
        """
        Initialize EVcomplex score model
        
        Parameters
        ----------
        x : np.array (or list-like)
            EC scores from which to infer the mixture model
        """
        self.x = np.array(x)

    def probability(self, x, plot=False):
        """
        Calculates evcomplex score as cn_score / min_cn_score.
        TODO: plotting functionality not yet implemented

        Parameters
        ----------
        x : np.array (or list-like)
            List of scores
        plot: bool, optional (default: False)
            Plot score distribution

        Returns
        -------
        probability: np.array(float)
            EVcomplex score
        """
        # Calculate the minimum score
        min_score = abs(np.min(self.x))

        return x / min_score


def add_mixture_probability(ecs, model="skewnormal", score="cn",
                            clamp_mu=False, plot=False):
    """
    Add lognormal mixture model probability to EC table.

    Parameters
    ----------
    ecs : pd.DataFrame
        EC table with scores
    model : {"skewnormal", "normal"}, optional (default: skewnormal)
        Use model with skew-normal or normal distribution
        for the noise component of mixture model
    score : str, optional (default: "cn")
        Score on which mixture model will be based
    clamp_mu : bool, optional (default: False)
        Fix mean of Gaussian component to 0 instead of
        fitting it based on data
    plot : bool, optional (default: False)
        Plot score distribution and probabilities

    Returns
    -------
    ec_prob : pd.DataFrame
        EC table with additional column "probability"
        that for each EC contains the posterior
        probability of belonging to the lognormal
        tail of the distribution.
    """
    ec_prob = deepcopy(ecs)

    # learn mixture model
    if model == "skewnormal":
        mm = ScoreMixtureModel(ecs.loc[:, score].values)
    elif model == "normal":
        mm = LegacyScoreMixtureModel(ecs.loc[:, score].values, clamp_mu)
    elif model == "evcomplex":
        mm = EVComplexScoreModel(ecs.loc[:, score].values)
    else:
        raise ValueError(
            "Invalid model selection, valid options are: "
            "skewnormal, normal, evcomplex"
        )

    # assign probability
    ec_prob.loc[:, "probability"] = mm.probability(
        ec_prob.loc[:, score].values, plot=plot
    )

    return ec_prob


def logreg_classifier_to_dict(classifier, feature_names=None):
    """
    Serialize sklearn logistic regression classifier

    Inspired by https://stackoverflow.com/questions/48328012/python-scikit-learn-to-json

    Parameters
    ----------
    classifier : sklearn.linear_model.LogisticRegression
        Logistic regression classifier to be serialized
    feature_names : list(str)
        Feature names of coefficients in classifier, must
        be in the same order as classifier.coef_

    Returns
    -------
    model : dict
        Serialized classifier
    """
    params = {
        "classifier_settings": classifier.get_params(),
        "model_settings": {}
    }

    for attrib in ["classes_", "intercept_", "coef_", "n_iter_"]:
        params["model_settings"][attrib] = getattr(classifier, attrib).tolist()

    if feature_names is not None:
        params["feature_names"] = feature_names

    return params


def logreg_classifier_from_dict(params):
    """
    Deserialize model parameters into sklearn LogisticRegression classifier

    Inspired by https://stackoverflow.com/questions/48328012/python-scikit-learn-to-json

    Parameters
    -------
    model : dict
        Serialized classifier

    Parameters
    ----------
    classifier : sklearn.linear_model.LogisticRegression
        Deserialized logistic regression classifier
    feature_names : list(str)
        Feature names of coefficients in classifier
        (in same order as classifier.coef_)
    """
    classifier = LogisticRegression(
        **params["classifier_settings"]
    )

    for attrib, values in params["model_settings"].items():
        setattr(classifier, attrib, np.array(values))

    feature_names = params.get("feature_names")

    return classifier, feature_names


def add_freqs_to_ec_table(ecs, freqs):
    """
    Add residue and gap frequency as well as conservation info
    to data table by merging EC table and frequencies table

    Parameters
    ----------
    ecs : pd.DataFrame
        EC table generated by pipeline (e.g. _ECs.txt, _CouplingsScores.csv)
    freqs : pd.DataFrame
        Frequency table generated by pipeline (_frequencies.csv) using
        evcouplings.align.protocol.describe_frequencies()
    """
    # rename column names to be more descriptive
    freqs = freqs.rename(
        columns={
            "-": "gap_i",
            "conservation": "cons_i",
        }
    )

    # extract frequency of target residue for each position
    freqs.loc[:, "freq_i"] = [r[r["A_i"]] for idx, r in freqs.iterrows()]

    # limit to relevant columns and create second version
    # for mapping second position suffixed with _j
    freqs_sel_i = freqs[["i", "A_i", "freq_i", "gap_i", "cons_i"]]

    freqs_sel_j = freqs_sel_i.rename(columns={
        c: c.replace("i", "j") for c in freqs_sel_i.columns
    })

    # add frequency/conservation information to table
    ecs_with_freqs = ecs.merge(
        freqs_sel_i, on=["i", "A_i"]
    ).merge(
        freqs_sel_j, on=["j", "A_j"]
    )

    # check nothing went wrong and we kept full table
    assert len(ecs_with_freqs) == len(ecs)

    return ecs_with_freqs


def mad_outlier_score(x):
    """
    Median absolute deviation outlier scoring of
    ECs in tail vs. background noise distribution
    (robust z-scoring using median and median
    absolute deviation instead of mean and
    sample standard deviation)

    Parameters
    ----------
    x : np.array
        List of scores (note this must contain *all*
        scores for estimation of median and median
        absolute deviation, i.e. a list of top X
        ECs only will give wrong results)

    Returns
    -------
    np.array
        Outlier scores
    """
    med = np.median(x)
    mad = median_absolute_deviation(x)
    return (x - med) / mad


class LogisticRegressionScorer:
    """
    Rescore ECs based on logistic regression model
    fitted to large set of runs
    """
    def __init__(self, logreg_model_file=None, min_n_eff_over_l=0.375):
        """
        Create new logistic regression-based
        EC rescorer

        Parameters
        ----------
        logreg_model_file : str, optional (default: None)
            Specify path to yml file with logistic regression
            model parameters; if None, will use default
            model included with package
            (evcouplings/couplings/scoring_models/logistic_regression_all.yml)
        min_n_eff_over_l : float, optional (default: 0.3)
            Minimum number of effective sequences per model site required
            for rescoring to be applied; otherwise standard score will
            be returned and all probabilities will be set to 0. The
            default value will be divided by theta for the rescored run,
            the default of 0.375 derives from N_eff/L = 0.3 at theta = 0.8
        """
        # by default load internal classifier included with package
        if logreg_model_file is None:
            logreg_model_file = resource_filename(
                __name__, "scoring_models/logistic_regression_all.yml"
            )

        # load classifier from param file
        logreg_model_serialized = read_config_file(logreg_model_file)

        # deserialize and store classifier
        self.classifier, self.feature_names = logreg_classifier_from_dict(
            logreg_model_serialized
        )

        # store min N_eff/L requirement
        self.min_n_eff_over_l = min_n_eff_over_l

    @classmethod
    def _create_full_data_table(cls, ecs, freqs, theta, effective_sequences, num_sites):
        """
        Create full EC feature table

        Parameters
        ----------
        ecs : pd.DataFrame
            Full, unfiltered EC table
        freqs : pd.DataFrame
            Frequencies table
        theta : float
            Theta parameter used for sequence reweighting
            (used to adjust influence of different thetas in scoring model)
        effective_sequences : float
            Number of effective sequences after clustering and downweighting
        num_sites : float, optional (default: None)
            Number of sites in model / alignment; if None, will
            be automatically inferred from Ec table

        Returns
        -------
        pd.DataFrame
            Full EC table with annotated input features for model
        """
        meff_over_l = effective_sequences / num_sites
        meff_over_l2 = effective_sequences / num_sites ** 2

        meff_over_l_norm = meff_over_l / theta
        meff_over_l2_norm = meff_over_l2 / theta

        # add frequency and conservation info to EC table
        ecs = add_freqs_to_ec_table(
            ecs, freqs
        )

        # assign additional features to EC table
        # needed by prediction model
        ecs = ecs.assign(
            num_sites_log=np.log10(num_sites),
            min_gap=np.minimum(ecs.gap_i, ecs.gap_j),
            max_gap=np.maximum(ecs.gap_i, ecs.gap_j),
            min_cons=np.minimum(ecs.cons_i, ecs.cons_j),
            max_cons=np.maximum(ecs.cons_i, ecs.cons_j),
            meff_over_l_norm_log=np.log10(meff_over_l_norm),
            meff_over_l2_norm_log=np.log10(meff_over_l2_norm)
        )

        return ecs

    def score(self, ecs, freqs, theta, effective_sequences, num_sites=None, score="cn"):
        """
        Rescore EC table (must be full, unfiltered table) using logistic regression
        based probabilistic scoring. If number of effective sequences is too low for reliable
        scoring, score will default to input score and probabilities will all be set to 0.

        Parameters
        ----------
        ecs : pd.DataFrame
            Full, unfiltered EC table
        freqs : pd.DataFrame
            Frequencies table
        theta : float
            Theta parameter used for sequence reweighting
            (used to adjust influence of different thetas in scoring model)
        effective_sequences : float
            Number of effective sequences after clustering and downweighting
        num_sites : float, optional (default: None)
            Number of sites in model / alignment; if None, will
            be automatically inferred from Ec table
        score : string, optional (default: "cn")
            Input score to be input into logistic regression model.
            Model has been trained on CN scores only, so is not applicable
            to other types of scoring like DI score

        Returns
        -------
        ecs_final : pd.DataFrame
            Rescored ECs, same table as ecs but with three
            additional columns (mad_score, probability, score) containing
            1) median absolute deviation outlier score, 2) logistic regression
            model probability for EC to be true EC, and 3) pair score
            derived from decision function of logistic regression model
            (this can reorder the original ECs based on positional conservation
            and gap frequencies).
            Final table is sorted by "score" in descending order.
        """
        # infer num_sites from EC table if not defined
        if num_sites is None:
            num_sites = len(
                set(ecs.i.unique()) | set(ecs.j.unique())
            )

        # if N_eff/L is too low for model to give reliable results,
        # keep standard score and set all probabilities to 0
        if effective_sequences / num_sites / theta < self.min_n_eff_over_l:
            return ecs.assign(
                score=ecs[score],
                probability=0
            )

        # add medium absolute deviation outlier score to copy of EC table
        ecs = ecs.assign(
            mad_score=mad_outlier_score(ecs[score])
        )

        # add additional information to table
        ecs_full = self._create_full_data_table(
            ecs, freqs, theta, effective_sequences, num_sites
        )

        # extract features for classifier ("X")
        feature_table = ecs_full.reindex(self.feature_names, axis=1)

        # predict probabilities of contact and decision function
        # (target class/true EC := 1)
        probs = self.classifier.predict_proba(feature_table)[:, 1]
        decision_func = self.classifier.decision_function(feature_table)

        # assign to EC table
        ecs_final = ecs_full.assign(
            score=decision_func,
            probability=probs
        ).sort_values(
            by="score", ascending=False
        )

        # remove temporary feature columns and return updated EC table
        return ecs_final[
            list(ecs.columns) + ["probability", "score"]
        ]
Functions for handling evolutionary couplings data.  .. todo::      1. clean up     2. add Pompom score     3. add mapping tools (multidomain, complexes)     4. ECs to matrix     5. APC on subsets of positions (e.g. for complexes)  Authors:   <NAME>   <NAME> (original mixture model code)   <NAME> (skew normal mixture model)   <NAME> (EVComplex Score code) Read a raw EC file (e.g. from plmc) and sort     by scores      Parameters     ----------     filename : str         File containing evolutionary couplings     sort : bool, optional (default: True)         If True, sort pairs by coupling score in         descending order     score_column : str, optional (default: True)         Score column to be used for sorting      Returns     -------     ecs : pd.DataFrame         Table of evolutionary couplings Calculate EC "enrichment" as first described in     Hopf et al., Cell, 2012.      .. todo::          Make this handle segments if they are in EC table      Parameters     ----------     ecs : pd.DataFrame         Dataframe containing couplings     num_pairs : int or float, optional (default: 1.0)         Number of ECs to use for enrichment calculation.         - If float, will be interpreted as fraction of the         length of the sequence (e.g. 1.0*L)         - If int, will be interpreted as         absolute number of pairs     score : str, optional (default: cn)         Pair coupling score used for calculation     min_seqdist : int, optional (default: 6)         Minimum sequence distance of couplings that will         be included in the calculation      Returns     -------     enrichment_table : pd.DataFrame         Sorted table with enrichment values for each         position in the sequence # determine how many positions ECs are over # calculate absolute number of pairs if # fraction of length is given # get longrange ECs and sort by score # select top EC pairs # stack dataframe so it contains each # EC twice as forward and backward pairs # (i, j) and (j, i) # now sum cumulative strength of EC for each position # average EC strength for top ECs # "enrichment" is ratio how much EC strength on # an individual position exceeds average strength in top Assign to each EC score the probability of being in the     lognormal tail of a normal-lognormal mixture model.      .. note::         this is the original version of the score mixture model         with a normal distribution noise component, this has been         superseded by a model using a skew normal distribution Mixture model of evolutionary coupling scores to         determine significant scores that are in high-scoring,         positive tail of distribution.          Parameters         ----------         x : np.array (or list-like)             EC scores from which to infer the mixture model         clamp_mu : bool, optional (default: False)             Fix mean of Gaussian component to 0 instead of             fitting it based on data         max_fun : int             Maximum number of function evaluations         max_iter : int             Maximum number of iterations # Infer parameters of mixture model Infer parameters of mixture model.          Parameters         ----------         x : np.array (or list-like)             EC scores from which to infer the mixture model         clamp_mu : bool, optional (default: False)             Fix mean of Gaussian component to 0 instead of             fitting it based on data         max_fun : int             Maximum number of function evaluations         max_iter : int             Maximum number of iterations          Returns         -------         mu : float             Mean of normal distribution         sigma : float             Standard deviation of normal distribution         q : float             Relative weights of each distribution         logmu : float             Mean of lognormal distribution         logsigma : float             Standard deviation of lognormal distribution # Initial starting parameters for normal # and lognormal distributions # q: relative contribution of each distribtuion # Target function for minimization # Minimize function # If clamping mu, also set to 0 in the end, so this value # is used for probability calculations # Check if fit worked Gaussian-lognormal mixture probability         density function.          Parameters         ----------         x : np.array             Scores for which PDF is calculated         params : tuple             Parameters of lognormal-Gaussian mixture             (mu, sigma, class weight q, loglog, logsigma)          Returns         -------         np.array             Probabilities Normal probability density (multiplied         by class weight).          Parameters         ----------         x : np.array             Scores for which PDF is calculated         params : tuple             Parameters of lognormal-Gaussian mixture             (mu, sigma, class weight q, loglog, logsigma)          Returns         -------         np.array             Probabilities Log normal probability density (multiplied         by class weight).          Parameters         ----------         x : np.array             Scores for which PDF is calculated         params : tuple             Parameters of lognormal-Gaussian mixture             (mu, sigma, class weight q, loglog, logsigma)          Returns         -------         np.array             Probabilities # only assign probability to defined (i.e. positive) values, # set all others to zero Calculate posterior probability of EC pair to         be located in positive (lognormal) tail of the         distribution.          Parameters         ----------         x : np.array (or list-like)             List of scores         plot : bool, optional (default: False)             Plot score distribution and probabilities          Returns         -------         posterior : np.array(float)             Posterior probability of being in signal             component of mixture model # fig = plt.figure(figsize=(4,4)) Assign to each EC score the probability of being in the     lognormal tail of a skew normal-lognormal mixture model. Mixture model of evolutionary coupling scores to         determine signifcant scores that are in high-scoring,         positive tail of distribution.          Parameters         ----------         x : np.array (or list-like)             EC scores from which to infer the mixture model # Infer parameters of mixture model Probability density of skew normal distribution         (noise component)          Parameters         ---------         x : np.array(float)             Data for which probability density should be calculated         location : float             Location parameter of skew normal distribution         scale : float             Scale parameter of skew normal distribution         skew : float             Skew parameter of skew normal distribution          Returns         ------         density : np.array(float)             Probability density for input array x Probability density of lognormal distribution         (signal component)          Parameters         ---------         x : np.array(float)             Data for which probability density should be calculated         logmu : float             Location of lognormal distribution (signal component)         logsig : float             Scale parameter of lognormal distribution (signal component)          Returns         ------         density : np.array(float)             Probability density for input array x Given scale and skew, returns location parameter to yield mean zero          Parameters         ----------         scale : float             Scale parameter of skew normal distribution         skew : float             Skew parameter of skew normal distribution          Returns         -------         location : float             Location parameter of skew normal distribution s.t.             mean of distribution is equal to 0 Compute mixture probability          Parameters         ----------         x : np.array(float)             Data for which probability density should be calculated         p : float             Mixing fraction between components for noise component             (signal component will be 1-p)         scale : float             Scale parameter of skew normal distribution (noise component)         skew : float             Skew parameter of skew normal distribution (noise component)         logmu : float             Location of lognormal distribution (signal component)         logsig : float             Scale parameter of lognormal distribution (signal component)          Returns         -------         density : np.array(float)             Probability density for input array x Compute posterior probability of being in signal component          Parameters         ----------         x : np.array(float)             Data for which probability density should be calculated         p : float             Mixing fraction between components for noise component             (signal component will be 1-p)         scale : float             Scale parameter of skew normal distribution (noise component)         skew : float             Skew parameter of skew normal distribution (noise component)         logmu : float             Location of lognormal distribution (signal component)         logsig : float             Scale parameter of lognormal distribution (signal component)          Returns         -------         posterior : np.array(float)             Posterior probability of being in signal component             for input array x Infer parameters of mixture model.          Parameters         ----------         x : np.array(float)             EC scores from which to infer the mixture model          Returns         -------         theta : np.array(float)             Array with inferred parameters of model             (mixing fraction, skew-normal scale, skew-normal skew,             log-normal mean, log-normal stddev) # mixing fraction, skew-normal scale, skew-normal skew, # log-normal mean, log-normal stddev # E step # M step # MLE of the mixing fraction is the mean z # Log-normal component # MLE is the z-weighted mean and std deviation of the log-scores # Skew-normal distribution # MLE requires numerical optimization # Test for EM convergence # status update Calculate posterior probability of EC pair to         be located in positive (lognormal) tail of the         distribution.          Parameters         ----------         x : np.array (or list-like)             List of scores          Returns         -------         posterior : np.array(float)             Posterior probability of being in signal             component of mixture model Assign to each EC score a (unnormalized) EVcomplex score as     described in Hopf, Schärfe et al. (2014).      TODO: this implementation currently does not take into account     score normalization for the number of sequences and length of     the model Initialize EVcomplex score model                  Parameters         ----------         x : np.array (or list-like)             EC scores from which to infer the mixture model Calculates evcomplex score as cn_score / min_cn_score.         TODO: plotting functionality not yet implemented          Parameters         ----------         x : np.array (or list-like)             List of scores         plot: bool, optional (default: False)             Plot score distribution          Returns         -------         probability: np.array(float)             EVcomplex score # Calculate the minimum score Add lognormal mixture model probability to EC table.      Parameters     ----------     ecs : pd.DataFrame         EC table with scores     model : {"skewnormal", "normal"}, optional (default: skewnormal)         Use model with skew-normal or normal distribution         for the noise component of mixture model     score : str, optional (default: "cn")         Score on which mixture model will be based     clamp_mu : bool, optional (default: False)         Fix mean of Gaussian component to 0 instead of         fitting it based on data     plot : bool, optional (default: False)         Plot score distribution and probabilities      Returns     -------     ec_prob : pd.DataFrame         EC table with additional column "probability"         that for each EC contains the posterior         probability of belonging to the lognormal         tail of the distribution. # learn mixture model # assign probability Serialize sklearn logistic regression classifier      Inspired by https://stackoverflow.com/questions/48328012/python-scikit-learn-to-json      Parameters     ----------     classifier : sklearn.linear_model.LogisticRegression         Logistic regression classifier to be serialized     feature_names : list(str)         Feature names of coefficients in classifier, must         be in the same order as classifier.coef_      Returns     -------     model : dict         Serialized classifier Deserialize model parameters into sklearn LogisticRegression classifier      Inspired by https://stackoverflow.com/questions/48328012/python-scikit-learn-to-json      Parameters     -------     model : dict         Serialized classifier      Parameters     ----------     classifier : sklearn.linear_model.LogisticRegression         Deserialized logistic regression classifier     feature_names : list(str)         Feature names of coefficients in classifier         (in same order as classifier.coef_) Add residue and gap frequency as well as conservation info     to data table by merging EC table and frequencies table      Parameters     ----------     ecs : pd.DataFrame         EC table generated by pipeline (e.g. _ECs.txt, _CouplingsScores.csv)     freqs : pd.DataFrame         Frequency table generated by pipeline (_frequencies.csv) using         evcouplings.align.protocol.describe_frequencies() # rename column names to be more descriptive # extract frequency of target residue for each position # limit to relevant columns and create second version # for mapping second position suffixed with _j # add frequency/conservation information to table # check nothing went wrong and we kept full table Median absolute deviation outlier scoring of     ECs in tail vs. background noise distribution     (robust z-scoring using median and median     absolute deviation instead of mean and     sample standard deviation)      Parameters     ----------     x : np.array         List of scores (note this must contain *all*         scores for estimation of median and median         absolute deviation, i.e. a list of top X         ECs only will give wrong results)      Returns     -------     np.array         Outlier scores Rescore ECs based on logistic regression model     fitted to large set of runs Create new logistic regression-based         EC rescorer          Parameters         ----------         logreg_model_file : str, optional (default: None)             Specify path to yml file with logistic regression             model parameters; if None, will use default             model included with package             (evcouplings/couplings/scoring_models/logistic_regression_all.yml)         min_n_eff_over_l : float, optional (default: 0.3)             Minimum number of effective sequences per model site required             for rescoring to be applied; otherwise standard score will             be returned and all probabilities will be set to 0. The             default value will be divided by theta for the rescored run,             the default of 0.375 derives from N_eff/L = 0.3 at theta = 0.8 # by default load internal classifier included with package # load classifier from param file # deserialize and store classifier # store min N_eff/L requirement Create full EC feature table          Parameters         ----------         ecs : pd.DataFrame             Full, unfiltered EC table         freqs : pd.DataFrame             Frequencies table         theta : float             Theta parameter used for sequence reweighting             (used to adjust influence of different thetas in scoring model)         effective_sequences : float             Number of effective sequences after clustering and downweighting         num_sites : float, optional (default: None)             Number of sites in model / alignment; if None, will             be automatically inferred from Ec table          Returns         -------         pd.DataFrame             Full EC table with annotated input features for model # add frequency and conservation info to EC table # assign additional features to EC table # needed by prediction model Rescore EC table (must be full, unfiltered table) using logistic regression         based probabilistic scoring. If number of effective sequences is too low for reliable         scoring, score will default to input score and probabilities will all be set to 0.          Parameters         ----------         ecs : pd.DataFrame             Full, unfiltered EC table         freqs : pd.DataFrame             Frequencies table         theta : float             Theta parameter used for sequence reweighting             (used to adjust influence of different thetas in scoring model)         effective_sequences : float             Number of effective sequences after clustering and downweighting         num_sites : float, optional (default: None)             Number of sites in model / alignment; if None, will             be automatically inferred from Ec table         score : string, optional (default: "cn")             Input score to be input into logistic regression model.             Model has been trained on CN scores only, so is not applicable             to other types of scoring like DI score          Returns         -------         ecs_final : pd.DataFrame             Rescored ECs, same table as ecs but with three             additional columns (mad_score, probability, score) containing             1) median absolute deviation outlier score, 2) logistic regression             model probability for EC to be true EC, and 3) pair score             derived from decision function of logistic regression model             (this can reorder the original ECs based on positional conservation             and gap frequencies).             Final table is sorted by "score" in descending order. # infer num_sites from EC table if not defined # if N_eff/L is too low for model to give reliable results, # keep standard score and set all probabilities to 0 # add medium absolute deviation outlier score to copy of EC table # add additional information to table # extract features for classifier ("X") # predict probabilities of contact and decision function # (target class/true EC := 1) # assign to EC table # remove temporary feature columns and return updated EC table