import aquila_borg as borg
import configparser
import numpy as np
import jax.numpy as jnp
from scipy.special import hyp2f1
import symbolic_pofk.linear

cons = borg.console()
myprint = lambda x: cons.print_std(x) if type(x) == str else cons.print_std(repr(x))

def get_cosmopar(ini_file):
    """
    Extract cosmological parameters from an ini file
    
    Args:
        :ini_file (str): Path to the ini file
    
    Returns:
        :cpar (borg.cosmo.CosmologicalParameters): Cosmological parameters
    """
    
    config = configparser.ConfigParser()
    config.read(ini_file)
        
    cpar = borg.cosmo.CosmologicalParameters()
    cpar.default()
    cpar.fnl = float(config['cosmology']['fnl'])
    cpar.omega_k = float(config['cosmology']['omega_k'])
    cpar.omega_m = float(config['cosmology']['omega_m'])
    cpar.omega_b = float(config['cosmology']['omega_b'])
    cpar.omega_q = float(config['cosmology']['omega_q'])
    cpar.h = float(config['cosmology']['h100'])
    cpar.sigma8 = float(config['cosmology']['sigma8'])
    cpar.n_s = float(config['cosmology']['n_s'])
    cpar.w = float(config['cosmology']['w'])
    cpar.wprime = float(config['cosmology']['wprime'])
    cpar.A_s = 1.e-9 * symbolic_pofk.linear.sigma8_to_As(
            cpar.sigma8, cpar.omega_m, cpar.omega_b, cpar.h, cpar.n_s)
    
    return cpar

def initial_pos(L,N,order="F"):
    values = np.linspace(0,L,N+1)[:-1] #ensure LLC
    xx,yy,zz = np.meshgrid(values,values,values)

    if order=="F":
        pos_mesh = np.vstack((yy.flatten(),xx.flatten(),zz.flatten())).T
    if order=="C":
        pos_mesh = np.vstack((zz.flatten(),xx.flatten(),yy.flatten())).T
        
    return pos_mesh


def correct_displacement_over_periodic_boundaries(disp,L,max_disp_1d=125):
    # Need to correct for positions moving over the periodic boundary
    
    moved_over_bound = L - max_disp_1d
    axis = ['x','y','z']
    
    for i in [0,1,2]:
        idx_sup, idx_sub = check(disp,L,moved_over_bound,max_disp_1d,i,axis)

        # Correct positions
        disp[:,i][idx_sup] -= L
        disp[:,i][idx_sub] += L

        check(disp,L,moved_over_bound,max_disp_1d,i,axis)

        assert np.amin(disp[:,i]) >= -max_disp_1d and np.amax(disp[:,i]) <= max_disp_1d, "Particles outside allowed region"
        
    return disp


def check(disp,L,moved_over_bound,max_disp_1d,i,axis):
    idxsup = disp[:,i]>moved_over_bound
    idx = np.abs(disp[:,i])<=max_disp_1d
    idxsub = disp[:,i]<-moved_over_bound

    sup = len(disp[:,i][idxsup])
    did_not_cross_boundary = len(disp[:,i][idx])
    sub = len(disp[:,i][idxsub])

    if not sub+did_not_cross_boundary+sup == len(disp[:,i]):
        myprint(f'Disp in {axis[i]} direction under -{moved_over_bound} Mpc/h is = '+str(sub))
        myprint(f'|Disp| in {axis[i]} direction under {max_disp_1d} Mpc/h is = '+str(did_not_cross_boundary))
        myprint(f'Disp in {axis[i]} direction over {moved_over_bound} Mpc/h is = '+str(sup))
        myprint('These add up to: '+str(sub+did_not_cross_boundary+sup))
        myprint(f"Should add up to: len(disp[:,i]) {len(disp[:,i])}")
        myprint('\n')

        assert sub+did_not_cross_boundary+sup == len(disp[:,i]), "Incorrect summation" # cannot lose/gain particles

    return idxsup, idxsub


def dis(x, undo=False, z=0.0, dis_std=6.0, **kwargs):
    dis_norm = dis_std * linear_D(z)  # [Mpc/h]

    if not undo:
        dis_norm = 1 / dis_norm

    x *= dis_norm
    
    
def linear_D(z, Om=0.31):
    """linear growth function for flat LambdaCDM, normalized to 1 at redshift zero
    """
    OL = 1 - Om
    a = 1 / (1+z)
    return a * hyp2f1(1, 1/3, 11/6, - OL * a**3 / Om) \
             / hyp2f1(1, 1/3, 11/6, - OL / Om)


def get_cell_coord(x):
    ix = jnp.floor(x)
    return x - ix, ix

def get_cell_coord_np(x):
    ix = np.floor(x)
    return x - ix, ix