import numpy as np
from astropy.table import Table
from astropy import units as u
from astropy.coordinates import SkyCoord, search_around_sky

from mltier1 import generate_random_catalogue, Field, Q_0

%load_ext autoreload

%pylab inline

field = Field(170.0, 190.0, 45.5, 56.5)

"""
Explanation: PanSTARRS - WISE cross-match: Compute the $Q_0$
End of explanation
"""


panstarrs_full = Table.read("panstarrs_u2.fits")

wise_full = Table.read("wise_u2.fits")

"""
Explanation: Load the data
Load the catalogues
End of explanation
"""


panstarrs = field.filter_catalogue(
    panstarrs_full, 
    colnames=("raMean", "decMean"))

# Free memory
del panstarrs_full

wise = field.filter_catalogue(
    wise_full, 
    colnames=("raWise", "decWise"))

# Free memory
del wise_full

"""
Explanation: Restrict the study to the well sampled area
End of explanation
"""


coords_panstarrs = SkyCoord(panstarrs['raMean'], panstarrs['decMean'], unit=(u.deg, u.deg), frame='icrs')

coords_wise = SkyCoord(wise['raWise'], wise['decWise'], unit=(u.deg, u.deg), frame='icrs')

"""
Explanation: Coordinates
As we will use the coordinates to make a cross-match we to load them
End of explanation
"""


# Example function (not used, we use a class that contains this code)
def q_0_r(coords_wise, coords_panstarrs, field, radius=5):
    """Compute the Q_0 for a given radius"""
    random_wise = field.random_catalogue(len(coords_wise))
    
    idx_random_wise, idx_panstarrs, d2d, d3d = search_around_sky(
        random_wise, coords_panstarrs, radius*u.arcsec)
    
    nomatch_random = len(coords_wise) - len(np.unique(idx_random_wise))
    
    idx_wise, idx_panstarrs, d2d, d3d = search_around_sky(
        coords_wise, coords_panstarrs, radius*u.arcsec)
    
    nomatch_wise = len(coords_wise) - len(np.unique(idx_wise))
                                          
    return (1. - float(nomatch_wise)/float(nomatch_random))

q_0_comp = Q_0(coords_wise, coords_panstarrs, field)

q_0_comp(radius=5)

"""
Explanation: Cross-match and random cross-match
We create an instance of Q_0 using as the input the two catalogues and the area. It will compute the $Q_0$ for different radius.
The following function is not used but shows the internal code use for the computing of the $Q_0$.
End of explanation
"""


n_iter = 10

"""
Explanation: $Q_0$ dependence on the radius
We will iterate 10 times for each radius. However, the error is so small that a direct computation can be accurate to 4 significant figures.
End of explanation
"""


rads = list(range(1,26))

q_0_rad = []
for radius in rads:
    q_0_rad_aux = []
    for i in range(n_iter):
        out = q_0_comp(radius=radius)
        q_0_rad_aux.append(out)
    q_0_rad.append(np.mean(q_0_rad_aux))
    print("{:2d} {:7.5f} +/- {:7.5f} [{:7.5f} {:7.5f}]".format(radius, 
            np.mean(q_0_rad_aux), np.std(q_0_rad_aux), 
            np.min(q_0_rad_aux), np.max(q_0_rad_aux)))

plt.rcParams["figure.figsize"] = (5,5)
plot(rads, q_0_rad)
xlabel("Radius (arcsecs)")
ylabel("$Q_0$")
ylim([0, 1]);

"""
Explanation: The radius tested ranges from 1 to 25
End of explanation
"""


