import numpy as np
from coordinates_bgs import Coordinates


def cut_sky(position, velocity, cosmology, Lbox=2000, zmin=None, zmax=None, replication=(0,0,0)):
    """
    Creates a cut sky mock by converting the cartesian coordiantes of a cubic box mock to ra, dec, z
    
    Args:
        position:  array of comoving position vectors (Mpc/h), in the range -Lbox/2 < pos < Lbox/2
        velocity:  array of proper velocity vectors (km/s)
        cosmology: instance of astropy.cosmology class
        [Lbox]:    comoving box length of simulation (Mpc/h), default is 2000 Mpc/h
        [zmin]:    If provided, will only return galaxies with z>=zmin.
                   By default will return all galaxies.
        [zmax]:    If provided, will only return galaxies with z<=zmax. 
                   By default will return all galaxies.
        [replication]: tuple indicating which periodic replication to use. 
                   Default value is (0,0,0) (no replications).
    """
    
    # array of the index of each galaxy in the original cube
    index = np.arange(position.shape[0])
    
    coords = Coordinates(cosmology)
    
    position_rep = position.copy()
    if replication==(0,0,0):
        print(f"{replication} No periodic replications")
    else:
        print(f"{replication} Applying periodic replications")
        for i in range(3):
            print(f"{replication} %.1f < %s < %.1f"%((-1+2*replication[i])*Lbox/2., chr(120+i), (1+2*replication[i])*Lbox/2.))
            position_rep[:,i] += Lbox*replication[i]
    print(np.shape(position_rep))

    ra, dec, zcos = coords.pos3d_to_equitorial(position_rep)
    print(min(ra), max(ra), min(dec), max(dec), min(zcos), max(zcos))
    #print(f'{replication} {len(ra)=}')
    vlos = coords.vel_to_vlos(position_rep, velocity)
    zobs = coords.zcos_to_zobs(zcos, vlos)
    #print(f'{replication} {min(zobs)=}, {max(zobs)=}')
    print(zmin, zmax)
    print(min(zobs), max(zobs))
    keep = np.ones(len(zobs), dtype=bool)
    if not zmin is None:
        print(f"{replication} Applying redshift cut z > {zmin:.3f}")
        keep &= zobs >= zmin

    if not zmax is None:
        print(f"{replication} Applying redshift cut z < {zmax:.3f}")
        keep &= zobs <= zmax

    pkeep = 100*np.sum(keep)/len(keep)

    print(f'{replication} Keeping {np.sum(keep)}/{len(keep)} ({pkeep:.2f}%) galaxies')
    ra, dec, zcos, zobs, index = ra[keep], dec[keep], zcos[keep], zobs[keep], index[keep]
    
    print(f"Redshift min/max de todas las galaxias antes del corte: {zobs.min()} - {zobs.max()}")
    print(f"Redshift min/max permitidos: {zmin} - {zmax}")


    print(f"Total de galaxias antes del corte: {len(zobs)}")
    print(f"Total de galaxias después del corte: {np.sum(keep)}")

    return ra, dec, zcos, zobs, index