from vasttools.pipeline import Pipeline
from vasttools.query import Query
from bokeh.io import output_notebook
from bokeh.plotting import show
from astropy.coordinates import Angle, SkyCoord
from astropy import units as u
import matplotlib.pyplot as plt
import pandas as pd
import os
import numpy as np

%matplotlib inline
output_notebook()

#from vasttools.utils import get_logger
#logger = get_logger(debug=True, quiet=False, logfile=None)

pipe = Pipeline()

pipe.list_piperuns()

run_name = 'workshop-2023-run'
run = pipe.load_run(run_name)

run.images

run.measurements

run.sources

run.sources.columns

run.sources[['avg_compactness', 'min_snr', 'max_snr','avg_flux_peak', 'max_flux_peak', 'min_flux_peak', 'v_peak', 'eta_peak', 'n_neighbour_dist', 'n_measurements', 'n_selavy', 'n_forced']]

run.sources_skycoord

single_det_query_str = (
    "n_measurements >= 10" # Require at least 10 measurements
    "& n_selavy == 1" # Only interested in sources that are detected once
    "& n_neighbour_dist > 1./60. " # nearest neighbour should be at least 1 arcmin away
    "& 0.8 < avg_compactness < 1.4 " # Only interested in compact sources
    "& n_relations == 0" # Only interested in sources that are not islands
    "& max_snr >= 10.0"
)

run.sources.query(single_det_query_str).sort_values('max_snr', ascending=False)

eta_v_query_str = (
    "n_measurements >= 10 " # require at least 10 measurements (selavy or forced)
    "& n_selavy > n_measurements / 2" # sources should be detected in at least half the observations
    "& n_neighbour_dist > 1./60. " # nearest neighbour should be at least 1 arcmin away
    "& 0.8 < avg_compactness < 1.4 "
    "& n_relations == 0"
    "& max_snr >= 10.0"
)

eta_thresh, v_thresh, eta_v_candidates, plot = run.run_eta_v_analysis(1.0, 1.0, query=eta_v_query_str)
print(eta_thresh, v_thresh)

eta_v_candidates

show(plot)

v_peak_cut = 0.35
eta_peak_cut = 50
strong_cands = eta_v_candidates.query(f"v_peak > {v_peak_cut} | eta_peak > {eta_peak_cut}").sort_values("v_peak", ascending=False)
strong_cands

fig = plt.figure()
ax = fig.add_subplot(111)
ax.scatter(run.sources.query(eta_v_query_str).eta_peak, run.sources.query(eta_v_query_str).v_peak, marker='.')
ax.scatter(strong_cands.eta_peak, strong_cands.v_peak, marker='.')
#ax.axhline(v_peak_cut)
#ax.axvline(eta_peak_cut)
ax.set_xscale('log')
ax.set_yscale('log')

cols_to_include = ['wavg_ra', 'wavg_dec', 'avg_flux_peak', 'max_flux_peak', 'min_flux_peak', 'v_peak', 'eta_peak', 'n_neighbour_dist', 'n_measurements', 'n_selavy', 'n_forced']

strong_cands[cols_to_include].to_csv('workshop-2023-cands.csv')

idx = 35430820
source = run.get_source(idx)

print(f"https://dev.pipeline.vast-survey.org/sources/{idx}/")

source.measurements

source.measurements.columns

lc = source.plot_lightcurve()

cutouts = source.show_all_png_cutouts(columns=5, figsize=(30,30))

cutouts = source.show_png_cutout(13)

source.simbad_search()

source.ned_search()

contour_plot = source.skyview_contour_plot(13, 'DSS')

idx = 34951897
stokes_i_source = run.get_source(idx)
stokes_i_lc = stokes_i_source.plot_lightcurve()

epochs = source.measurements.image.str.extract("epoch_([0-9]+)", expand=True)[0].values.astype(str)

query = Query(
    coords=stokes_i_source.coord,
    source_names=[stokes_i_source.name], # doesn't have to be provided
    matches_only=False, # return all sources regardless of if they are detected
    epochs=list(epochs), # Automatically query all observed VAST epochs, although not all are available on nimbus or in the form we require
    crossmatch_radius=10., # Search for matches within a 10" radius
    output_dir='example-output',
    forced_fits=False, # Do not run forced fitting - I recommend turning this off for initial explorations because it is slow due to poor I/O on nimbus
    use_tiles=True, # We only have TILES data for the full survey
    corrected_data=False, # Do not use corrected data - this is not yet available for the full survey
    post_processed_data=False,
    stokes='V' # Search Stokes V
)

query.find_sources()

query.results

stokes_v_source = query.results.iloc[0]

stokes_v_lc = stokes_v_source.plot_lightcurve()

stokes_i_lc

stokes_i_source.measurements['sbid'] = stokes_i_source.measurements.island_id.str.split('_').str[0].str[2:].astype(int)
full_stokes_measurements = stokes_i_source.measurements.merge(stokes_v_source.measurements, how='left', left_on='sbid', right_on='sbid', suffixes=("_I", "_V"))

full_stokes_measurements['polarisation_frac'] = np.abs(full_stokes_measurements.flux_peak_V/full_stokes_measurements.flux_peak_I)
full_stokes_measurements['polarisation_frac_err'] = full_stokes_measurements.polarisation_frac*np.sqrt((full_stokes_measurements.flux_peak_err_V/full_stokes_measurements.flux_peak_V)**2+(full_stokes_measurements.flux_peak_err_I/full_stokes_measurements.flux_peak_I)**2)

fig = plt.figure(figsize=(16,9))
ax = fig.add_subplot(111)
ax.errorbar(full_stokes_measurements.dateobs_I,
            full_stokes_measurements.polarisation_frac,
            yerr=full_stokes_measurements.polarisation_frac_err,
            marker='o',
            ls='')
ax.set_ylim(bottom=0)

from vasttools.moc import VASTMOCS
mocs = VASTMOCS()

psrcat_vast_sources = mocs.query_vizier_vast_pilot('B/psr', max_rows=200000) # Get all PSRCat sources that are within the VAST footprint

psrcat_vast_sources

psrcat_pd = psrcat_vast_sources.to_pandas()

psr_scs = SkyCoord(psrcat_pd._RAJ2000, psrcat_pd._DEJ2000, unit=u.deg)

psr_names = list(psrcat_pd.PSRJ)

idx, d2d, _ = psr_scs.match_to_catalog_sky(run.sources_skycoord)
radius_limit = 15 * u.arcsec
d2d_mask = d2d <= radius_limit

# How many pulsars are detected?
d2d_mask.sum() 

# Select the crossmatches less than 15 arcsec
psrcat_crossmatch_result = psrcat_pd.loc[d2d_mask].copy()

# Append the id and distance of the VAST crossmatch to the PSRCAT sources
psrcat_crossmatch_result['vast_xmatch_id'] = run.sources.iloc[idx[d2d_mask]].index.values
psrcat_crossmatch_result['vast_xmatch_d2d_asec'] = d2d[d2d_mask].arcsec

# Join the result
psrcat_crossmatch_result = psrcat_crossmatch_result.merge(run.sources, how='left', left_on='vast_xmatch_id', right_index=True, suffixes=("_psrcat", "_vast"))
psrcat_crossmatch_result

psrcat_query = Query(
    coords=psr_scs,
    source_names=psr_names, 
    matches_only=False,
    epochs="all-vast",
    crossmatch_radius=10., # Search for matches within a 10" radius
    output_dir='example-output',
    forced_fits=True,
    use_tiles=True,
    corrected_data=False,
    post_processed_data=False
)

#psrcat_query.find_sources()