#!/usr/bin/env python
from __future__ import division
import matplotlib.pyplot as plt
from scipy.interpolate import interp1d
import numpy as np

A = 0.3222
p = 0.3
q = 0.707
qw = 0.5
deltaC = 1.686
Om_noint = 0.272
Om_int = 0.272

correct_h0 = True

if correct_h0:
    h0_noint = 0.704
    h0_int = 0.704
else:
    h0_noint = 1.0
    h0_int = 1.0


rhocrit_noint = 27.755e10
rhocrit_int = rhocrit_noint
sigma8_int = [0.76, 0.8, 0.8, 0.81, 0.81, 0.81]
sigma8_noint = 0.81
#sigma8_int = [0.81]*6
#sigma8_noint = 0.81
mass_in_8_int   = 4/3 * np.pi * 8**3 * rhocrit_int * Om_int
mass_in_8_noint = 4/3 * np.pi * 8**3 * rhocrit_noint * Om_noint

M_10 =[1e14, 5.3e12, 2.3e11, 8.3e9, 2.9e8, 9.6e6]

def Markus_fit_ratio(mass, M_10):
    alpha = 0.9
    return 1.0/(1.0 + (M_10/mass)**alpha)

def normed_sigma(data_file, target_sigma, int=False):

    "Calcualate the sigma(M) normalised to sigma_8 for this cosmology."

    sigma = np.genfromtxt(data_file, usecols = 1)
    mass = np.genfromtxt(data_file, usecols=0)
    f = interp1d(mass, sigma)
    if int:
        norm = f(mass_in_8_int)
        sigma = sigma / norm * target_sigma
    else:
        norm = f(mass_in_8_noint)
        sigma = sigma / norm * target_sigma
    #print(mass_in_8)
    return mass, sigma


def sheth_tormen(mass, sigma, dlnsigma, q, A, deltaC, int=False):

    nu = (deltaC / sigma) ** 2
    if int:
        first = -0.5 * (rhocrit_int*Om_int / mass)
    else:
        first = -0.5 * (rhocrit_noint*Om_noint / mass)
    second = A * np.sqrt(2 * q / np.pi * nu) * (1 + (q * nu) ** (-p )) * np.exp(-q * nu / 2.)
    third = dlnsigma

    dndlM = first * second * third

    return dndlM

# CDM
mass_cdm, sigma_cdm = normed_sigma("../Data/NoInt_extrap.dat.spline", target_sigma=sigma8_noint)
mass_cdm, sigma_cdm = mass_cdm[2:], sigma_cdm[2:]
dlnsigmaC = np.genfromtxt("../Data/NoInt_extrap.dat.spline", usecols=3, skip_header=2)
'''
# Int-u4
mass_u4, sigma_u4  = normed_sigma("../Data/u-4_extrap.dat.spline.FiltFac_beta1.00", target_sigma=sigma8_int[0], int=True)
mass_u4, sigma_u4 = mass_u4[2:], sigma_u4[2:]
dlnsigmau4 = np.genfromtxt("../Data/u-4_extrap.dat.spline.FiltFac_beta1.00", usecols=3, skip_header=2)
# Int-u5
mass_u5, sigma_u5  = normed_sigma("../Data/u-5_extrap.dat.spline.FiltFac_beta1.00", target_sigma=sigma8_int[1], int=True)
mass_u5, sigma_u5 = mass_u5[2:], sigma_u5[2:]
dlnsigmau5 = np.genfromtxt("../Data/u-5_extrap.dat.spline.FiltFac_beta1.00", usecols=3, skip_header=2)
# Int-u6
mass_u6, sigma_u6  = normed_sigma("../Data/u-6_extrap.dat.spline.FiltFac_beta1.00", target_sigma=sigma8_int[2], int=True)
mass_u6, sigma_u6 = mass_u6[2:], sigma_u6[2:]
dlnsigmau6 = np.genfromtxt("../Data/u-6_extrap.dat.spline.FiltFac_beta1.00", usecols=3, skip_header=2)
# Int-u7
mass_u7, sigma_u7  = normed_sigma("../Data/u-7_extrap.dat.spline.FiltFac_beta1.00", target_sigma=sigma8_int[3], int=True)
mass_u7, sigma_u7 = mass_u7[2:], sigma_u7[2:]
dlnsigmau7 = np.genfromtxt("../Data/u-7_extrap.dat.spline.FiltFac_beta1.00", usecols=3, skip_header=2)
# Int-u8
mass_u8, sigma_u8  = normed_sigma("../Data/u-8_extrap.dat.spline.FiltFac_beta1.00", target_sigma=sigma8_int[4], int=True)
mass_u8, sigma_u8 = mass_u8[2:], sigma_u8[2:]
dlnsigmau8 = np.genfromtxt("../Data/u-8_extrap.dat.spline.FiltFac_beta1.00", usecols=3, skip_header=2)
'''
# Int-u9
mass_u9, sigma_u9  = normed_sigma("../Data/NoInt_extrap.dat.spline.FiltFac_beta1.00", target_sigma=sigma8_int[5], int=True)
mass_u9, sigma_u9 = mass_u9[2:], sigma_u9[2:]
dlnsigmau9 = np.genfromtxt("../Data/NoInt_extrap.dat.spline.FiltFac_beta1.00", usecols=3, skip_header=2)

