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

modelNames = ["NoInt", "Int_u-9"]
filters    = ["tophat", "kspace"]
filter_ks  = [2.5, 2.0]

cols   = ["k", "r"]
lstyle = ["-","-."]
alphas = [1.0, 0.6]

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

rhocrit       = 27.755e10
sigma8        = 0.816
mass_in_8     = 4/3 * np.pi * 8**3 * rhocrit * Om

def normed_sigma(data_file, target_sigma):

    "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)
    norm = f(mass_in_8)
    sigma = sigma / norm * target_sigma
    return mass, sigma


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

    nu = (deltaC / sigma) ** 2
    first = -0.5 * (rhocrit*Om / 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

# Figure
plt.close("all")
fig = plt.figure(figsize=(7,5))
ax  = fig.add_subplot(111)
ax.set_xlim(1e6, 3e15)
ax.set_ylim(1e5, 4e9)
ax.set_xlabel(r"Mass")
ax.set_ylabel(r"HMF")


for ii in range(len(modelNames)):

    for jj in range(len(filters)):

        if filters[jj] == "tophat":
            mass, sigma = normed_sigma("../Data/filter_tests/sigma_M_%s_%s.dat"%(modelNames[ii],filters[jj]), target_sigma=sigma8)
            dlnsigma    = np.genfromtxt("../Data/filter_tests/sigma_M_%s_%s.dat"%(modelNames[ii],filters[jj]), usecols=2)

            MF = sheth_tormen(mass, sigma, dlnsigma, q, A, 1.686)

            ax.loglog(mass, mass * MF, color = cols[ii], ls = lstyle[jj], label = r"%s %s"%(modelNames[ii], filters[jj]))

        if filters[jj] == "kspace":
            for mm in range(len(filter_ks)):
                mass, sigma = normed_sigma("../Data/filter_tests/sigma_M_%s_%s_ks_%2.1f.dat"%(modelNames[ii],filters[jj],filter_ks[mm]), target_sigma=sigma8)
                dlnsigma    = np.genfromtxt("../Data/filter_tests/sigma_M_%s_%s_ks_%2.1f.dat"%(modelNames[ii],filters[jj],filter_ks[mm]), usecols=2)

                MF = sheth_tormen(mass, sigma, dlnsigma, q, A, 1.686)
        
                ax.loglog(mass, mass * MF, color = cols[ii], ls = lstyle[jj], alpha = alphas[mm], label = r"%s %s ks=%2.1f"%(modelNames[ii], filters[jj], filter_ks[mm]))

ax.legend(loc="lower right")
plt.savefig("test.pdf")
#plt.show()