#import astropy
#import h5py
import math
import numpy as np
#from scipy.constants import c 
#from astropy.io import fits
import subprocess
import statistics as stat
from operator import mul
from operator import truediv

# plotting variables
from setup_matplotlib import *
from matplotlib.ticker import MaxNLocator
width = 8.8
# Create the plot
fig = plt.figure(figsize=(1.25*cm2inch(width), 6./8.*cm2inch(width)))
#end plotting variables

def get_invvarweighted(x,y) :
    return sum(x/y/y)/sum(1./y/y);


#initial variables
dire='/mn/stornext/d5/unnif/sindex_bp/coswmap23_coswmap33ab/combab_011-195/'
pre='b_'     #b_
filpre='ut'  #ut
alphaarr=('a00000','a00005','a00010','a00015','a00020','a00025','a00030','a00035','a00040','a00045','a00050','a00055','a00060','a00065','a00070','a00075','a00080','a00085')
alphanum=[]
for alphas in range(0,90,5):
    alphanum.append(alphas)
print(alphanum)

meanfilnavn=dire+filpre+'_sample_betas_mean.txt'
regmeanfile = open(meanfilnavn, 'w') 
invvarfilnavn=dire+filpre+'_sample_betas_invvar.txt'
reginvvarfile = open(invvarfilnavn, 'w') 

#howmanyregions=0
reg_betas=[]
reg_uncerts=[]

#the loop over regions
for reg in range(1,25):
    betas=[]
    uncerts=[]
    alphabetas_mat=[]
    alphauncerts_mat=[]
    #the loop over samples (cosmoglobe)
    for samp in range(11,25):
        print('reg='+str(reg)+' samp='+str(samp))
        #read the spectral index in region
        datadir=dire+'/../s0'+str(samp)+'/'+pre+str(reg)+'/'
        spfilnavn=datadir+filpre+'_value_beta_scatter.txt'
        #print(spfilnavn)
        with open(spfilnavn) as spfil:
            beta=float(spfil.read())
        print(beta)
        betas.append(beta)

        #read the uncert in region
        alphafilnavn=datadir+filpre+'_scatter_alpha_spectral_index.txt'
        u=[]
        b=[]
        with open(alphafilnavn) as f:
            for line in f:
                data = line.split()
                b.append(float(data[1]))
                u.append(float(data[2]))
        uncert=min(u)
        print(uncert)
        #could write out the uncert to a file here
        uncerts.append(uncert)

        #write the sample alpha values to a array of arrays
        alphabetas_mat.append(b)
        alphauncerts_mat.append(u)

    #end of sample loop
    ############################################
    # stuff for the (main) sample big alpha and scatter plot
    ############################################
    #ut_sample_scatter_alpha_spectral_index.txt
    nalphabetas_mat=np.array(alphabetas_mat)
    nalphauncerts_mat=np.array(alphauncerts_mat)
    print(nalphabetas_mat)
    print(nalphabetas_mat.shape) #shape har formen (#samples,#alphavalues)

    # creating files
    samplealphafilnavn=dire+pre+str(reg)+'/'+filpre+'_sample_scatter_alpha_spectral_index.txt'
    samplebetavaluefilnavn=dire+pre+str(reg)+'/'+filpre+'_sample_value_beta_scatter.txt'
    print(samplealphafilnavn)
    reg_alphafile = open(samplealphafilnavn, 'w') 

    # looping over the 18 alpha values
    for a in range(0,18):
        print(a)
        reg_alpha_beta=get_invvarweighted(nalphabetas_mat[:,a],nalphauncerts_mat[:,a])
        reg_alpha_uncert=min(nalphauncerts_mat[:,a]) +stat.stdev(nalphabetas_mat[:,a])
        print(reg_alpha_beta)
        print(reg_alpha_uncert)
        print(min(nalphauncerts_mat[:,a]) ,stat.stdev(nalphabetas_mat[:,a]))
        #writeto file
        reg_alphafile.write(str(alphanum[a])+'\t'+str(reg_alpha_beta)+'\t'+str(reg_alpha_uncert)+'\n')

    ############################
    # stuff for the region plot
    ############################
    print(" region plot values")
    s_betas=np.array(betas)
    s_uncerts=np.array(uncerts)
    final_beta=sum(s_betas/s_uncerts/s_uncerts)/sum(1./s_uncerts/s_uncerts)
    invvar_uncert=math.sqrt(1./(sum(1./s_uncerts/s_uncerts)))
    mean_beta=stat.mean(s_betas)
    min_uncert=min(s_uncerts)
    print('uncert(min)='+str(min_uncert))
    std_beta=stat.stdev(s_betas)
    print('std='+str(std_beta))
    final_uncert=math.sqrt(min_uncert*min_uncert + std_beta*std_beta)
    print('final uncert (quarature)='+str(final_uncert))

    print('The straight mean beta='+str(stat.mean(s_betas))+str(min(s_uncerts)))
    print('The inverse weighted beta='+str(final_beta)+'+-'+str(final_uncert))
    #print('howmanyregions='+str(howmanyregions))

    regmeanfile.write(str(reg)+'\t'+str(mean_beta)+'\t'+str(min(s_uncerts))+'\n')
    reginvvarfile.write(str(reg)+'\t'+str(final_beta)+'\t'+str(final_uncert)+'\n')

    #write beta to ut_sample_value_beta_scatter.txt
    sample_betafile=open(samplebetavaluefilnavn, 'w')
    sample_betafile.write(str(final_beta))

    #add them to a region array, for plotting purposes
    reg_betas.append(final_beta)
    reg_uncerts.append(final_uncert)

    #plot them
    if reg==1:
        lab=r"Cosmoglobe T-T plot {\it WMAP}"
    else:
        lab=""
    plt.errorbar(reg,final_beta, yerr=final_uncert, marker='x', label=lab,linewidth=0.5, color='red',markersize=3,capsize=2,ls='none')

