import aquila_borg as borg
import numpy as np
from astropy.coordinates import SkyCoord
import astropy.units as apu
import configparser
import ast

import borg_velocity.utils as utils
from borg_velocity.utils import myprint, parse_file_to_dict
import borg_velocity.forwards as forwards
import borg_velocity.poisson_process as poisson_process
import borg_velocity.projection as projection

def radially_scatter(xtrue, frac_sig_x):
    """
    Radially perturb true positions by a fractional amount
    frac_sig_x, where the observer sits at x = (0, 0, 0).
    
    Args:
        :xtrue (np.ndarray): The true coordinates (Mpc/h), shape=(3,Nt)
        :frac_sig_x (float): The fractional uncertainty in the radial direction
        
    Returns:
        :xmeas (np.ndarray): The observed coordiantes (Mpc/h) of the tracers.
        :sigma_mu (float):  The uncertainty in the distance moduli of the tracers.
        
    """
    
    # Convert to RA, Dec, Distance
    rtrue = np.sqrt(np.sum(xtrue** 2, axis=0))  # Mpc/h
    c = SkyCoord(x=xtrue[0], y=xtrue[1], z=xtrue[2], representation_type='cartesian')
    RA = c.spherical.lon.degree
    Dec = c.spherical.lat.degree
    r_hat = np.array(SkyCoord(ra=RA*apu.deg, dec=Dec*apu.deg).cartesian.xyz)

    # Add noise to radial direction
    sigma_mu = 5. / np.log(10) * frac_sig_x
    mutrue = 5 * np.log10(rtrue * 1.e6 / 10)
    mumeas = mutrue + np.random.normal(size=len(mutrue)) * sigma_mu
    rmeas = 10 ** (mumeas / 5.) * 10 / 1.e6
    xmeas, ymeas, zmeas = rmeas[None,:] * r_hat

    xmeas = np.array([xmeas, ymeas, zmeas])
    
    return xmeas, sigma_mu


