plot_field.py

#!/usr/bin/env python
import os, sys, logging, io
import numpy as np
import pyvo as vo
import casacore.tables as pt
from astropy.table import Table, Column, vstack, unique, hstack
import argparse
#from lofarpipe.support.data_map import DataMap
#from lofarpipe.support.data_map import DataProduct
import requests
from astropy.coordinates import SkyCoord
from astropy import units as u
from astropy.io import fits
from astropy.wcs import WCS
from astropy.table import Table
from time import sleep


def sum_digits(n):
    s = 0
    c = 0
    for ii in n:
        if ii != '-':
            s += int(ii)
            c += 1
    norm = float(s)/float(c)
    return( norm )


def count_p(n):
    s = 0
    c = 0
    for ii in n:
        if ii != '-':
            c += 1
            if ii == 'P':
                s += 1
    norm = float(s)/float(c)
    return(norm)

def count_s(n):
    s = 0
    c = 0
    for ii in n:
        if ii != '-':
            c += 1
            if ii == 'S':
                s += 1
    norm = float(s)/float(c)
    return(norm)

def count_x(n):
    s = 0
    c = 0
    for ii in n:
        if ii != '-':
            c += 1
            if ii == 'X':
                s += 1
    norm = float(s)/float(c)
    return(norm)

def grab_coo_MS(MS):
    """
    Read the coordinates of a field from one MS corresponding to the selection given in the parameters
    Parameters
    ----------
    MS : str
        Full name (with path) to one MS of the field
    Returns
    -------
    RA, Dec : float,float
        coordinates of the field (RA, Dec in deg , J2000)
    """

    # reading the coordinates ("position") from the MS
    # NB: they are given in rad,rad (J2000) 
    [[[ra,dec]]] = pt.table(MS+'/FIELD', readonly=True, ack=False).getcol('PHASE_DIR')
    # [[[ra,dec]]] = pt.table(MS, readonly=True, ack=False).getcol('PHASE_DIR')

    # RA is stocked in the MS in [-pi;pi]
    # => shift for the negative angles before the conversion to deg (so that RA in [0;2pi])
    if ra<0:
        ra=ra+2*np.pi

    # convert radians to degrees
    ra_deg =  ra/np.pi*180.
    dec_deg = dec/np.pi*180.

    # and sending the coordinates in deg
    return ra_deg,dec_deg

def input2strlist_nomapfile(invar):
    """ from bin/download_IONEX.py
    give the list of MSs from the list provided as a string
    """

    str_list = None
    if type(invar) is str:
        if invar.startswith('[') and invar.endswith(']'):
            str_list = [f.strip(' \'\"') for f in invar.strip('[]').split(',')]
        else:
            str_list = [invar.strip(' \'\"')]
    elif type(invar) is list:
        str_list = [str(f).strip(' \'\"') for f in invar]
    else:
        raise TypeError('input2strlist: Type '+str(type(invar))+' unknown!')
    return str_list

