"""
Created on Fri Aug 31 2012 14:05

Copyright (c) 2012, Martin S. Lindner and Maximilian Kollock, LindnerM@rki.de,
Robert Koch-Institut, Germany,
All rights reserved.

Redistribution and use in source and binary forms, with or without
modification, are permitted provided that the following conditions are met:
    * Redistributions of source code must retain the above copyright
      notice, this list of conditions and the following disclaimer.
    * Redistributions in binary form must reproduce the above copyright
      notice, this list of conditions and the following disclaimer in the
      documentation and/or other materials provided with the distribution.
    * The name of the author may not be used to endorse or promote products
      derived from this software without specific prior written permission.

THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" AND
ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED
WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
DISCLAIMED. IN NO EVENT SHALL MARTIN S. LINDNER OR MAXIMILIAN KOLLOCK BE LIABLE
FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED
AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT
(INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS
SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
"""



import numpy as np
import scipy.stats as stats

from scipy.special import digamma, betainc
from scipy.optimize import newton


"""----------------------------------------------------------------------------
    Define the distributions
----------------------------------------------------------------------------"""

class Distribution:
    _p1 = None
    _p2 = None
    _name = "General Distribution"
    _dof = 0 # Number of degrees of freedom

    alpha = 1.
    def __str__(self):
        """ return the name of the distribution """
        return self._name

    def set_par(self,p1=None, p2=None):
        """ explicitly set a parameter """
        if p1 != None:
            self._p1 = p1
        if p2 != None:
            self._p2 = p2

    def pmf(self, x):
        """ return the value of the probability mass function at x """
        return x*0.

    def estimate_par(self, data=None, weights=None):
        """ estimate the distribution parameters from the data and the weights
        (if provided) """
        pass

    def init_par(self, mean=None, var=None):
        """ estimate initial distribution parameters given the mean and
        variance of the data"""
        pass

    def report_stats(self, width=20):
        """ return a string that reports information about the distribution """
        return str(self._name).ljust(width) + str(self.alpha).ljust(width) + \
           str(self._p1).ljust(width) + str(self._p2).ljust(width)


class Zero(Distribution):
    _name = "Zero"
    _dof = 1

    def pmf(self,x):
        if isinstance(x,np.ndarray):
            return (x==0).astype(np.float)
        else:
            return float(x==0)


class NBinom(Distribution):
    _name = "NBinom"
    _dof = 3
    _use_MOM = False

    def pmf(self, x):
        return stats.nbinom.pmf(x,self._p1,self._p2)

    def estimate_par(self, data, weights=None):
        if weights == None:
            weights = data*0. + 1.
        norm = np.sum(weights)
        mean = np.sum(data*weights)/(norm + 10**(-25))
        var = np.sum((data - mean)**2 * weights) / (norm + 10**(-25))

        if self._use_MOM:
            if var < mean:
                var = 1.01*mean
            self._p1 = mean**2 / (var - mean)
            self._p2 = mean / var
        else:
            def dp1_llnbinom(param,obs,obs_w):
                # param: parameter 1
                # obs:   observed values
                # obs_w: weight of each value
                N = np.sum(obs_w)
                data_mean = np.sum(obs*obs_w)/(N)
                return np.sum(digamma(obs+param)*obs_w) - N*digamma(param) + \
                   N*np.log(data_mean/(param+data_mean))
            try:
                self._p1 = newton(dp1_llnbinom,self._p1, args=(data,weights),
                                                  maxiter=10000)
                self._p2 = (self._p1)/(self._p1+mean)
            except:
                print('Warning: MLE for negative binomial failed. Using MOM.')
                if var < mean:
                    print('Warning: var < mean')
                    var = 1.01*mean
                self._p1 = mean**2 / (var - mean)
                self._p2 = mean / var



    def report_stats(self, width=20):
        """ return a string that reports information about the distribution """
        return str(self._name).ljust(width) + str(self.alpha).ljust(width) + \
           str(self._p1).ljust(width) + str(self._p2).ljust(width) + \
                   str(stats.nbinom.mean(self._p1,self._p2)).ljust(width)



class Poisson(Distribution):
    _name = "Poisson"
    _dof = 2

    def pmf(self,x):
        return stats.poisson.pmf(x,self._p1)

    def estimate_par(self, data, weights=None):
        mean = np.sum(data*weights) / (np.sum(weights) )
        self._p1 = mean



