"""Module containing a stand-alone script which compares the power spectrum of ICs
to the power spectrum fed into MP-GenIC, read from CLASS format files."""
import argparse
import os
import scipy.interpolate as interp
import numpy as np
import matplotlib
matplotlib.use("PDF")
import matplotlib.pyplot as plt
from nbodykit.lab import BigFileCatalog,FFTPower

def modecount_rebin(kk, pk, modes, pkc, minmodes=250, ndesired=200):
    """Rebins a power spectrum so that there are sufficient modes in each bin"""
    assert np.all(kk) > 0
    logkk=np.log10(kk)
    mdlogk = (np.max(logkk) - np.min(logkk))/ndesired
    istart=iend=1
    count=0
    pk_div = pk /pkc(kk)
    k_list=[kk[0]]
    pk_list=[pk_div[0]]
    targetlogk=mdlogk+logkk[istart]
    while iend < np.size(logkk)-1:
        count+=modes[iend]
        iend+=1
        if count >= minmodes and logkk[iend-1] >= targetlogk:
            pk1 = np.sum(modes[istart:iend]*pk_div[istart:iend])/count
            kk1 = np.sum(modes[istart:iend]*kk[istart:iend])/count
            k_list.append(kk1)
            pk_list.append(pk1)
            istart=iend
            targetlogk=mdlogk+logkk[istart]
            count=0
    k_list = np.array(k_list)
    pk_list = np.array(pk_list) * pkc(k_list)
    return (k_list, pk_list)

class CLASSPowerSpectrum(object):
    """Class to store some routines for manipulating and storing power spectra as generated by CLASS."""
    def __init__(self, camb_matter, camb_transfer, omega0, omegab, omeganu=0):
        pk_camb = np.loadtxt(camb_matter)
        assert np.shape(pk_camb)[1] == 2
        # Build an interpolator for the matter power spectrum
        self.dpk = interp.interp1d(pk_camb[:, 0], pk_camb[:, 1], kind='cubic')
        # Build interpolators for various species of transfer functions.
        tk_camb = np.loadtxt(camb_transfer)
        self.dtk = {}
        omegacdm = omega0 - omegab - omeganu
        tdmby = (omegab * tk_camb[:,2] + omegacdm * tk_camb[:,3])
        ttot = np.array(tdmby)
        if omeganu > 0:
            ttot += omeganu * tk_camb[:,6]
        ttot /= omega0
        tdmby /= (omegab + omegacdm)

        self.dtk[1] = interp.interp1d(tk_camb[:,0], pk_camb[:,1] * (tk_camb[:, 3]/ttot)**2, kind='cubic')
        #Baryons
        self.dtk[0] = interp.interp1d(tk_camb[:,0], pk_camb[:,1] * (tk_camb[:, 2]/ttot)**2, kind='cubic')
        #DM + baryon
        self.dtk[3] = interp.interp1d(tk_camb[:,0], pk_camb[:,1] * (tdmby/ttot)**2, kind='cubic')

    def get_class_power(self, species=-1):
        """Get a matter power spectrum for DM, baryons from CAMB."""
        if species < 0:
            return self.dpk
        return self.dtk[species]

def plot_ic_power(kk_ic, Pk_ic, Pk_camb, npart, sp=1, outdir="."):
    """Make the plot"""
    #Make some useful figures
    #Check that they agree between 1/4 the box and 1/4 the nyquist frequency
    mink = np.min(kk_ic)
    imax = np.searchsorted(kk_ic, npart*mink/4)
    imin = np.searchsorted(kk_ic, mink*4)
    plt.semilogx(kk_ic, Pk_ic/Pk_camb,linewidth=2)
    plt.semilogx([kk_ic[0]*0.9,kk_ic[-1]*1.1], [0.95,0.95], ls="--",linewidth=2)
    plt.semilogx([kk_ic[0]*0.9,kk_ic[-1]*1.1], [1.05,1.05], ls="--",linewidth=2)
    plt.semilogx([kk_ic[imin],kk_ic[imin]], [0,1.5], ls=":",linewidth=2)
    plt.semilogx([kk_ic[imax],kk_ic[imax]], [0,1.5], ls=":",linewidth=2)
    plt.ylim(0.7, 1.3)
    plt.savefig(os.path.join(outdir,"ICS/PK-IC-"+str(sp)+"-diff.pdf"))
    plt.clf()
    plt.loglog(kk_ic, Pk_ic,linewidth=2)
    plt.loglog(kk_ic, Pk_camb,ls="--", linewidth=2)
    plt.ylim(ymax=Pk_camb[0]*10)
    plt.savefig(os.path.join(outdir,"ICS/PK-IC-"+str(sp)+"-abs.pdf"))
    plt.clf()
    error = Pk_ic[imin:imax]/Pk_camb[imin:imax] -1
    return error