def my_lotss_catalogue( RATar, DECTar,  Radius=1.5, bright_limit_Jy=5., faint_limit_Jy = 0.0, outfile='' ):

    """
    Download the LoTSS skymodel for the target field
    Parameters
    ----------
    ms_input : str
        String from the list (map) of the target MSs
    Radius : float (default = 1.5)
        Radius for the LOTSS cone search in degrees
    
    """
    ## first check if the file already exists
    if os.path.isfile( outfile ):
        print("LOTSS Skymodel for the target field exists on disk, reading in.")
        tb_final = Table.read( outfile, format='csv' )
        if 'Source_Name' in tb_final.colnames:
            tb_final.rename_column('Source_Name','Source_id')
            tb_final.write( outfile, format='csv' )
    else:
        print("DOWNLOADING LOTSS Skymodel for the target field")

        # Reading a MS to find the coordinate (pyrap)
        #RATar, DECTar = grab_coo_MS(input2strlist_nomapfile(ms_input)[0])

    ## this is the tier 1 database to query
        #url = 'http://vo.astron.nl/lofartier1/q/cone/scs.xml'
        # HETDEX database.
        #url = 'https://vo.astron.nl/hetdex/lotss-dr1/cone/scs.xml'
        url = 'https://vo.astron.nl/lotss_dr2/q/src_cone/scs.xml'

        ## query the database
        query = vo.dal.scs.SCSQuery( url )
        query['RA'] = float( RATar )
        query['DEC'] = float( DECTar )
        query.radius = float( Radius )
        t = query.execute()

        ## convert to VO table
        try:
            tb = t.votable.to_table()
        except AttributeError:
            # Above statement didn't work, try the alternative.
            tb = t.to_table()
        flux_sort = tb.argsort('Total_flux')
        tb_sorted = tb[flux_sort[::-1]]
        ## and keep only some of the columns
        colnames = tb_sorted.colnames
        ## first check for a resolved column
        if 'Resolved' not in colnames:
            resolved = np.where(is_resolved(tb_sorted['Total_flux'], tb_sorted['Peak_flux'], tb_sorted['Isl_rms']), 'R', 'U')
            tb_sorted['Resolved'] = resolved
        ## rename source_id column if necessary   
        if 'Source_Name' in tb_sorted.colnames:
            tb_sorted.rename_column('Source_Name', 'Source_id')
        keep_cols = ['Source_id', 'RA', 'DEC','Total_flux','Peak_flux', 'Majax', 'Minax', 'PA', 'DC_Maj', 'DC_Min', 'DC_PA', 'Isl_rms', 'Resolved']
        if 'LGZ_Size' in colnames:
            keep_cols = keep_cols + ['LGZ_Size', 'LGZ_Width', 'LGZ_PA']

        tb_final = tb_sorted[keep_cols]

        tb_final = tb_final[tb_final['Total_flux'] >= faint_limit_Jy*1e3]

        tb_final = tb_final[tb_final['Total_flux'] <= bright_limit_Jy*1e3]

        tb_final.write( outfile, format='csv' )

    return tb_final

def my_lbcs_catalogue( RATar, DECTar, Radius=1.5, outfile='' ):

    """
    Download the LBCS skymodel for the target field
    Parameters
    ----------
    ms_input : str
        String from the list (map) of the target MSs
    Radius : float (default = 1.5)
        Radius for the LOTSS cone search in degrees
    
    """
    ## first check if the file already exists
    #print( outfile )
    if os.path.isfile( outfile ):
        print("LBCS Skymodel for the target field exists on disk, reading in.")
        tb = Table.read( outfile, format='csv' )
        return(tb)
    else:
        print("DOWNLOADING LBCS Skymodel for the target field")

        # Reading a MS to find the coordinate (pyrap)
        #RATar, DECTar = grab_coo_MS(input2strlist_nomapfile(ms_input)[0])

    ## construct an html query and try to connect
    #print(RATar, DECTar)
    url = 'https://lofar-surveys.org/lbcs-search.fits?ra=%f&dec=%f&radius=%f' % (float(RATar), float(DECTar), float(Radius))
    connected=False
    while not connected:
            try:
                response = requests.get(url, stream=True,verify=True,timeout=60)
                if response.status_code!=200:
                    print (response.headers)
                    raise RuntimeError('Code was %i' % response.status_code)
            except requests.exceptions.ConnectionError:
                print ('Connection error! sleeping 30 seconds before retry...')
                sleep(30)
            except (requests.exceptions.Timeout,requests.exceptions.ReadTimeout):
                print( 'Timeout! sleeping 30 seconds before retry...')
                sleep(30)
            else:
                connected=True

    mem = io.BytesIO()
    for chunk in response.iter_content(chunk_size=8192):
        if chunk:
            mem.write(chunk)
    mem.seek(0)
    tb = Table.read(mem)
    mem.close()
    del(response)

    #print( len( tb ) )

    ### Workaround to sort and pick good calibrator info from tb array ###########
    
    if not len(tb) > 1:
        logging.critical('There are no LBCS sources within the given radius. Check your source is within the LBCS footprint and increase the search radius. Exiting...')
        return
    else:
            ## calculate the total FT goodness
        ft_total = []
        for xx in range(len(tb)):
                ft_total.append( sum_digits( tb[xx]['FT_Goodness'] ) )
        ft_col = Column( ft_total, name='FT_total' )
        tb.add_column( ft_col )

    tb.write( outfile, format='csv' )

    return tb