class TailDistribution(Distribution):
    _name = "Tail"
    _dof = 1
    _norm = False  # normalization; recalculate only if necessary
    _parent = None
    _active = True # switch tail on/off

    def set_par(self,p1=None, p2=None):
        """ explicitly set a parameter """
        if p1 != None:
            self._p1 = p1
        if p2 != None:
            self._p2 = p2
        self._norm = False

    def estimate_par(self, data=None, weights=None):
        """ Do not estimate parameters, but obtain parameters from parent
        distribution """
        self._p1 = self._parent._p1
        self._p2 = self._parent._p2
        self._norm = False



class NbTail(TailDistribution):
    _name = "Tail"
    _dof = 1
    def __init__(self, nbinom):
        """ NbTail distribution must be connected to a negative binomial"""
        if isinstance(nbinom, NBinom):
            self._parent = nbinom
        else:
            raise(Exception("NbTail must be connected to a NBinom object"))

    def pmf(self, x):
        if np.isscalar(x) and x == 0:
            return 0.
        if stats.nbinom.mean(self._p1, self._p2) < 2.:
            self._active = False # switch tail permanently off
        if self._active == False:
            return 0*x


        def betaincreg(x,p1,p2):
            return 1-betainc(p1,x+1,p2)

        # calculate normalization up to certain precision
        if not self._norm:
            ks = int(max(2,stats.nbinom.ppf(0.999999,self._p1,self._p2)))
            norm = np.sum(betaincreg(np.arange(1,ks),self._p1,self._p2))
        else:
            norm = self._norm
        # now return the value(s)
        ret = betaincreg(x,self._p1,self._p2) / norm
        if not np.isscalar(x):
            ret[np.where(x==0)] = 0.
        return ret



class PoissonTail(TailDistribution):
    _name = "Tail of Poisson"
    _dof = 1
    def __init__(self, poisson):
        """ PoissonTail distribution must be connected to a poisson
        distribution """
        if isinstance(poisson, Poisson):
            self._parent = poisson
        else:
            raise(Exception("PoissonTail must be connected to a Poisson object"))

    def pmf(self, x):
        if self._p1 < 2.:
            self._active = False # switch tail permanently off
        if self._active == False:
            return 0*x

        if np.isscalar(x) and x == 0:
            return 0.

        xmax = int(max(np.max(x),5,stats.poisson.ppf(0.999999,self._p1)))+1
        xs = np.arange(0,xmax, dtype=np.float)
        #backward cumulative sum
        with np.errstate(divide='ignore', invalid='ignore'):
            tx = np.cumsum((stats.poisson.pmf(xs, self._p1)/xs)[::-1])[::-1]
        tx[0] = 0
        tx /= np.sum(tx)
        return tx[x]



class Geom(Distribution):
    """ Geometric distribution for fitting distances between reads. """
    _name = "Geometric"
    _dof = 2

    def pmf(self,x):
        return np.power(1-self._p1,x)*self._p1

    def estimate_par(self, data, weights=None):
        mean = np.sum(data*weights) / (np.sum(weights) )
        self._p1 = 1./(1+mean)

    def report_stats(self, width=20):
        """ return a string that reports information about the distribution """
        return str(self._name).ljust(width) + str(self.alpha).ljust(width) + \
            str(self._p1).ljust(width) + '[None]'.ljust(width) + \
                  str(1./self._p1).ljust(width)


def build_mixture_model(dist_str):
    """ Build a mixture model from a string """
    num_dist = len(dist_str)
    mm = np.array([Zero()]*num_dist)
    it = 0
    for dist in dist_str:
        if dist == "z":
            it += 1
        elif dist=="n":
            mm[it] = NBinom()
            mm[it]._use_MOM = True
            it += 1
        elif dist=="N":
            mm[it] = NBinom()
            mm[it]._use_MOM = False
            it += 1
        elif dist == "t":
            if it > 0 and isinstance(mm[it-1], NBinom):
                mm[it] = NbTail(mm[it-1])
                it += 1
            elif it > 0 and isinstance(mm[it-1], Poisson):
                mm[it] = PoissonTail(mm[it-1])
                it += 1
            else:
                raise Exception("Error: wrong input '%s'. A Tail distribution "
                                                "(t) must follow a Negative Binomial (n|N) or "
                                                "Poisson (p)."%dist_str)
        elif dist == "p":
            mm[it] = Poisson()
            it+=1
        elif dist == "g":
            mm[it] = Geom()
            it+=1
        else:
            raise Exception("Input Error: distribution %s not recognized!"%dist)

    # initialize all distributions with equal weights
    alpha = 1./len(mm)
    for dist in mm:
        dist.alpha = alpha

    return mm


def init_gamma(mixture_model, dataset):
    """ create initial responsibilities gamma. The probability that a coverage
    belongs to a distribution is equal for all distributions."""
    N_mm = len(mixture_model)
    N_cov = len(dataset.value)

    return np.array([[1./N_mm for d in mixture_model] for i in range(N_cov)])