def borg_mock(s_hat, state, fwd_model, fwd_vel, ini_file, seed=None):
    
    myprint('Making mock from BORG')
    
    config = configparser.ConfigParser()
    config.read(ini_file)
    
    nsamp = int(config['run']['nsamp'])
    
    # Ensure cosmology is correct
    cosmo = utils.get_cosmopar(ini_file)
    fwd_model.setCosmoParams(cosmo)
    
    # Run BORG density field
    output_density = np.zeros(fwd_model.getOutputBoxModel().N)
    fwd_model.forwardModel_v2(s_hat)
    fwd_model.getDensityFinal(output_density)
    state["BORG_final_density"][:] = output_density
    
    # Get growth rate
    cosmo = utils.get_cosmopar(ini_file)
    output_vel = fwd_vel.getVelocityField()
    
    # Add bulk flow
    bulk_flow = np.array(ast.literal_eval(config['model']['bulk_flow']))
    output_vel = output_vel + bulk_flow.reshape((3, 1, 1, 1))
    
    if seed is not None:
        np.random.seed(seed)
    
    # Sample positions according to bias model
    bias_epsilon = float(config['model']['bias_epsilon'])
    R_max = float(config['mock']['R_max'])
    if config['model']['R_lim'] == 'none':
        R_lim = fwd_model.getOutputBoxModel().L[0]/2
    else:
        R_lim = float(config['model']['R_lim'])
    coord_true = [None] * nsamp
    coord_meas = [None] * nsamp
    sig_mu = [None] * nsamp
    for i in range(nsamp):
        
        myprint(f'Making mock for sample {i}')
        
        frac_sig_x = float(config[f'sample_{i}']['frac_sig_rhMpc'])
        alpha = float(config[f'sample_{i}']['alpha'])
        lam = float(config[f'sample_{i}']['lam'])
        Nt = int(config[f'sample_{i}']['Nt'])
        phi = (1. + output_density + bias_epsilon) ** alpha
        
        # Weight poisson mean by radial selection P(r) \propto r^2 exp(-lam r/R_max)
        # Do not include r^2 factor as we are sampling in Cartesian 3D,
        # so do not need the Jacobian term
        N = fwd_model.getOutputBoxModel().N
        L = fwd_model.getOutputBoxModel().L
        xmin = fwd_model.getOutputBoxModel().xmin
        x = np.linspace(xmin[0], xmin[0] + L[0], N[0])
        y = np.linspace(xmin[1], xmin[1] + L[1], N[1])
        z = np.linspace(xmin[2], xmin[2] + L[2], N[2])
        X, Y, Z = np.meshgrid(x, y, z, indexing='ij')
        r = np.sqrt(X**2 + Y**2 + Z**2)
        phi *= np.exp(- lam * r / R_max)

        # Sample coordinates
        
        # coord_true[i] = poisson_process.sample_3d(phi, Nt, 
        #       fwd_model.getOutputBoxModel().L[0], fwd_model.getOutputBoxModel().xmin)
        # coord_meas[i], sig_mu[i] = radially_scatter(coord_true[i], frac_sig_x)

        # Initialize lists to store valid positions and corresponding sig_mu values
        coord_meas_valid = np.empty((3, Nt))
        coord_true_valid = np.empty((3, Nt))

        # Counter for accepted positions
        accepted_count = 0

        # Loop until we have Nt valid positions
        while accepted_count < Nt:
            # Generate positions
            xtrue = poisson_process.sample_3d(phi, Nt, 
                                                   fwd_model.getOutputBoxModel().L[0], 
                                                   fwd_model.getOutputBoxModel().xmin)

            # Apply radial scattering
            xmeas, sig_mu[i] = radially_scatter(xtrue, frac_sig_x)

            # Compute radial distances of each measured coordinate
            radial_distances = np.sqrt(np.sum(xmeas ** 2, axis=0))
            
            # Filter coordinates that meet the R_lim condition
            valid_indices = radial_distances < R_lim
            valid_coords = xmeas[:,valid_indices]
            valid_coords_true = xtrue[:,valid_indices]

            # Calculate how many valid positions we need to reach Nt
            remaining_needed = Nt - accepted_count
            selected_count = min(valid_coords.shape[1], remaining_needed)

            # Append only the needed number of valid positions
            coord_meas_valid[:,accepted_count:accepted_count+selected_count] = valid_coords[:,:selected_count]
            coord_true_valid[:,accepted_count:accepted_count+selected_count] = valid_coords_true[:,:selected_count]
            
            # Update the accepted count
            accepted_count += selected_count
            
            myprint(f'\tMade {accepted_count} of {Nt}')
        
        coord_true[i] = coord_true_valid.copy()
        coord_meas[i] = coord_meas_valid.copy()

    # Interpolate velocities to tracers
    interp_order = int(config['model']['interp_order'])
    vr_true = [None] * nsamp
    for i in range(nsamp):
        tracer_vel = projection.interp_field(
                                output_vel,
                                np.expand_dims(coord_true[i], axis=2),
                                fwd_model.getOutputBoxModel().L[0],
                                np.array(fwd_model.getOutputBoxModel().xmin),
                                interp_order
                                )
        vr_true[i] = np.squeeze(projection.project_radial(
                                tracer_vel,
                                np.expand_dims(coord_true[i], axis=2),
                                np.zeros(3,)
                                ))
        
    # Compute observed redshifts (including noise)
    sig_v = float(config['model']['sig_v'])
    cz_obs = [None] * nsamp
    for i in range(nsamp):
        rtrue = np.sqrt(np.sum(coord_true[i] ** 2, axis=0))
        zco = utils.z_cos(rtrue, cosmo.omega_m)
        cz_obs[i] = utils.speed_of_light * zco + (1 + zco) * vr_true[i]
        cz_obs[i] += np.random.normal(size=cz_obs[i].shape) * sig_v  # CHECK THIS LINE!!!!!
        
    # Add observational systematic due to incorrect distance estimate 
    # \mu -> \mu + 5 log10(A), or equivalently d -> A d
    for i in range(nsamp):
        muA = float(config[f'sample_{i}']['muA'])
        coord_meas[i] = coord_meas[i] * muA
        
    return coord_true, coord_meas, sig_mu, vr_true, cz_obs