regmeanfile.close()
reginvvarfile.close()
print("All beta values:")
print(reg_betas)
print("All uncert values:")
print(reg_uncerts)
print("** FINISHED LOOPING OVER REGIONS **")

#calculating the mean of cosmoglobe betas
c_betas=np.array(reg_betas)
c_uncerts=np.array(reg_uncerts)
#cos_invvar_beta=sum(c_betas/c_uncerts/c_uncerts)/sum(1./c_uncerts/c_uncerts)
cos_invvar_beta=get_invvarweighted(c_betas,c_uncerts)
cos_mean_beta=stat.mean(c_betas)
cos_min_uncert=min(c_uncerts)
print('The straight mean beta='+str(cos_mean_beta)+' +-'+str(cos_min_uncert))
print('The inverse weighted beta='+str(cos_invvar_beta)+' +-'+str(cos_min_uncert))
plt.axhline(y=cos_invvar_beta, color='red',linestyle='-',linewidth=0.5)

#plot the old wmap
# this is the combined ttplot and maxlike, do not use this!
#oldwmapbetafilnavn="/mn/stornext/d5/unnif/unnif-mirfakd1/spectral/w9_n64_1deg_kopi/ut_paper_finalfinalregion_alphamean.txt"
# this is the old ttplot wmap result
oldwmapbetafilnavn="/mn/stornext/d5/unnif/sindex_bp/wmap23_wmap33_bootn/ut_region_scatter_alphamean.txt"
with open(oldwmapbetafilnavn) as w:
    rx=[]
    bx=[]
    ux=[]
    for line in w:
        data = line.split()
        rx.append(float(data[0])+0.1)  #plot the point next to the first
        bx.append(float(data[1]))
        ux.append(float(data[2]))

plt.errorbar(rx,bx, yerr=ux, marker='x', label=r"Original T-T plot {\it WMAP}",linewidth=0.5, color='black',markersize=3,capsize=2,ls='none')
plt.axhline(y=-2.99, color='black',linestyle='--',linewidth=0.5) #-2.99 is the old inv var weigh mean

# general plot parameters
plt.ylim(-4.,-2.2)
#plt.yticks(-4., -2.2, step=0.2)
plt.yticks([-4.,-3.8, -3.6, -3.4, -3.2, -3.0, -2.8, -2.6, -2.4, -2.2],[-4.0,-3.8, -3.6, -3.4, -3.2, -3.0, -2.8, -2.6, -2.4, -2.2])
plt.ylabel(r"Spectral index", fontsize=10);plt.xlabel(r"Region number", fontsize=10);
leg = plt.legend(frameon=True, loc='lower right', prop={'size':7})
# remove box around legend
leg.get_frame().set_edgecolor("white")
leg.get_frame().set_alpha(.0)
plt.show()
plt.savefig("region_beta_cosmoglobe_vs_wmap.png", bbox_inches='tight', bbox_extra_artists=[],pad_inches=0.03)


#ut_sample_scatter_alpha_spectral_index.txt