"""----------------------------------------------------------------------------
    Define DataSet to store all relevant information (incl. GCP)
----------------------------------------------------------------------------"""
class DS_cov:
    def __init__(self,ref):
        """ fill dataset with data provided in Reference ref """
        t_lns = [t.length for t in iter(ref.targets.values())]
        # r_info contains read position and length for each read
        r_info = [[[r[0],r[1]] for r in iter(t.reads.values())] for t in iter(ref.targets.values())]

        # pile up coverage for each target and summarize in one array
        cov = [np.zeros((l,),dtype=np.int32) for l in t_lns]
        rlen = 0
        rds = 0
        for c,ri in zip(cov,r_info):
            for p,l in ri:
                c[p:(p+l)] += 1
                rds += 1
                rlen += l
        one_cov = np.concatenate(cov)
                
        self.value = np.unique(one_cov)
        self.count = np.array( [np.sum(one_cov==v) for v in self.value] )
        self.total = rds                # Number of reads
        self.rlen = rlen/float(rds)     # average read length
        self.glen = sum(t_lns)  # length of target genome

class DS_dst:
    def __init__(self,ref):
        """ fill dataset with data provided in Reference ref """
        t_lns = [t.length for t in iter(ref.targets.values())]
        r_pos = [[r[0] for r in iter(t.reads.values())] for t in iter(ref.targets.values())]
        r_len = [[r[1] for r in iter(t.reads.values())] for t in iter(ref.targets.values())]
        r_len = np.concatenate([np.array(ls) for ls in r_len])
        positions = [np.sort(np.array(ps)) for ps in r_pos]
        distances = np.concatenate([p[1:] - p[:-1] for p in positions])

        self.value = np.unique(distances)
        self.count = np.array( [np.sum(distances==v) for v in self.value] )
        self.total = sum([len(p) for p in positions])
        self.rlen = np.mean(r_len)
        self.glen = sum(t_lns)





"""----------------------------------------------------------------------------
    Iterative Method function definitions
----------------------------------------------------------------------------"""


def update_gamma(data_set, mixture_model, gamma):
    """ Update the probability gamma_i(x), that a position with coverage x
    belongs to distribution i
    """

    num_dist = len(mixture_model)
    for it in range(len(data_set.value)):
        # calculate probability that coverages[it] belongs to a distribution
        prob = [dist.alpha*dist.pmf(data_set.value[it]) for dist in mixture_model]
        if sum(prob) <= 0:
            gamma[it,:] = 0.
        else:
            gamma[it,:] = np.array([prob[i]/sum(prob) for i in range(num_dist)])

    return 1


def update_alpha(data_set, mixture_model, gamma):
    """ Update alpha_i, the proportion of data that belongs to
    distribution i """

    for i in range(len(mixture_model)-1):
        w_probs = np.array([p*w for p,w in zip(gamma[:,i],data_set.count)])
        mixture_model[i].alpha = np.sum(w_probs) / np.sum(data_set.count)

    mixture_model[-1].alpha = 1 - sum([d.alpha for d in mixture_model[:-1]])

    return 1


def iterative_fitting(data_set, mixture_model, gamma, iterations):
    """ Generator fuction to run the iterative method. Operates directly on the
    data structures mixture_model and gamma. """

    for i in range(iterations):
        # Expectation-step: update gammas
        update_gamma(data_set, mixture_model, gamma)

        # temporarily store old parameters
        old_p1 = np.array([d._p1 or 0 for d in mixture_model])
        old_p2 = np.array([d._p2 or 0 for d in mixture_model])
        old_alpha = np.array([d.alpha for d in mixture_model])

        # Parameter estimation step
        for j,dist in enumerate(mixture_model):
            dist_weights = gamma[:,j]*data_set.count
            dist.estimate_par(data_set.value, dist_weights)
        update_alpha(data_set, mixture_model, gamma)

        # calculate relative change of the parameters
        new_p1 = np.array([d._p1 or 0 for d in mixture_model])
        new_p2 = np.array([d._p2 or 0 for d in mixture_model])
        new_alpha = np.array([d.alpha for d in mixture_model])

        rel_p1 = np.max(np.abs(new_p1-old_p1) / (new_p1 + 1e-20))
        rel_p2 = np.max(np.abs(new_p2-old_p2) / (new_p2 + 1e-20))
        rel_alpha = np.max(np.abs(new_alpha-old_alpha) / (new_alpha + 1e-20))

        max_change = np.max([rel_p1, rel_p2, rel_alpha])

        # maximum CDF difference
        xs = np.arange(np.max(data_set.value)+1,dtype=np.int)
        ref_pdf = np.zeros((np.max(data_set.value)+1,))
        for dist in mixture_model:
            ref_pdf += dist.pmf(xs)*dist.alpha

        obs_pdf = np.zeros((np.max(data_set.value)+1,))
        obs_pdf[data_set.value.astype(np.int)] = data_set.count/float(np.sum(data_set.count))
        max_cdf_diff = np.max(np.abs(np.cumsum(ref_pdf)-np.cumsum(obs_pdf)))

        yield i, max_change, max_cdf_diff