def velmass_mock(ini_file, seed=None):
    
    myprint('Making mock from VELMASS')
    
    if seed is not None:
        np.random.seed(seed)
    
    config = configparser.ConfigParser()
    config.read(ini_file)
    
    dirname = config['mock']['velmass_dirname']
    
    with open(dirname + "/halos_0.0.ascii", "rb") as f:
        for _ in range(30):
            r = str(f.readline())
            if "Units" in r:
                myprint(r)

    # Get box size
    with open(dirname + '/auto-rockstar.cfg') as f:
        data = [r for r in f]
    Lbox = [r for r in data if r.startswith('BOX_SIZE')][0].strip()
    Lbox = float(Lbox.split('=')[1])
    origin = np.array([Lbox/2, Lbox/2, Lbox/2])    
    omega_m = [r for r in data if r.startswith('Om')][0].strip()
    omega_m = float(omega_m.split('=')[1])
    
    halo_file = dirname + '/out.npy'
    halos = np.load(halo_file)

    # Cut by mass
    Mmin = float(config['mock']['velmass_Mmin'])
    halos = halos[halos['Mvir'] > Mmin]

    # Load positions
    xtrue = halos['X']  - origin[0] # Mpc / h
    ytrue = halos['Y']  - origin[1] # Mpc / h
    ztrue = halos['Z']  - origin[2] # Mpc / h
    rtrue = np.sqrt(xtrue**2 + ytrue**2 + ztrue**2) # Mpc / h
    xtrue = np.array([xtrue, ytrue, ztrue]) # shape = (3, nhalos)
    
    # Load velocities
    vtrue = np.array([halos['VX'], halos['VY'], halos['VZ']])
    
    # Get relative probability of each object according to selection function
    R_max = float(config['mock']['R_max'])
    
    # Define probability array. This is needed to prevent an object being in two samples
    prob = np.ones(len(rtrue))
    
    R_max = float(config['mock']['R_max'])
    if config['model']['R_lim'] == 'none':
        R_lim = fwd_model.getOutputBoxModel().L[0]/2
    else:
        R_lim = float(config['model']['R_lim'])
    
    # Sample 
    nsamp = int(config['run']['nsamp'])
    coord_true = [None] * nsamp
    coord_meas = [None] * nsamp
    vr_true = [None] * nsamp
    sig_mu = [None] * nsamp
    for i in range(nsamp):
        
        myprint(f'Making mock for sample {i}')
        lam = float(config[f'sample_{i}']['lam'])
        frac_sig_x = float(config[f'sample_{i}']['frac_sig_rhMpc'])
        Nt = int(config[f'sample_{i}']['Nt'])
        
        # Radial selection function
        m = prob != 0  # objects already found
        prob[m] = np.exp(- lam * rtrue[m] / R_max)
        prob /= np.sum(prob)
        
        # Initialize lists to store valid positions and corresponding sig_mu values
        coord_meas_valid = np.empty((3, Nt))
        coord_true_valid = np.empty((3, Nt))
        tracer_vel = np.empty((3, Nt))

        # Counter for accepted positions
        accepted_count = 0

        # Loop until we have Nt valid positions
        while accepted_count < Nt:

            ids = np.random.choice(len(rtrue), size=Nt, p=prob, replace=False)

            # Apply radial scattering
            xmeas, sig_mu[i] = radially_scatter(xtrue[:,ids], frac_sig_x)

            # Compute radial distances of each measured coordinate
            radial_distances = np.sqrt(np.sum(xmeas ** 2, axis=0))
            
            # Filter coordinates that meet the R_lim condition
            valid_indices = radial_distances < R_lim
            ids = ids[valid_indices]
            xmeas = xmeas[:,valid_indices]
            radial_distances = radial_distances[valid_indices]

            # Calculate how many valid positions we need to reach Nt
            remaining_needed = Nt - accepted_count
            selected_count = min(len(ids), remaining_needed)
            ids = ids[:selected_count]
            
            # Append only the needed number of valid halos
            coord_meas_valid[:,accepted_count:accepted_count+selected_count] = xmeas[:,:selected_count]
            coord_true_valid[:,accepted_count:accepted_count+selected_count] = xtrue[:,ids]
            tracer_vel[:,accepted_count:accepted_count+selected_count] = vtrue[:,ids]

            # Update the accepted count
            accepted_count += selected_count

            myprint(f'\tMade {accepted_count} of {Nt}')

            # Set up for next iteration
            prob[ids] = 0
            prob /= np.sum(prob)

        coord_true[i] = coord_true_valid.copy()
        coord_meas[i] = coord_meas_valid.copy()
        
        vr_true[i] = np.squeeze(projection.project_radial(
                                np.expand_dims(tracer_vel, axis=2),
                                np.expand_dims(coord_true[i], axis=2),
                                np.zeros(3,)
                                ))
        
    # Get the correct omega_m
    config_dict = parse_file_to_dict(f'{dirname}/rockstar.cfg')
    omega_m = config_dict['Om']
    
    # Compute observed redshifts (including noise - ACTUALLY NOT YET)
    sig_v = float(config['model']['sig_v'])
    cz_obs = [None] * nsamp
    for i in range(nsamp):
        rtrue = np.sqrt(np.sum(coord_true[i] ** 2, axis=0))
        zco = utils.z_cos(rtrue, omega_m)
        cz_obs[i] = utils.speed_of_light * zco + (1 + zco) * vr_true[i]
        # cz_obs[i] += np.random.normal(size=cz_obs[i].shape) * sig_v  # CHECK THIS LINE!!!!!
        
    # Add observational systematic due to incorrect distance estimate 
    # \mu -> \mu + 5 log10(A), or equivalently d -> A d
    for i in range(nsamp):
        muA = float(config[f'sample_{i}']['muA'])
        coord_meas[i] = coord_meas[i] * muA
        
    return coord_true, coord_meas, sig_mu, vr_true, cz_obs
 