import numpy as np
from astropy.table import Table, join
from astropy import units as u
from astropy.coordinates import SkyCoord, search_around_sky
from IPython.display import clear_output
import pickle
import os

from mltier1 import (get_center, get_n_m, estimate_q_m, Field, MultiMLEstimatorU,
                     parallel_process, get_sigma_all, get_q_m)

%load_ext autoreload

%autoreload

from IPython.display import clear_output

%pylab inline

"""
Explanation: ML match for LOFAR and the combined PanSTARRS WISE catalogue: Get all sources above the threshold
In this notebook the maximum likelihood cross-match between the LOFAR HETDEX catalogue and the combined PansSTARRS WISE catalogue is computed. However, we take all the sources that have likelihood ratios above the threshold found in the first run.
Configuration
Load libraries and setup
End of explanation
"""


save_intermediate = True
plot_intermediate = True

idp = "idata/final"

if not os.path.isdir(idp):
    os.makedirs(idp)

"""
Explanation: General configuration
End of explanation
"""


# Busy week Edinburgh 2017
ra_down = 172.09
ra_up = 187.5833
dec_down = 46.106
dec_up = 56.1611

# Busy week Hatfield 2017
ra_down = 170.
ra_up = 190.
dec_down = 46.8
dec_up = 55.9

# Full field July 2017
ra_down = 160.
ra_up = 232.
dec_down = 42.
dec_up = 62.

field = Field(170.0, 190.0, 46.8, 55.9)

field_full = Field(160.0, 232.0, 42.0, 62.0)

"""
Explanation: Area limits
End of explanation
"""


combined_all = Table.read("pw.fits")

lofar_all = Table.read("data/LOFAR_HBA_T1_DR1_catalog_v0.9.srl.fits")

np.array(combined_all.colnames)

np.array(lofar_all.colnames)

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


lofar = field_full.filter_catalogue(lofar_all, colnames=("RA", "DEC"))

combined = field_full.filter_catalogue(combined_all, 
                               colnames=("ra", "dec"))

"""
Explanation: Filter catalogues
End of explanation
"""


combined["colour"] = combined["i"] - combined["W1mag"]

combined_aux_index = np.arange(len(combined))

"""
Explanation: Additional data
End of explanation
"""


coords_combined = SkyCoord(combined['ra'], 
                           combined['dec'], 
                           unit=(u.deg, u.deg), 
                           frame='icrs')

coords_lofar = SkyCoord(lofar['RA'], 
                       lofar['DEC'], 
                       unit=(u.deg, u.deg), 
                       frame='icrs')

"""
Explanation: Sky coordinates
End of explanation
"""


combined_matched = (~np.isnan(combined["i"]) & ~np.isnan(combined["W1mag"])) # Matched i-W1 sources

combined_panstarrs = (~np.isnan(combined["i"]) & np.isnan(combined["W1mag"])) # Sources with only i-band

combined_wise =(np.isnan(combined["i"]) & ~np.isnan(combined["W1mag"])) # Sources with only W1-band

combined_i = combined_matched | combined_panstarrs
combined_w1 = combined_matched | combined_wise
#combined_only_i = combined_panstarrs & ~combined_matched
#combined_only_w1 = combined_wise & ~combined_matched

print("Total    - ", len(combined))
print("i and W1 - ", np.sum(combined_matched))
print("Only i   - ", np.sum(combined_panstarrs))
print("With i   - ", np.sum(combined_i))
print("Only W1  - ", np.sum(combined_wise))
print("With W1  - ", np.sum(combined_w1))

"""
Explanation: Class of sources in the combined catalogue
The sources are grouped depending on the available photometric data.
End of explanation
"""


colour_limits = [-0.5, 0.0, 0.5, 1.0, 1.25, 1.5, 1.75, 2.0, 2.25, 2.5, 2.75, 3.0, 3.5, 4.0]

# Start with the W1-only, i-only and "less than lower colour" bins
colour_bin_def = [{"name":"only W1", "condition": combined_wise},
                  {"name":"only i", "condition": combined_panstarrs},
                  {"name":"-inf to {}".format(colour_limits[0]), 
                   "condition": (combined["colour"] < colour_limits[0])}]

# Get the colour bins
for i in range(len(colour_limits)-1):
    name = "{} to {}".format(colour_limits[i], colour_limits[i+1])
    condition = ((combined["colour"] >= colour_limits[i]) & 
                 (combined["colour"] < colour_limits[i+1]))
    colour_bin_def.append({"name":name, "condition":condition})

# Add the "more than higher colour" bin
colour_bin_def.append({"name":"{} to inf".format(colour_limits[-1]), 
                       "condition": (combined["colour"] >= colour_limits[-1])})

combined["category"] = np.nan
for i in range(len(colour_bin_def)):
    combined["category"][colour_bin_def[i]["condition"]] = i

np.sum(np.isnan(combined["category"]))

"""
Explanation: Colour categories
The colour categories will be used after the first ML match
End of explanation
"""


numbers_combined_bins = np.array([np.sum(a["condition"]) for a in colour_bin_def])

numbers_combined_bins

"""
Explanation: We get the number of sources of the combined catalogue in each colour category. It will be used at a later stage to compute the $Q_0$ values
End of explanation
"""