def find_close_objs(lo, lbcs, tolerance=5.):

    ## first filter the LBCS data on Flags
    lbcs_idx = np.where( np.logical_or(lbcs['Flags'] == 'O', lbcs['Flags'] == 'A' ) )
    lbcs = lbcs[lbcs_idx]
    ## get rid of anything with only X's
    nump = []
    nums = []
    for xx in range(len(lbcs)):
        nump.append(count_p(lbcs['Goodness'][xx]))
        nums.append(count_s(lbcs['Goodness'][xx]))
    #print( np.array(nump) + np.array(nums) )
    lbcs_idx = np.where( np.array(nump)+np.array(nums) > 0 )[0]
    lbcs = lbcs[lbcs_idx]

    ## get RA and DEC for the catalogues
    lotss_coords = SkyCoord( lo['RA'], lo['DEC'], frame='icrs', unit='deg' )
    lbcs_coords = SkyCoord( lbcs['RA'], lbcs['DEC'], frame='icrs', unit='deg' )

    ## search radius 
    search_rad = tolerance / 60. / 60. * u.deg

    ## loop through the lbcs coordinates -- this will be much faster than looping through lotss
    lotss_idx = []
    lbcs_idx = []    
    for xx in range(len(lbcs)):
        seps = lbcs_coords[xx].separation(lotss_coords)
        match_idx = np.where( seps < search_rad )[0]
        if len( match_idx ) == 0:
                # there's no match, move on to the next source
                m_idx = [-1]
                pass
        else:
            if len( match_idx ) == 1:
                    ## there's only one match
                    m_idx = match_idx[0]
                    lbcs_idx.append(xx)
                    lotss_idx.append(m_idx)
            if len( match_idx ) > 1:
                    ## there's more than one match, pick the brightest
                    tmp = lo[match_idx]
                    m_idx = np.where( tmp['Total_flux'] == np.max( tmp['Total_flux'] ) )[0]
                    if not isinstance(m_idx,int):
                        m_idx = m_idx[0]
                    lbcs_idx.append(xx)
                    lotss_idx.append(m_idx) 
    lbcs_matches = lbcs[lbcs_idx]
    lotss_matches = lo[lotss_idx]

    combined = hstack( [lbcs_matches, lotss_matches], join_type='exact' )
    ## check if there are duplicate lbcs observations for a lotss source
    if len( np.unique( combined['Source_id'] ) ) != len( combined ):
        # there are duplicates
        print( 'There are duplicate LBCS sources, selecting the best candidate(s).' )
        src_ids = np.unique( combined['Source_id'] )
        good_idx = []
        for src_id in src_ids:
            idx = np.where( combined['Source_id'] == src_id )[0]
            if len(idx) > 1:
            ## multiple matches found.  Count P's first and then break ties with Goodness_FT 
                num_P = []
                total_ft = []
                for yy in range( len( idx ) ):
                    tmp = combined[idx[yy]]['Goodness']
                    num_P.append( count_p( tmp ) )

                    tmp = combined[idx[yy]]['FT_Goodness']
                    total_ft.append( sum_digits( tmp ) )
                ## check that the total_ft values are non-zero before looking for a best cal
                if np.max( total_ft ) > 0:
                ## pick the one with the highest number of P's -- if tie, use total_ft
                    best_idx = np.where( num_P == np.max( num_P ) )[0]  ## this is an array
                if len( best_idx ) == 1:
                    good_idx.append(idx[best_idx][0])  ## idx[best_idx][0] is a number
                if len( best_idx ) > 1:
                    currentmax = 0.0 
                    for i in range(0,len(best_idx)):
                        if total_ft[best_idx[i]] > currentmax:
                            currentmax = total_ft[best_idx[i]]
                            ft_idx = i
                    good_idx.append( idx[best_idx[ft_idx]] )
                else:
                    print( 'Duplicate sources have total_ft = 0, removing from results.' )
            else:
                good_idx.append(idx[0])

        result = combined[good_idx]
    else:
        print( 'No duplicate sources found' )
        result = combined
    ## rename RA columns
    result.rename_column('RA_1','RA_LBCS')
    result.rename_column('DEC_1','DEC_LBCS')
    result.rename_column('RA_2','RA')
    result.rename_column('DEC_2','DEC')

    return result