"""----------------------------------------------------------------------------
    Convenience functions to calculate the validity
----------------------------------------------------------------------------"""

def validity_from_coverage(DS, poisson=True, tail=True, max_iter=50, iter_threshold=0.01,
                                          init_alpha=None, data_cutoff=0.95):

    # weight cutoff. Coverages above this value have no weight in parameter estimation
    try:
        cutoff = max(int(np.max(DS.value[np.where(np.cumsum(DS.count)<= \
                         data_cutoff*DS.glen)])+1),10)
    except:
        cutoff = max(np.max(DS.value)+1,10)

    DS.count = DS.count[np.where(DS.value < cutoff)]
    DS.value = DS.value[np.where(DS.value < cutoff)]
    DS.glen = np.sum(DS.count)

    distributions = 'z'
    if poisson:
        distributions += 'p'
    else:
        distributions += 'n'
    if tail:
        distributions += 't'

    MM = build_mixture_model(distributions)

    # set initial proportions for each distribution
    if init_alpha:
        alpha = init_alpha
        if len(init_alpha) < len(MM):
            rest_alpha = 1. - sum(init_alpha)
            rest_models = len(MM) - len(init_alpha)
            alpha.append([rest_alpha / rest_models] * rest_models)

        for dist,a in zip(MM,alpha):
            dist.alpha = float(a)

    # initialize distribution parameters
    data_mean = np.mean(DS.value[1:]*DS.count[1:]/np.mean(DS.count))
    data_var = np.sum((DS.value-data_mean)**2*DS.count)/DS.glen

    for dist in MM:
        if isinstance(dist,Zero) or isinstance(dist,NbTail) \
        or isinstance(dist,PoissonTail):
            dist.estimate_par()
            continue
        if isinstance(dist,NBinom):
            m,v = data_mean,data_var
            if m > v:
                v = 1.01*m
            dist.set_par(m**2 / (v - m) , m / v )
            continue
        if isinstance(dist,Poisson):
            dist.set_par(data_mean)

    # set initial values for the responsibilities gamma
    GAMMA = init_gamma(MM,DS)

    # run the iteration, repeatedly update the variables MM and GAMMA
    for i,change,diff in iterative_fitting(DS, MM, GAMMA, max_iter):
        if change < iter_threshold:
            break

    # estimate the validity as one minus the alpha of the zero distribution
    val = 1. - MM[0].alpha

    if isinstance(MM[1],NBinom):
        cov = stats.nbinom.mean(MM[1]._p1,MM[1]._p2)
    elif isinstance(MM[1],Poisson):
        cov = MM[1]._p1
    else:
        cov = -1

    return val,cov




def validity_from_spaces(DS, num_dist = 3, max_iter = 50, iter_threshold = 0.01,
                init_mean = [2,1000,100000], init_alpha = [0.5,0.49,0.01]):

    MM = build_mixture_model("g"*num_dist)

    # set initial proportions for each distribution
    if init_alpha:
        alpha = init_alpha
        if len(init_alpha) < len(MM):
            # create the missing alpha values
            rest_alpha = 1. - sum(init_alpha)
            rest_models = len(MM) - len(init_alpha)
            alpha.append([rest_alpha / rest_models] * rest_models)

        for dist,a in zip(MM,alpha):
            dist.alpha = float(a)

    # initialize distribution parameters
    factors = np.arange(num_dist)
    data_mean = np.power(10.,factors)

    if init_mean:
        # use mean values provided by the user, where possible
        for i,m in enumerate(init_mean):
            if i < num_dist:
                data_mean[i] = m

    for i,dist in enumerate(MM):
        dist.set_par(1./(1.+data_mean[i]))

    # set initial values for the responsibilities gamma
    GAMMA = init_gamma(MM,DS)

    # run the iteration, repeatedly update the variables MM and GAMMA
    for i,change,diff in iterative_fitting(DS, MM, GAMMA, max_iter):
        if change < iter_threshold:
            break

    # calculate the validity
    if num_dist == 3:
        corr = 1.-MM[0].alpha
    else:
        corr = 1.
    val = corr*DS.total/(MM[num_dist-2]._p1*DS.glen)
    cov = MM[num_dist-2]._p1*DS.rlen

    return val,cov
