import healpix_util as hu
import astropy as ap
import numpy as np
from astropy.io import fits
from astropy.table import Table
import astropy.io.ascii as ascii
from astropy.io import fits
from astropy.constants import c
import matplotlib.pyplot as plt
import math as m
from math import pi
#from scipy.constants import c
import scipy.special as sp
from astroML.decorators import pickle_results
from scipy import integrate
import warnings
from sklearn.neighbors import BallTree
import pickle
import multiprocessing as mp
import time
from lcdmmetric import *
from progressbar import *
from tqdm import *
from functools import partial
import pymangle
from scipy.optimize import curve_fit
#from astroML.datasets import fetch_sdss_specgals
#from astroML.correlation import bootstrap_two_point_angular
%matplotlib inline

dr12q=fits.open("./input/DR12Q.fits")

dr12q

dr12qdat=dr12q[1].data

dr12qdat.columns

z=dr12qdat['Z_PIPE']

ra=dr12qdat['RA']

dec=dr12qdat['DEC']

fdata = open("./output/DR12Q.dat",'w')
fdata.write("z\t ra\t dec \n")
for i in range(0,len(ra)-1):
    fdata.write("%f\t" %z[i])
    fdata.write("%f\t" %ra[i])
    fdata.write("%f\n" %dec[i])
fdata.close()

z

ra,dec

dat=np.array([z,ra,dec])

dat

dat=dat.transpose()

dat

"""
Explanation: Correlation function of DR12Q SDSS Catalog
First import all the modules such as healpy and astropy needed for analyzing the structure
End of explanation
"""


# Saving the objects:
with open('datDR12Q.pkl', 'w') as f:  # Python 3: open(..., 'wb')
    pickle.dump(dat, f)

# Getting back the objects:
with open('datDR12Q.pkl') as f:  # Python 3: open(..., 'rb')
    dat = pickle.load(f)
dat

Ez = lambda x: 1/m.sqrt(0.3*(1+x)**3+0.7)

np.vectorize(Ez)
#Calculate comoving distance of a data point using the Redshift - This definition is based on the cosmology model we take. Here the distance for E-dS universe is considered. Also note that c/H0 ratio is cancelled in the equations and hence not taken.

def DC_LCDM(z):
  return integrate.quad(Ez, 0, z)[0]
DC_LCDM=np.vectorize(DC_LCDM)

dr12f = open("./output/DR12Qsrarf.dat",'w')
dr12f.write("z\t ra\t dec\t s\t rar\t decr \n")

for i in range(0,len(z)):
    dr12f.write("%f\t " %z[i])
    dr12f.write("%f\t %f\t " %(ra[i],dec[i]))
    dr12f.write("%f\t " %DC_LCDM(z[i]))
    dr12f.write("%f\t %f\n " %(ra[i]*pi/180.0,dec[i]*pi/180.0))
dr12f.close()

dr12Qdat=ascii.read("./output/DR12Qsrarf.dat")

s=dr12Qdat['s']
rar=dr12Qdat['ra']
decr=dr12Qdat['dec']

dat=np.array([s,rar,decr])

dat

dat=dat.transpose()

dat

"""
Explanation: Read the data file (taken from http://cosmo.nyu.edu/~eak306/SDSS-LRG.html ) converted to ascii with comoving distance etc. in V01 reading from pkl files for faster read
End of explanation
"""


# Saving the objects:
with open('datDR12Q.pkl', 'w') as f:  # Python 3: open(..., 'wb')
    pickle.dump(dat, f)

# Getting back the objects:
with open('datDR12Q.pkl') as f:  # Python 3: open(..., 'rb')
    dat = pickle.load(f)
dat

bins=np.arange(0.,0.08,0.005)

print bins

binsq=bins**2

binsq

len(dat)

LCDMmetricsq(dat[0],dat[1])

%%time
BTD = BallTree(dat,metric='pyfunc',func=LCDMmetricsq,leaf_size=5) 

with open('BTDDR12QLCDM.pkl', 'w') as f:
    pickle.dump(BTD,f)

with open('BTDDR12QLCDM.pkl') as f:
    BTD = pickle.load(f)
    
BTD

%%time
start_time=time.time()
counts_DD=BTD.two_point_correlation(dat,binsq)
print counts_DD
end_time=time.time()
tottime=end_time-start_time
print "Total run time:"
print tottime

with open('BTD12QcDDLCDM.pkl', 'w') as f:
    pickle.dump(counts_DD,f)

with open('BTD12QcDDLCDM.pkl') as f:
    counts_DD = pickle.load(f)
    
counts_DD

DD=np.diff(counts_DD)

DD

plt.plot(bins[1:len(bins)],DD,'ro-')

