"""
Adaption of similarity correction method from GASiC (plus normalized abundances)

@author: Anke Penzlin, June 2013
"""
import numpy as np
import scipy.optimize as opt

def similarity_correction(sim, reads, N, negative_correction=True):
    """ Calculate corrected abundances given a similarity matrix and observations using optimization.
    Copied from GASiC, for metaproteomic experiments replace the term 'reads' with 'spectra'.

    Input:
    sim:   [numpy.array (M,M)] with pairwise similarities between species
    reads: [numpy.array (M,)] with number of observed reads for each species
    N:     [int] total number of reads
    negative_correction: [boolean] whether numerical errors (negative values) should be corrected or not

    Output:
    abundances: [numpy.array (M,)] estimated abundance of each species in the sample
    """

    # transform reads to abundances and rename similarity matrix
    A = np.transpose(sim)
    r = reads.astype(np.float) / N

    rng = range(len(reads))

    # Now solve the optimization problem: min_c |Ac-r|^2 s.t. c_i >= 0 and 1-sum(c_i) >= 0
    # construct objective function
    def objective (c):
        # numpy implementation of objective function |A*c-r|^2
        return np.sum(np.square(np.dot(A,c)-r))

    # construct constraints
    def build_con(k):
        # constraint: k'th component of x should be >0
        return lambda c: c[k]
    cons = [build_con(k) for k in rng]
    # constraint: the sum of all components of x should not exceed 1
    cons.append( lambda c: 1-np.sum(c) )

    # initial guess
    c_0 = np.array([0.5 for i in range(len(reads))])

    # finally: optimization procedure
    abundances = opt.fmin_cobyla(objective, c_0, cons, disp=0, rhoend=1e-10, maxfun=10000)

    if negative_correction:
        for i in range(len(abundances)):
            if abundances[i] < 0 :
                print "warning: correcting negative similarity ", abundances[i], " -> 0"
                abundances[i] = 0

    return abundances


def PASiC(simiMat, counts, N):
    """
    Call similarity_correction and additionally calculate normalized abundances.

    Input:
    simiMat: [numpy.array (M,M)] with pairwise similarities between species
    counts:  [numpy.array (M,)] with number of observed PSMs for each species
    N:       [int] total number of spectra

    Output:
    abundances: [numpy.array (M,)] estimated abundance of each species in the sample
    normAbundances: [numpy.array (M,)] abundances normalized to one
    """
    abundances = similarity_correction(sim=simiMat, reads=counts, N=N)
    normAbundances = np.divide(abundances, float(sum(abundances)))
    return abundances, normAbundances

'''if __name__ == "__main__":
    # Idee: np-Dateien einlesen von Eingabepfaden -d common directory
    #c = np.load("/data/NG4/anke/spectra/acid_mine/counts_Biofilm.dat")
    #M = np.load("/data/NG4/anke/spectra/acid_mine/unw_similarityMatrix_seq.dat")
    #M_wtd = np.load("/data/NG4/anke/spectra/acid_mine/wtd_similarityMatrix_seq.dat")
    corrAbunds, normCorr = PASiC(simiMat=M, counts=c, N=10001)
    wtdCorr, n_wtdCorr = PASiC(simiMat=M_wtd, counts=c, N=10001)
    print "observed  abundances: ", np.round(np.divide(c, 10001),5)
    print "observed, normalized: ", np.round(np.divide(c, float(sum(c))),5)
    print "corrected  abundances:", np.round(corrAbunds,5)
    print "corrected, normalized:", np.round(normCorr,5)
    print "with weighting"
    print "corrected  abundances:", np.round(wtdCorr,5)
    print "corrected, normalized:", np.round(n_wtdCorr,5)
'''