import posteriors
import pylab
from numpy import pi

lcdm_data=posteriors.load('lcdm_prediction.dat',1,3)
lcdm_ell=lcdm_data[0].astype(int)
lcdm_cl=lcdm_data[1]
lcdm_start=lcdm_ell[0]
maxipol_low=151
maxipol_high=693
maxipol_mid=(maxipol_high+maxipol_low)/2
maxipol_T_cal_err=0.13

shape_functions={
    'square': lambda x: 4*pi/(2*x+1),
    'flat' : lambda x: x,
    'llp1': lambda x: 2*pi/(x*x+x),
    'lcdm': lambda x: lcdm_cl[x-lcdm_start]
    }

pretty_names={
    "square":"(2*l+1)/4pi",
    "flat" : "Flat",
    "llp1" : "l(l+1)/2pi",
    "lcdm" : "LCDM"
    } 

def all_any_stats():
    bins=[2,5,8]
    bin_names={2:'EE',5:'BB',8:'EB'}
    all_results={}
    for bin in bins:
        name=bin_names[bin]
        all_results[name]=all_bin_stats(bin)
    return all_results

def all_bin_stats(bin):
    shapes=shape_functions.keys()
    calibrate=[True,False]
    methods={'maxlike':posteriors.maxlike,'post':posteriors.postlike}
    results={}
    for shape in shapes:
        for method in methods.keys():
            for calibrating in calibrate:
                x,y=posteriors.load('plot_data/'+shape+'_p'+str(bin)+'.dat',1,2)
                x,y=posteriors.normalize(x,y)
                if calibrating:
                    x,y=posteriors.calibration(x,y,maxipol_T_cal_err)
                results[(pretty_names[shape],method,calibrating and "Calibrated" or "Uncalibrated")]=methods[method](x,y)
    return results


def output_bin_stats(r,filename):
    #f=open('filename','w')
    for shape in pretty_names.values():
            for cal in ['Calibrated','Uncalibrated']:
                print ("\nShape:%s,%s\n"%(shape,cal))
                ml=r[(shape,'maxlike',cal)]
                pl=r[(shape,'post',cal)]
                print ('mean=%d\nmode=%d\nMax-like 68%% range=%d to %d\nMax-like 95%% range=%d to %d\nStandard 68%% range=%d to %d\nStandard 95%% range=%d to %d\n\n' \
                        % (pl[2]   ,   ml[2],                 ml[1], ml[3],                ml[0], ml[4],                pl[1], pl[3],                pl[0],pl[4]))
    #f.close()


def allstats(base,spectra=["EE","BB","EB"],bins=["1","2","3"],units="\ell(\ell+1)C_\ell/2\pi uK^2",mode="post",plotting=False):
    base_path="plot_data/"+base
    filenames={}
    fn=1
    if mode == "post":
        get_limits=posteriors.postlike
    elif mode=="onetail":
        get_limits=posteriors.onetail
    else :
        get_limits=posteriors.maxlike
        
    for spectrum in spectra:
        for bin in bins:
            filenames[spectrum+bin]=base_path+"_p"+str(fn)+".dat"
            fn+=1
    limits={}
    sp=1
    for spectrum in spectra:
        for bin in bins:
            filename=filenames[spectrum+bin]
            x,y=posteriors.load(filename,1,2)
            pylab.subplot(len(spectra),len(bins),sp)
            results=get_limits(x,y,plotting=plotting,stepfactor=300)
            ax=pylab.gca()
            ax.set_xlabel(spectrum+" "+bin,fontsize=8)
            pylab.xticks(fontsize=8)
            pylab.yticks([])
            limits[spectrum+bin]=results
            sp+=1
    return limits

def get_lcdm_prediction(shape,bin_low,bin_high):
    """
    Returns the binned data dor a given shape function (ie in C_l)
    """
    f=shape_functions[shape]
    l=lcdm_ell
    cl=lcdm_cl
    start=l[0]
    index_low=bin_low-start
    index_high=bin_high-start
    #This should be plain C_l, no factors
    mid=(index_high+index_low)/2
    s=f(l)
    result=sum(cl[index_low:index_high]/s[index_low:index_high])*s[mid]/(index_high-index_low)
    
    return result
    
def allnumbers():
    shapes=shape_functions.keys()
    calibrations=[True,False]
    calibration_error=0.13
    results={}
    norm=(maxipol_mid+maxipol_mid*maxipol_mid)/(2*pi)
    for shape in shapes:
        prediction=get_lcdm_prediction(shape,maxipol_low,maxipol_high)*norm
        for calibrating in calibrations:
            x,y=posteriors.load('plot_data/'+shape+'_p2.dat',1,2)
            if calibrating:
                x,y=posteriors.calibration(x,y,maxipol_T_cal_err)
            x,y=posteriors.normalize(x,y)
            p=posteriors.interp(x,y)
            zero_integral=posteriors.tab_area(x,y,x[0],0.0)
            lcdm_integral=posteriors.tab_area(x,y,x[0],prediction)
            zero_ratio=(p(0.0)/max(y))[0]
            lcdm_ratio=(p(prediction)/max(y))[0]
            mode=x[y.argmax()]
            results[(shape,calibrating and "Calibrated" or "Uncalibrated")]=(zero_integral,zero_ratio,lcdm_integral,lcdm_ratio,prediction)
    return results



def output_results(r,filename):
    r2={}
    for k in r.keys():
        r2[k]=[100*x for x in r[k]]
        r2[k][4]=r2[k][4]/100
    f=open(filename,'w')
    for k in r2.keys():
        name=pretty_names[k[0]]
        f.write('Shape: %s, %s\n' % (name,k[1]))
        f.write('P(C_l<0):\t\t\t %.1f%%\nP(C_l=0)/P(C_l=mode):\t\t %.1f%%\nP(C_l<LCDM):\t\t\t %.1f%%\nP(C_l=LCDM)/P(C_l=mode):\t %.1f%%\n\n\n'%tuple(r2[k][0:4]))
    f.close()