def is_resolved(Sint, Speak, rms):
    """ Determines if a source is resolved or unresolved.
    The calculation is presented in Shimwell et al. 2018 of the LOFAR DR1 paper splash.
    
    Args:
        Sint (float or ndarray): integrated flux density.
        Speak (float or ndarray): peak flux density.
        rms (float or ndarray): local rms around the source.
    Returns:
        resolved (bool or ndarray): True if the source is resolved, False if not.
    """
    resolved = ((Sint / Speak) > 1.25 + 3.1 * (Speak / rms) ** (-0.53))
    return resolved

def remove_multiples_position( mycat, racol='RA', decol='DEC' ):
    radecstrings = []
    for i in np.arange(0,len(mycat)):
        radecstrings.append(str(mycat[i][racol]) + str(mycat[i][decol]) )
    radecstrings = np.asarray(radecstrings)
    if len( np.unique( radecstrings ) ) != len( mycat ):
        radecs = np.unique( radecstrings )
        good_idx = []
        for radec in radecs:
            idx = np.where( radecstrings == radec )[0]
            if len(idx) > 1:
                ## multiple matches found.  Count P's first and then break ties with Goodness_FT 
                num_P = []
                total_ft = []
                for yy in range( len( idx ) ):
                    tmp = mycat[idx[yy]]['Goodness']
                    num_P.append( count_p( tmp ) )
                    tmp = mycat[idx[yy]]['FT_Goodness']
                    total_ft.append( sum_digits( tmp ) )
                ## check that the total_ft values are non-zero before looking for a best cal
                if np.max( total_ft ) > 0:
                    ## pick the one with the highest number of P's -- if tie, use total_ft
                    best_idx = np.where( num_P == np.max( num_P ) )[0]  ## this is an array
                    if len( best_idx ) == 1:
                        good_idx.append(idx[best_idx][0])  ## idx[best_idx][0] is a number
                    if len( best_idx ) > 1:
                        currentmax = 0.0
                        for i in range(0,len(best_idx)):
                            if total_ft[best_idx[i]] > currentmax:
                                currentmax = total_ft[best_idx[i]]
                                ft_idx = i
                        good_idx.append( idx[best_idx[ft_idx]] )
                else:
                    print( 'Duplicate sources have total_ft = 0, removing from results.' )
            else:
                good_idx.append(idx[0])
        mycat = mycat[good_idx]
    else:
        print( 'All LBCS sources are unique' )

    return( mycat )

def mkfits (rasiz,decsiz,imsiz,pixsiz):
    hdu=fits.PrimaryHDU(np.zeros((int(imsiz),int(imsiz))))
    hdu.header.update({'CTYPE1':'RA---SIN'})
    hdu.header.update({'CRVAL1':rasiz})
    hdu.header.update({'CRPIX1':imsiz/2.})
    hdu.header.update({'CDELT1':-pixsiz})
    hdu.header.update({'CTYPE2':'DEC--SIN'})
    hdu.header.update({'CRVAL2':decsiz})
    hdu.header.update({'CRPIX2':imsiz/2.})
    hdu.header.update({'CDELT2':pixsiz})
    hdu.header.update({'EQUINOX':2000.0})
    hdu.header.update({'EPOCH':2000.0})
#    hdu.data = np.random.random(imsiz*imsiz).reshape(imsiz,imsiz)
    os.system('rm temp.fits')
    hdu.writeto('temp.fits')
    
def smearing_calculation(nchan = 16, obs_freq=144000000, radii=np.arange(0.001,4,0.00001), resolution=0.3/206265 * 180 / np.pi,
              av_time=1, thresholds = [0.2, 0.4, 0.6, 0.8] ):
    from scipy.special import erf

    bandwidth=1.95e3/nchan

    beta_fac = bandwidth/obs_freq * radii / resolution
    gamma_fac = 2. * np.sqrt( np.log(2.) ) 
    reduction_bandwidth = np.sqrt(np.pi)/(gamma_fac*beta_fac) * erf( gamma_fac*beta_fac/2)
    reduction_time = 1 - 1.22e-9 * ( radii / resolution )**2 * av_time**2
    reduction_total = reduction_bandwidth*reduction_time
    #print(reduction_bandwidth, reduction_total)
    thres_radii = []
    for thresh in thresholds:
        idx = np.argwhere(reduction_total < thresh)[0]
        #print(radii[idx])
        thres_radii.append(radii[idx])
    return thres_radii