bin_list, centers, Q_0_colour, n_m, q_m = pickle.load(open("lofar_params.pckl", "rb"))

likelihood_ratio_function = MultiMLEstimator(Q_0_colour, n_m, q_m, centers)

"""
Explanation: Maximum Likelihood
End of explanation
"""


radius = 15

lr_threshold = 0.36

selection = ~np.isnan(combined["category"]) # Avoid the dreaded sources with no actual data
catalogue = combined[selection]

def apply_ml(i, likelihood_ratio_function):
    idx_0 = idx_i[idx_lofar == i]
    d2d_0 = d2d[idx_lofar == i]
    
    category = catalogue["category"][idx_0].astype(int)
    mag = catalogue["i"][idx_0]
    mag[category == 0] = catalogue["W1mag"][idx_0][category == 0]
    
    lofar_ra = lofar[i]["RA"]
    lofar_dec = lofar[i]["DEC"]
    lofar_pa = lofar[i]["PA"]
    lofar_maj_err = lofar[i]["E_Maj"]
    lofar_min_err = lofar[i]["E_Min"]
    c_ra = catalogue["ra"][idx_0]
    c_dec = catalogue["dec"][idx_0]
    c_ra_err = catalogue["raErr"][idx_0]
    c_dec_err = catalogue["decErr"][idx_0]
    
    sigma, sigma_maj, sigma_min = get_sigma_all(lofar_maj_err, lofar_min_err, lofar_pa, 
                      lofar_ra, lofar_dec, 
                      c_ra, c_dec, c_ra_err, c_dec_err)

    lr_0 = likelihood_ratio_function(mag, d2d_0.arcsec, sigma, sigma_maj, sigma_min, category)
    
    chosen_index = (lr_0 >= lr_threshold)
    result = [combined_aux_index[selection][idx_0[chosen_index]], # Index
              (d2d_0.arcsec)[chosen_index],                        # distance
              lr_0[chosen_index]]                                  # LR
    return result

"""
Explanation: ML match
End of explanation
"""


idx_lofar, idx_i, d2d, d3d = search_around_sky(
    coords_lofar, coords_combined[selection], radius*u.arcsec)

idx_lofar_unique = np.unique(idx_lofar)

total_sources = len(idx_lofar_unique)

"""
Explanation: Run the cross-match
End of explanation
"""


def ml(i):
    return apply_ml(i, likelihood_ratio_function)

res = parallel_process(idx_lofar_unique, ml, n_jobs=8)

"""
Explanation: Run the ML matching
End of explanation
"""


lofar_aux_index = np.arange(len(lofar))

lr = []
lr_dist = []
lr_index = []
lr_order = []
lr_lofar_index = []

for i, idx in enumerate(idx_lofar_unique):
    result = res[i]
    n = len(result[0])
    lofar_index = lofar_aux_index[idx]
    if n > 0:
        order = np.argsort(result[2])[::-1]
        lr.extend(result[2][order])
        lr_dist.extend(result[1][order])
        lr_index.extend(result[0][order])
        lr_order.extend(np.arange(n, dtype=int) + 1)
        lr_lofar_index.extend(np.ones(n, dtype=int)*lofar_index)
    else:
        lr.append(np.nan)
        lr_dist.append(np.nan)
        lr_index.append(np.nan)
        lr_order.append(np.nan)
        lr_lofar_index.append(lofar_index)

aux_table = Table()
aux_table['aux_index'] = lr_lofar_index
aux_table['lr'] = lr
aux_table['lr_dist'] = lr_dist
aux_table['lr_index'] = lr_index
aux_table['lr_order'] = lr_order

aux_table

lofar["aux_index"] = lofar_aux_index

lofar_lr = join(lofar, aux_table, join_type='outer', keys='aux_index')

lofar_lr

"""
Explanation: Selection and match
End of explanation
"""


combined["lr_index"] = combined_aux_index.astype(float)

for col in ['lr', 'lr_dist', 'lr_index', 'lr_order']:
    lofar_lr[col].fill_value = np.nan

pwl = join(lofar_lr.filled(), combined, join_type='left', keys='lr_index')

len(pwl)

pwl_columns = pwl.colnames

for col in pwl_columns:
    fv = pwl[col].fill_value
    #print(col, fv)
    if (isinstance(fv, np.float64) and (fv != 1e+20)):
        print(col, fv)
        pwl[col].fill_value = 1e+20

columns_save = ['Source_Name', 'RA', 'E_RA', 'E_RA_tot', 'DEC', 'E_DEC', 'E_DEC_tot',
 'Peak_flux', 'E_Peak_flux', 'E_Peak_flux_tot', 'Total_flux', 'E_Total_flux', 'E_Total_flux_tot',
 'Maj', 'E_Maj', 'Min', 'E_Min', 'PA', 'E_PA', 'Isl_rms', 'S_Code', 'Mosaic_ID', 'Isl_id',
 'AllWISE', 'objID', 'ra', 'dec', 'raErr', 'decErr',
 'W1mag', 'W1magErr', 'i', 'iErr', 'colour', 'category',
 'lr', 'lr_dist', 'lr_order']

pwl[columns_save].filled().write('lofar_multi_lr_pw.fits', format="fits")

"""
Explanation: Save combined catalogue
End of explanation
"""