'''
# Int-u9 kspace 2.5                                                                                                                                                                                       
mass_u9_2p5, sigma_u9_2p5  = normed_sigma("../Data/filter_tests/sigma_M_Int_u-9_kspace_ks_2.5.dat", target_sigma=sigma8_int[5], int=True)
mass_u9_2p5, sigma_u9_2p5 = mass_u9_2p5[:], sigma_u9_2p5[:]
dlnsigmau9_2p5 = np.genfromtxt("../Data/filter_tests/sigma_M_Int_u-9_kspace_ks_2.5.dat", usecols=2, skip_header=0)

# Int-u11
mass_u11, sigma_u11  = normed_sigma("../Data/u-11_extrap.dat.spline.FiltFac_beta1.00", target_sigma=sigma8_int[5], int=True)
mass_u11, sigma_u11 = mass_u11[2:], sigma_u11[2:]
dlnsigmau11 = np.genfromtxt("../Data/u-11_extrap.dat.spline.FiltFac_beta1.00", usecols=3, skip_header=2)

# Int-u11 kspace 2.5
mass_u11_2p5, sigma_u11_2p5  = normed_sigma("../Data/u-11_extrap.dat.spline.FiltFac_beta1.00_ks2p5", target_sigma=sigma8_int[5], int=True)
mass_u11_2p5, sigma_u11_2p5 = mass_u11_2p5[2:], sigma_u11_2p5[2:]
dlnsigmau11_2p5 = np.genfromtxt("../Data/u-11_extrap.dat.spline.FiltFac_beta1.00_ks2p5", usecols=3, skip_header=2)
'''
MF_CDM = sheth_tormen(mass_cdm, sigma_cdm, dlnsigmaC, q, A, 1.686)
'''
MF_u4  = sheth_tormen(mass_u4, sigma_u4, dlnsigmau4, q, A, 1.686, int=True)
MF_u4_mod  = sheth_tormen(mass_u4, sigma_u4, dlnsigmau4, q, A*0.6, 1.686)
MF_u5  = sheth_tormen(mass_u5, sigma_u5, dlnsigmau5, q, A, 1.686, int=True)
MF_u6  = sheth_tormen(mass_u6, sigma_u6, dlnsigmau6, q, A, 1.686, int=True)
MF_u7  = sheth_tormen(mass_u7, sigma_u7, dlnsigmau7, q, A, 1.686, int=True)
MF_u8  = sheth_tormen(mass_u8, sigma_u8, dlnsigmau8, q, A, 1.686, int=True)
'''
MF_u9  = sheth_tormen(mass_u9, sigma_u9, dlnsigmau9, q, A, 1.686, int=True)
#MF_u9_2p5 = sheth_tormen(mass_u9_2p5, sigma_u9_2p5, dlnsigmau9_2p5, q, A, 1.686, int=True)

#MF_u11     = sheth_tormen(mass_u11, sigma_u11, dlnsigmau11, q, A, 1.686, int=True)
'''
MF_u11_2p5 = sheth_tormen(mass_u11_2p5, sigma_u11_2p5, dlnsigmau11_2p5, q, A, 1.686, int=True)
'''

# Markus data
#files  = ["../Data/MF_Ratio_Markus_u4.dat", "../Data/MF_Ratio_Markus_u5.dat", "../Data/MF_Ratio_Markus_u6.dat"]
colours = ["b", "g", "r", "m", "grey", "y"]

# Figure
plt.close("all")
'''
plt.plot(mass_u4/h0_int, MF_u4/MF_CDM, color = "b", label = "u4")
#plt.loglog(mass_u4/h0, MF_u4_mod/MF_CDM, color = "cyan", label = "u4-mod")
plt.plot(mass_u5/h0_int, MF_u5/MF_CDM, color = "g", label = "u5")
plt.plot(mass_u6/h0_int, MF_u6/MF_CDM, color = "r", label = "u6")
plt.plot(mass_u7/h0_int, MF_u7/MF_CDM, color = "m", label = "u7")
plt.plot(mass_u8/h0_int, MF_u8/MF_CDM, color = "grey", label = "u8")
'''
plt.plot(mass_u9/h0_int, MF_u9/MF_CDM, color = "y", label = "u9")
#plt.plot(mass_u9_2p5/h0_int, MF_u9_2p5/MF_CDM, color = "y", ls = ":", label = "u9 2.5")

#plt.plot(mass_u11/h0_int, MF_u11/MF_CDM, color = "y", label = "u11 kspace")
#plt.plot(mass_u11_2p5/h0_int, MF_u11_2p5/MF_CDM, color = "m", label = "u11 kspace 2.5")


x = np.logspace(7, 16, 40)

for jj in range(len(M_10)):
    #data_markus = np.genfromtxt(files[jj])
    #plt.scatter(data_markus[:,0], data_markus[:,1], color = colours[jj], s=25, marker="o")
    plt.plot(x, Markus_fit_ratio(x, M_10[jj]), color=colours[jj], ls="--", lw=1.0)

# Get Galform estimate
cdm_gal   = np.genfromtxt("../Data/halo_weights_CDM")

'''
for nn in range(len(M_10)):
    intu4_gal = np.genfromtxt("../Data/halo_weights_Intu%d"%(4+nn))
    plt.plot(cdm_gal[:,0], (intu4_gal[:,1]*intu4_gal[:,2])/(cdm_gal[:,1]*cdm_gal[:,2]), color = colours[nn], ls = ":")
'''

plt.xlim(1e4, 3e15)
plt.ylim(0.7, 1.2)
plt.axhline(y=1, color="k", ls="--", lw=1.0)
plt.xlabel(r"Mass")
plt.ylabel(r"Ratio")
plt.yscale("log")
plt.xscale("log")
plt.minorticks_on()
plt.legend(loc="lower right")
plt.show()
#plt.savefig("MassFunction_Ratios.pdf")