#!python
import numpy as np
from numpy import inf
from numpy import nan
from scipy.optimize import fmin
from scipy.stats import beta
from scipy.special import beta as B
from scipy.special import comb
import argparse
import sys

def parseArgs():
    '''Function for parsing arguments'''
    parser = argparse.ArgumentParser(description="Pipeline for analyzing barcoded amplicon \
                                                  sequencing data with Unique molecular \
                                                  identifiers (UMI)")
    parser.add_argument('-cons', '--cons_file', dest='cons_file', help='Path to cons file, for fitting parameters of the bgmodel')
    parser.add_argument('-nonbgposfile', '--non-background-positions', dest='nonbgposfile',
                        help='Path to file with non-background positions')
    parser.add_argument('-out', '--out_file',dest='out_file',help="name of output file, default = %(default)s]",default="bgmodel.params") 
    parser.add_argument('-f','--fsize',dest='fsize', help='Family size cutoff (consensus cutoff) for variant calling. [default = %(default)s]', default=3)
    args = parser.parse_args(sys.argv[1:])
    return(args)

def parse_cons_file(filename,fsize=3):
    n1=[]
    f1=[]
    c1=[]
    posx=[]
    data=[]
    with open(filename) as f:
        for line in f:
            if not line.startswith('Sample Name'):
                line=line.rstrip('\n')
                parts=line.split('\t')
                pos=parts[1]+':'+parts[2]
                name=parts[3]
                #print(name)
                if name not in "":
                    famsize=parts[-4]
                    if int(famsize)==fsize:
                        frac=float(parts[-2])
                        alt=parts[-1]
                        count=parts[-3]
                        if frac > 0 and alt not in 'N':
                            cov=int(parts[-5])
                            f1.append(float(frac))
                            n1.append(int(cov))
                            c1.append(int(count))
                            posx.append(pos)
                            data.append(line)
                #print(name)
                #print(famsize)
    return(f1,n1,c1,posx,data)

def betaNLL(params,*args):
    a,b = params
    data = np.array(args[0])
    pdf=beta.pdf(data,a,b,loc=0,scale=1)
    lg=np.log(pdf)
    #lg=np.where(lg==-np.inf,0,lg)
    mask = np.isfinite(lg)
    nll = -lg[mask].sum()
    nll=-1*np.sum(lg)
    return(nll)


def get_beta_parameters(data):
    m=np.mean(data)
    v=np.var(data)
    a0=m*(m * (1-m) / v-1 )
    b0=(1-m)*(m * (1-m) / v-1 )
    result=fmin(betaNLL,[a0,b0],args=(data,))
    return(result)

def run_fit_bgmodel(args):
    spikepositions=[178952085,55599321,7577558,7577547,7577538,7577120]
    if args.nonbgposfile:
        nonbgpos=[]
        with open(args.nonbgposfile) as f:
            for line in f:
                line=line.rstrip()
                nonbgpos.append(line)
    else:
        nonbgpos=spikepositions
    if not args.cons_file:
        args.cons_file=glob.glob(args.output_path+'/*cons.tsv')[0]
    args.fsize=int(args.fsize)
    f1,n1,a1,pos,data=parse_cons_file(args.cons_file,args.fsize)
    f1 = np.array(f1)
    n1 = np.array(n1)
    a1 = np.array(a1)
    pos = np.array(pos)
    data = np.array(data)
    result=get_beta_parameters(f1[np.isin(pos,nonbgpos)!=True])
    #a=prob_bb(n1,a1,result[0],result[1])
    print(pos,nonbgpos,np.isin(pos,nonbgpos))
    with open(args.out_file,'w') as g:
        g.write('{}\n'.format(result[0]))
        g.write('{}\n'.format(result[1]))
    #a[a==inf]=1e-10
    #a[np.isnan(a)]=1e-10
    #Q = -10*np.log10(a)
    #data=np.array(data)
    #plot_histogram(Q,args.output_path+'/'+args.sample_name+'.histogram.png')
    #if args.vc_method.lower()=='bbmodel':
    #    rout=data[Q >= float(args.qvalue_threshold)]
    #    Qsig=Q[Q >= float(args.qvalue_threshold)]
    #else:
    #    rout=data[a1 >= float(args.count_cutoff)]
    #    Qsig=Q[a1 >= float(args.count_cutoff)]
    #outfilename=args.output_path+'/'+args.sample_name+'2.vcf'
    #write_vcf(outfilename,rout,Qsig,args.reference_file)

if __name__=='__main__':
    args=parseArgs()
    run_fit_bgmodel(args)

[((14, 13, 16, 69), 'argparse.ArgumentParser', 'argparse.ArgumentParser', (), '', False, 'import argparse\n'), ((58, 11, 58, 28), 'numpy.array', 'np.array', ({(58, 20, 58, 27): 'args[0]'}, {}), '(args[0])', True, 'import numpy as np\n'), ((59, 8, 59, 40), 'scipy.stats.beta.pdf', 'beta.pdf', (), '', False, 'from scipy.stats import beta\n'), ((60, 7, 60, 18), 'numpy.log', 'np.log', ({(60, 14, 60, 17): 'pdf'}, {}), '(pdf)', True, 'import numpy as np\n'), ((62, 11, 62, 26), 'numpy.isfinite', 'np.isfinite', ({(62, 23, 62, 25): 'lg'}, {}), '(lg)', True, 'import numpy as np\n'), ((69, 6, 69, 19), 'numpy.mean', 'np.mean', ({(69, 14, 69, 18): 'data'}, {}), '(data)', True, 'import numpy as np\n'), ((70, 6, 70, 18), 'numpy.var', 'np.var', ({(70, 13, 70, 17): 'data'}, {}), '(data)', True, 'import numpy as np\n'), ((73, 11, 73, 45), 'scipy.optimize.fmin', 'fmin', (), '', False, 'from scipy.optimize import fmin\n'), ((90, 9, 90, 21), 'numpy.array', 'np.array', ({(90, 18, 90, 20): 'f1'}, {}), '(f1)', True, 'import numpy as np\n'), ((91, 9, 91, 21), 'numpy.array', 'np.array', ({(91, 18, 91, 20): 'n1'}, {}), '(n1)', True, 'import numpy as np\n'), ((92, 9, 92, 21), 'numpy.array', 'np.array', ({(92, 18, 92, 20): 'a1'}, {}), '(a1)', True, 'import numpy as np\n'), ((93, 10, 93, 23), 'numpy.array', 'np.array', ({(93, 19, 93, 22): 'pos'}, {}), '(pos)', True, 'import numpy as np\n'), ((94, 11, 94, 25), 'numpy.array', 'np.array', ({(94, 20, 94, 24): 'data'}, {}), '(data)', True, 'import numpy as np\n'), ((64, 11, 64, 21), 'numpy.sum', 'np.sum', ({(64, 18, 64, 20): 'lg'}, {}), '(lg)', True, 'import numpy as np\n'), ((97, 23, 97, 44), 'numpy.isin', 'np.isin', ({(97, 31, 97, 34): 'pos', (97, 35, 97, 43): 'nonbgpos'}, {}), '(pos, nonbgpos)', True, 'import numpy as np\n'), ((95, 34, 95, 55), 'numpy.isin', 'np.isin', ({(95, 42, 95, 45): 'pos', (95, 46, 95, 54): 'nonbgpos'}, {}), '(pos, nonbgpos)', True, 'import numpy as np\n')]