def check_ic_power_spectra(genicfileout, camb_zstr, outdir=".", accuracy=0.07, m_nu=0):
    """Generate the power spectrum for each particle type from the generated simulation files
    and check that it matches the input. This is a consistency test on each simulation output."""
    #Generate power spectra
    output = os.path.join(outdir, genicfileout)
    #Now check that they match what we put into the simulation, from CAMB
    #Reload the CAMB files from disc, just in case something went wrong writing them.
    matterpow = os.path.join(outdir,"camb_linear/ics_matterpow_"+camb_zstr+".dat")
    transfer = os.path.join(outdir, "camb_linear/ics_transfer_"+camb_zstr+".dat")
    #Load DM catalog and try for a baryon catalog
    cats = {1: BigFileCatalog(output, dataset='1/', header='Header')}
    try:
        cats[0] = BigFileCatalog(output, dataset='0/', header='Header')
    except:
        pass
    omegab = cats[1].attrs['OmegaBaryon']
    omega0 = cats[1].attrs['Omega0']
    hubble = cats[1].attrs['HubbleParam']
    npart = int(np.round(np.cbrt(cats[1].attrs['TotNumPart'][1])))
    assert npart > 0
    cambpow = CLASSPowerSpectrum(matterpow, transfer,omega0=omega0, omegab=omegab, omeganu=m_nu/93.14/hubble**2)
    for sp in cats.keys():
        #GenPK output is at PK-[nu,by,DM]-basename(genicfileout)
        cats[sp].to_mesh(Nmesh=npart*2, window='cic', compensated=True, interlaced=True)
        pk = FFTPower(cats[sp], mode='1d', Nmesh=npart*2, dk=5.0e-6)
        #GenPK output is at PK-[nu,by,DM]-basename(genicfileout)
        #Load the power spectra
        #Convert units from kpc/h to Mpc/h
        kk_ic = pk.power['k'][1:]*1e3
        Pk_ic = pk.power['power'][1:].real/1e9
        modes_ic = pk.power['modes'][1:]
        ii = np.isfinite(kk_ic)
        kk_ic = kk_ic[ii]
        Pk_ic = Pk_ic[ii]
        #Load the power spectrum. Note that DM may be total.
        ccsp = sp
        if len(cats) == 1:
            ccsp = -1
            if m_nu > 0:
                ccsp = 3
        Pk_camb = cambpow.get_class_power(species=ccsp)
        (kk_ic, Pk_ic) = modecount_rebin(kk_ic, Pk_ic, modes_ic[ii], Pk_camb, ndesired=npart//2)
        error = plot_ic_power(kk_ic, Pk_ic, Pk_camb(kk_ic), sp=sp, npart=npart, outdir=outdir)
        #Don't worry too much about one failing mode.
        if np.size(np.where(error > accuracy)) > 3:
            raise RuntimeError("Pk accuracy check failed for "+str(sp)+". Max error: "+str(np.max(error)))

if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    parser.add_argument('genicfile', type=str, help='File with generated ICs')
    parser.add_argument('--czstr', type=str, help='Redshift string used in class files',required=True)
    parser.add_argument('--mnu', default=0, type=float,help='Sum of neutrino masses',required=False)
    args = parser.parse_args()
    check_ic_power_spectra(args.genicfile, camb_zstr = args.czstr, m_nu=args.mnu)