def angular_distance(RA1, DEC1, RA2, DEC2):
    return np.sqrt((RA1-RA2)**2 + (DEC1-DEC2)**2)

def smallest_distance(RA, DEC, lbcs_catalogue):
    distances = [angular_distance(RA, DEC, source['RA'], source['DEC']) for source in lbcs_catalogue]

    ids = np.argmin(distances)

    return ids, distances[ids]
    
    
def make_plot(RA, DEC,  lotss_catalogue, extreme_catalogue, lbcs_catalogue, targRA=None, 
          targDEC=None, nchan = 16, av_time = 1):
    import matplotlib.pyplot as plt
    from matplotlib.patches import Circle

    mkfits(RA, DEC, 2048, 2.6367)

    f = fits.open("temp.fits")
    w = WCS(f[0].header)

    lbcs = lbcs_catalogue

    if os.path.exists(lotss_catalogue):
        lotss = Table.read(lotss_catalogue)
        avg_flux = np.median(lotss['Total_flux'])
    scaling = 0.1

    thres_radii = smearing_calculation(nchan = nchan, av_time=av_time)

    plt.figure(figsize = (10,10))

    ax = plt.subplot(projection = w,)

    color = 0.5
    fraction = ["80%" , "60%", "40%", "20%"]
    for i,thresh in enumerate(thres_radii):
        color = color + 0.1
        c = Circle((RA, DEC), thresh, edgecolor=None, facecolor=str(color),
           transform=ax.get_transform('fk5'), zorder = -1)
        ax.text(RA + 0.98*thresh , DEC, fraction[i], 
              transform=ax.get_transform('fk5'), fontsize = "large")
        ax.add_patch(c)

    ax.scatter(RA, DEC, marker= 'x', s = 30, transform=ax.get_transform('fk5'), label = "Ptg Centre")

    if os.path.exists(lotss_catalogue):
        ax.scatter(lotss['RA'], lotss['DEC'], transform=ax.get_transform('fk5'), 
            s = lotss['Total_flux'] * scaling, label = "Potential Targets")
        
    if len(extreme_catalogue) > 0:
        for source in extreme_catalogue:
            vector_orig = [source['RA'] - RA, source['DEC'] - DEC]
            norm = np.sqrt(vector_orig[0]**2 + vector_orig[1]**2)
            vector = [vector_orig[0]/norm, vector_orig[1]/norm]
            alt = [vector[0] -RA, vector[1] - DEC]
            print(vector, alt)
            ax.arrow(RA+2*vector[0], DEC+2*vector[1], 0.5*vector[0], 0.5*vector[1], transform=ax.get_transform('fk5'), 
                width = 0.01, head_width = 0.05, head_length = 0.05, length_includes_head = True) 
            ax.text(RA+2.5*vector[0], DEC+2.5*vector[1], 
                    s = "%.2f"%(source['Total_flux']/1000) + " Jy - %.2f degrees"%norm, 
                    transform=ax.get_transform('fk5'))

    ax.scatter(lbcs['RA'], lbcs['DEC'], transform=ax.get_transform('fk5'), s = 60, label = "LBCS Sources")
    

    c = Circle((RA, DEC), 1.5, edgecolor='yellow', facecolor='none',
           transform=ax.get_transform('fk5'))
    ax.add_patch(c)


    ### Calculate and plot smallest distance to target 
    dist_ids, dist = smallest_distance(RA, DEC, lbcs)

    closest_calib = lbcs[dist_ids]
    ax.plot([RA, closest_calib['RA']], [DEC, closest_calib['DEC']], 
             linestyle = '--', linewidth = 2, transform=ax.get_transform('fk5'), label = "Distance = %.2f degrees"%dist)    

    if targRA != None:
        ax.scatter(targRA, targDEC, marker= 'x', s = 80, 
         transform=ax.get_transform('fk5'), label = "Target")
        dist = angular_distance(targRA, targDEC, closest_calib['RA'], closest_calib['DEC'])

        ax.plot([targRA, closest_calib['RA']], [targDEC, closest_calib['DEC']], 
                 linestyle = '--', color = 'black',  linewidth = 2, transform=ax.get_transform('fk5'), label = "Distance = %.2f degrees"%dist)
    
    for i in range(len(lbcs)):
        ax.text(lbcs[i]['RA']-0.05, lbcs[i]['DEC'], "%.2f Jy"%(lbcs[i]['Total_flux']/1000),  transform=ax.get_transform('fk5'))
    plt.legend(fontsize = 'x-large')

    plt.savefig("output.png")