"""
Explanation: Read the data file (taken from http://cosmo.nyu.edu/~eak306/SDSS-LRG.html ) converted to ascii with comoving distance etc. in V01 reading from pkl files for faster read
End of explanation
"""


dataR=ascii.read("./output/rand200kDR12Q.dat")

dataR

len(dataR)

len(dat)

rdr12f = open("./output/DR12Qsrarf.dat",'w')
rdr12f.write("z\t ra\t dec\t s\t rar\t decr \n")

for i in range(0,len(dataR)):
    rdr12f.write("%f\t " %dataR['z'][i])
    rdr12f.write("%f\t %f\t " %(dataR['ra'][i],dataR['dec'][i]))
    rdr12f.write("%f\t " %DC_LCDM(dataR['z'][i]))
    rdr12f.write("%f\t %f\n " %((dataR['ra'][i]*pi)/180.0,(dataR['dec'][i]*pi)/180.0))
rdr12f.close()

datR=ascii.read("./output/DR12Qsrarf.dat")

datR

rs=np.array(datR['s'])
rrar=np.array(datR['rar'])
rdecr=np.array(datR['decr'])

datR=np.array([rs,rrar,rdecr])

datR

datR.reshape(3,len(dataR))

datR=datR.transpose()

datR

# Saving the objects:
with open('./output/rDR12Qsrarf.pkl', 'w') as f:  # Python 3: open(..., 'wb')
    pickle.dump(datR, f)

# Getting back the objects:
with open('./output/rDR12Qsrarf.pkl') as f:  # Python 3: open(..., 'rb')
    datR = pickle.load(f)
datR

%%time
BT_R2 = BallTree(datR,metric='pyfunc',func=LCDMmetricsq,leaf_size=5) 

with open('./output/BTRDR12QLCDM.pkl', 'w') as f:
    pickle.dump(BT_R2,f)

with open('./output/BTRDR12QLCDM.pkl') as f:
    BTR = pickle.load(f)
BTR

%%time
start_time=time.time()
counts_RR=BTR.two_point_correlation(datR,bins)
print counts_RR
end_time=time.time()
tottime=end_time-start_time
print "Total run time:"
print tottime

with open('./output/BTRDR12QcRRLCDM.pkl', 'w') as f:
    pickle.dump(counts_RR,f)

with open('./output/BTRDR12QcRRLCDM.pkl') as f:
    counts_RR = pickle.load(f)
    
counts_RR

counts_RR

RR=np.diff(counts_RR)

RR

plt.plot(bins[1:len(bins)],RR,'bo-')

RR_zero = (RR == 0)
RR[RR_zero] = 1

%%time
start_time=time.time()
counts_DR=BTR.two_point_correlation(dat,bins)
print counts_DR
end_time=time.time()
tottime=end_time-start_time
print "Total run time:"
print tottime

with open('./output/BTRDR12QcDRLCDM.pkl', 'w') as f:
    pickle.dump(counts_DR,f)

with open('./output/BTRDR12QcDRLCDM.pkl') as f:
    counts_DR = pickle.load(f)
    
counts_DR

DR=np.diff(counts_DR)

corrells=(4.0 * DD - 4.0 * DR + RR) / RR

DR

corrells

plt.plot(bins[1:len(bins)],corrells,'go-')

plt.plot(bins[1:len(bins)],bins[1:len(bins)]*bins[1:len(bins)]*corrells*(c*1e-5)**2,'go-')

plt.plot(bins[2:len(bins)],bins[2:len(bins)]*bins[2:len(bins)]*corrells[1:len(bins)]*(c*1e-5)**2,'go-')

plt.plot(bins[2:len(bins)],corrells[1:len(bins)],'go-')

plt.plot(bins[2:len(bins)],corrells[1:len(bins)],'go-')
plt.savefig("correlDR12Qls.pdf")

plt.plot(bins[2:len(bins)]*c/1e5,corrells[1:len(bins)],'bo-')
plt.savefig("correl12Q1ls.pdf")

plt.yscale('log')
plt.plot(bins[1:len(bins)]*c/1e5,corrells,'bo-')
plt.savefig("correllsfiglogDR12Q.pdf")

plt.yscale('log')
plt.plot(bins[2:len(bins)]*c/1e5,corrells[1:len(bins)],'ro-')
plt.savefig("correllslog2xDR12Q.pdf")

plt.yscale('log')
plt.xscale('log')
plt.plot(bins[1:len(bins)]*c/1e5,corrells,'bo-')
plt.savefig("correllsloglogDR12Q.pdf")

"""
Explanation: BallTree.two_point_correlation works almost 10 times faster! with leaf_size=5 Going with it to the random catalog
End of explanation
"""


