import numpy as np, os
from glob import glob
from scipy.interpolate import LinearNDInterpolator,interp1d
from scipy import integrate
import matplotlib.pyplot as plt
import matplotlib.cm as cm
from astropy.cosmology import WMAP9 as cosmo

f0 = 1420405751.7667  # Hz

def build_model_interp(parm_array,delta2_array,k_array,redshift,
    regrid_ks=None):
    #input an array of models, the parameters for each model, list of k modes
    #   and the desired redshift
    #return a list of ks and a list of matching interpolation functions,
    # function has call f(log10(Nx),alphaX,log10(MminX))
    #parm_array expected to be nmodels,nparms
    #with columns (z,Nf,Nx,alphaX,MminX,other-stuff....)
    #delta2_array expected to be nmodels,nkmodes
    #NOTE: assumes all models are computed at the same k modes
    closest_redshift = parm_array[np.abs(parm_array[:,0]-redshift).argmin(),0]
    model_points = parm_array[parm_array[:,0]==closest_redshift,2:5]

    #interpolate NX and Mmin in log space
    model_points[:,0] = np.log10(model_points[:,0])
    model_points[:,2] = np.log10(model_points[:,2])
    #get the power spectrum values that go with this redshift
    raw_model_values = delta2_array[parm_array[:,0]==closest_redshift]
    model_values = [[]]*len(raw_model_values)
    if not regrid_ks is None: #regrid to a different set of k bins
        for i in xrange(len(raw_model_values)):
            model_values[i] = interp1d(k_array,raw_model_values[i,:])(regrid_ks)
        model_values = np.array(model_values)
        k_array = regrid_ks
        print "interpolated sim shape after regridding",model_values.shape
    else:
        model_values = raw_model_values.copy()
    #for a single redshift, build an interplation for each k mode
    Pk_models_atz = []
    for ki,k in enumerate(k_array):

        M = LinearNDInterpolator(model_points,model_values[:,ki])
        Pk_models_atz.append(M)
    return Pk_models_atz
def build_tau_interp_model(parm_array):
    #interpolate NX and Mmin in log space
    alphaXs = np.sort(list(set(parm_array[:,3])))
    Mmins = np.sort(list(set(parm_array[:,4])))
    Nxs = np.sort(list(set(parm_array[:,2])))
    taus = []
    for Nx in Nxs:
        for alphaX in alphaXs:
            for Mmin in Mmins:
                _slice = np.argwhere(all_and([
                                    parm_array[:,2]==Nx,
                                    parm_array[:,3]==alphaX,
                                    parm_array[:,4]==Mmin]
                                    ))
                taus.append([np.log10(Nx),alphaX,np.log10(Mmin),
                            nf_to_tau(parm_array[_slice,0].squeeze(),
                            parm_array[_slice,1].squeeze())])
    taus = np.array(taus)
    return LinearNDInterpolator(taus[:,:3],taus[:,3])

def all_and(arrays):
    #input a list or arrays
    #output the arrays anded together
    if len(arrays)==1:return arrays
    out = arrays[0]
    for arr in arrays[1:]:
        out = np.logical_and(out,arr)
    return out

def load_andre_models(fileglob, sortz=True):
    #input a string that globs to the list of input model files
    #return arrays of parameters,k modes, delta2,and delt2 error
    #parm_array expected to be nmodels,nparms
    #with columns (z,Nf,Nx,alphaX,Mmin,other-stuff....)
    #delta2_array expected to be nmodels,nkmodes
    filenames = glob(fileglob)
    parm_array = []
    k_array = []
    delta2_array = []
    delta2_err_array = []