def convert_vlass_fits(fitsfile):
    import matplotlib.pyplot as plt
    from matplotlib.colors import LogNorm
    from astropy.visualization import PercentileInterval, imshow_norm
    from astropy.wcs import WCS
    header = fits.open(fitsfile)[0].header
    wcs = WCS(header)

    image_data = fits.getdata(fitsfile)

    # Shape is (1,1, 3722, 3722). Plot the first image
    interval = PercentileInterval(99.9)
    process_data = interval(image_data)
    plt.subplot(projection = wcs, slices=('x','y',0,0))
    imshow_norm(process_data, cmap='gray')

    
    plt.savefig(fitsfile[:-5] + ".png")


def plugin_main( RA, DEC, **kwargs ):
    RATar = RA
    DECTar = DEC
    targRA = kwargs['targRA']
    targDEC = kwargs['targDEC']
    #mapfile_in = kwargs['mapfile_in']
    lotss_radius = kwargs['lotss_radius']
    lbcs_radius  = kwargs['lbcs_radius']
    im_radius = float(kwargs['im_radius'])
    image_limit_Jy = float(kwargs['image_limit_Jy'])
    bright_limit_Jy = float(kwargs['bright_limit_Jy'])
    lotss_result_file = kwargs['lotss_result_file']
    lotss_catalogue = kwargs['lotss_catalogue']
    lbcs_catalogue = kwargs['lbcs_catalogue']
    delay_cals_file = kwargs['delay_cals_file']
    match_tolerance = float(kwargs['match_tolerance'])
    fail_lotss_ok = kwargs['continue_no_lotss'].lower().capitalize()
    nchan = kwargs['nchan']
    av_time = kwargs['av_time']
    vlass = kwargs['vlass']
    #mslist = DataMap.load(mapfile_in)
    #MSname = mslist[0].file
    # For testing
    #MSname = kwargs['MSname']
 
    ## first check for a valid delay_calibrator file
    if os.path.isfile(delay_cals_file):
        print( 'Delay calibrators file {:s} exists! returning.'.format(delay_cals_file) )
        return

    ## look for or download LBCS
    print("Attempting to find or download LBCS catalogue.")
    lbcs_catalogue = my_lbcs_catalogue( RATar, DECTar, Radius=lbcs_radius, outfile=lbcs_catalogue )
    ## look for or download LoTSS
    print("Attempting to find or download LoTSS catalogue.")
    lotss_catalogue = my_lotss_catalogue( RATar, DECTar,Radius=lotss_radius, bright_limit_Jy=bright_limit_Jy, faint_limit_Jy = 0.0,
                                         outfile=lotss_catalogue )
    
    print("Finding bright sources outside field")
    extreme_catalogue = my_lotss_catalogue( RATar, DECTar,Radius=10.0, bright_limit_Jy=1000., faint_limit_Jy = 10.0, outfile = "extreme_catalogue.csv" )
    extreme_catalogue = remove_multiples_position(extreme_catalogue)
    ## if lbcs exists, and either lotss exists or continue_without_lotss = True, continue
    ## else provide an error message and stop
    if len(lbcs_catalogue) == 0:
        logging.error('LBCS coverage does not exist, and catalogue not found on disk.')
        return
    if len(lotss_catalogue) == 0 and not fail_lotss_ok:
        logging.error('LoTSS coverage does not exist, and contine_without_lotss is set to False.')
        return 
    
    


        

    ## if the LoTSS catalogue is empty, write out the delay cals only
    if len(lotss_catalogue) == 0:
        print('Target field not in LoTSS coverage yet! Only writing {:s} based on LBCS'.format(delay_cals_file))

        ## Add the radius from phase centre to the catalogue
        #RATar, DECTar = grab_coo_MS(input2strlist_nomapfile(MSname)[0])
        ptg_coords = SkyCoord( RATar, DECTar, frame='icrs', unit='deg' )

        src_coords = SkyCoord( lbcs_catalogue['RA'], lbcs_catalogue['DEC'], frame='icrs', unit='deg' )
        separations = src_coords.separation(ptg_coords )
        seps = Column( separations.deg, name='Radius' )
        lbcs_catalogue.add_column( seps )

        ## rename the source_id column
        #lbcs_catalogue.rename_column('Observation','Source_id')

        ## add in some dummy data
        Total_flux = Column( np.ones(len(lbcs_catalogue))*1e3, name='Total_flux', unit='mJy' )
        lbcs_catalogue.add_column( Total_flux )
        LGZ_Size = Column( np.ones( len(lbcs_catalogue) )*20., name='LGZ_Size', unit='arcsec' ) ## set to a default of 20 arcsec
        lbcs_catalogue.add_column( LGZ_Size )

        ## remove duplicate sources if necessary 
        lbcs_catalogue = remove_multiples_position( lbcs_catalogue )

        ## order based on radius from the phase centre
        lbcs_catalogue.sort('Radius')

        ## write the catalogue
        lbcs_catalogue.write(delay_cals_file, format='csv')

        result = lbcs_catalogue
    else:
        ## else continue 
        result = find_close_objs( lotss_catalogue, lbcs_catalogue, tolerance=match_tolerance )
        ## check if there are any matches
        if len(result) == 0:
            logging.error('LoTSS and LBCS coverage exists, but no matches found. This indicates something went wrong, please check your catalogues.')
        else:
        # add radius to the catalogue
            #RATar, DECTar = grab_coo_MS(input2strlist_nomapfile(MSname)[0])
            #result = lbcs_catalogue
            ptg_coords = SkyCoord( RATar, DECTar, frame='icrs', unit='deg' )

            src_coords = SkyCoord( result['RA'], result['DEC'], frame='icrs', unit='deg' )
            separations = src_coords.separation(ptg_coords )
            seps = Column( separations.deg, name='Radius' )
            result.add_column( seps )

            ## order by radius from the phase centre
            result.sort( 'Radius' )

            #result.rename_column('Observation','Source_id')

            ## Write catalogues
            ## 1 - delay calibrators -- from lbcs_catalogue
            result.write( delay_cals_file, format='csv' )
            print('Writing delay calibrator candidate file {:s}'.format(delay_cals_file))

            ## sources to image -- first remove things that are already in the delay_cals_file
            good_index = [ x for x, src_id in enumerate( lotss_catalogue['Source_id'] ) if src_id not in result['Source_id'] ]
        
            tmp_cat = lotss_catalogue[good_index]

            ## make a flux cut
            image_index = np.where( tmp_cat['Peak_flux'] >= image_limit_Jy*1e3 )[0]
            flux_cut_sources = tmp_cat[image_index]

            ## make a radius cut
            src_coords = SkyCoord( flux_cut_sources['RA'], flux_cut_sources['DEC'], frame='icrs', unit='deg' )
            separations = src_coords.separation( ptg_coords )
            seps = Column( separations.deg, name='Radius' )
            flux_cut_sources.add_column( seps )
            good_idx = np.where( flux_cut_sources['Radius'] <= im_radius )[0]
            sources_to_image = flux_cut_sources[good_idx]

            nsrcs = float( len( sources_to_image ) )
            print( "There are "+str(len(lbcs_catalogue))+" delay calibrators within " + str(im_radius) + " degrees of the pointing centre")
            print( "There are "+str(nsrcs)+" sources above "+str(image_limit_Jy*1000)+" mJy within "+str(im_radius)+" degrees of the phase centre.")
            sources_to_image.write( lotss_result_file, format='csv' )

    print("Assumed averaging - nchannels: %s; time averaging: %s"%(nchan, av_time))
    make_plot(RATar, DECTar,  lotss_result_file, extreme_catalogue, result, targRA, targDEC,nchan = nchan, av_time = av_time)

    
    
    if vlass:
        from vlass_search import search_vlass
        ## Get cutouts of all LBCS sources
        print("Getting cutouts of LBCS sources")
        for i, source in enumerate(result):
            ra, dec = source['RA'], source['DEC']
            c = SkyCoord(ra, dec, unit = (u.deg, u.deg))
            outfile = "%s_vlass.fits"%source['Observation']
            try:
                search_vlass(c, crop = True, crop_scale = 256)
                os.system("mv vlass_post**.fits  %s"%outfile)
                convert_vlass_fits(outfile)
            except:
                pass
    return




if __name__ == "__main__":

    parser = argparse.ArgumentParser()
    parser.add_argument( '--lotss_radius', dest='lotss_radius', type=float, help='Radius to search LoTSS', default=1.5 )
    parser.add_argument( '--lbcs_radius', dest='lbcs_radius', type=float, help='Radius to search LBCS', default=1.5 )
    parser.add_argument( '--im_radius', dest='im_radius', type=float, help='Radius in which to image', default=1.24 )
    parser.add_argument( '--lotss_catalogue', dest='lotss_catalogue', type=str, help='input file for LoTSS catalogue [will be downloaded if does not exist]', default='lotss_catalogue.csv' )
    parser.add_argument( '--lbcs_catalogue', dest='lbcs_catalogue', type=str, help='input file for LBCS catalogue [will be downloaded if does not exist]', default='lbcs_catalogue.csv' )
    parser.add_argument( '--lotss_result_file', dest='lotss_result_file', type=str, help='output file of sources to image', default='image_catalogue.csv' )
    parser.add_argument( '--delay_cals_file', dest='delay_cals_file', type=str, help='output file of delay calibrators', default='delay_calibrators.csv' )
    parser.add_argument( '--match_tolerance', dest='match_tolerance', type=float, help='radius for matching LBCS to LoTSS [arcsec]', default=5. )
    parser.add_argument( '--bright_limit_Jy', dest='bright_limit_Jy', type=float, help='Flux limit for bright sources [Jy]', default=5.0 )
    parser.add_argument( '--image_limit_Jy', dest='image_limit_Jy', type=float, help='Flux limit for which sources to image [Jy]', default = 0.01 )
    parser.add_argument( '--continue_no_lotss', dest='continue_no_lotss', action='store_true', help='Continue with the pipeline if no lotss cross-matches can be found?', default = False )
    parser.add_argument( '--targRA', type=float, dest='targRA', help='Target RA in deg' )
    parser.add_argument( '--targDEC', type=float, dest='targDEC', help='Target DEC in deg' )
    parser.add_argument( '--nchan', dest='nchan', type=int, help='Number of frequency channels', default=16 )
    parser.add_argument( '--av_time', dest='av_time', type=float, help='Time averaging', default = 1. )

    parser.add_argument( '--MS', type=str, dest='MS', help='Measurement Set' )
    parser.add_argument( '--vlass', dest='vlass', action='store_true', help='Get VLASS cutouts of delay_calibraors.csv sources', default = False)

    parser.add_argument( '--RA', type=float, dest = 'RA', help='Ptg RA in deg' )
    parser.add_argument( '--DEC', type=float, dest = 'DEC', help='Ptg DEC in deg' )


    args = parser.parse_args()

    if args.MS is not None:
        print("Using MS to get RA and DEC") 
        ptgRA, ptgDEC = grab_coo_MS(args.MS)
    else:
        ptgRA, ptgDEC = args.RA, args.DEC
        


    plugin_main( float(ptgRA), float(ptgDEC), targRA = args.targRA, targDEC = args.targDEC, 
         lotss_radius=args.lotss_radius, lbcs_radius=args.lbcs_radius, im_radius=args.im_radius,
           bright_limit_Jy=args.bright_limit_Jy, lotss_catalogue=args.lotss_catalogue, 
           lbcs_catalogue=args.lbcs_catalogue, lotss_result_file=args.lotss_result_file, 
           delay_cals_file=args.delay_cals_file, match_tolerance=args.match_tolerance, 
           image_limit_Jy=args.image_limit_Jy, continue_no_lotss = str(args.continue_no_lotss),
            nchan = args.nchan,  av_time = args.av_time, vlass=args.vlass)






### TO DO LIST

# Change size of targets based on flux - Add argument for setting minimum flux - DONE
# Look for bright sources within 10 degrees above a few Jy - use the bright_limit_Jy argument - DONE
# Add colour for quality of LBCS sources - Use fit or the compactness codes - Maybe not 
# Number the calibrators based on brightness - DONE


# Label shading of the smearing - DONE
# Print some statistics about the sources - Number etc. - DONE
# How many sources above the flux limit - DONE
# Target to phase center distance - DONE
# Make all the grey's lighter - DONE
# Print what averaging was assumed - Also add as an optional argument - Both channels and time - DONE