import os
import sys
import shutil
import numpy as np
import json

import urllib.request


# use polynomial coefficients to calculate the flux at the reference frequency
def vlacal_calc_flux_from_coeffs( coeff, freq ):
        x = np.log10( freq / 1e9 )
        p1 = np.poly1d( coeff[::-1] )
        return 10**(p1(x))



# interpolate fluxscale data
def vlacal_interpolate_fluxscale( fluxscale_coeffs, freq, target_date, verbose=False ):

    dates = np.array( [ fluxscale_coeffs[key]['date_observed'] for key in fluxscale_coeffs.keys() ], dtype='int')
    coeffs = [ fluxscale_coeffs[key]['coefficients'] for key in fluxscale_coeffs.keys() ]
    to_sort = dates.argsort()

    fluxes = np.array( [ vlacal_calc_flux_from_coeffs( coef, freq ) for coef in coeffs ] )

    if verbose:
        print(f'Available fluxscale data at reference frequency {freq/1e9} GHz:')
        for i,date in enumerate(dates[to_sort]):
            print('   {0}  {1}'.format(date, fluxes[to_sort][i]) )

    new_flux = np.interp( target_date, dates[to_sort], fluxes[to_sort] )

    return new_flux



def vlacal_setjy_calculate_ratio( source, band, date ):
    
    vlacal_url = "http://vlacal.pythonanywhere.com"

    band_table_txt = '''\
    P, 0.23, 0.47
    L, 1, 2
    S, 2, 4
    C, 4, 8
    X, 8, 12
    U, 12, 18
    K, 18, 26.5
    A, 26.5, 40
    Q, 40, 50
    '''

    band_table = np.genfromtxt( band_table_txt.split('\n'), dtype=None, delimiter=',', encoding='utf-8' )
    f_min, f_max = band_table[band_table['f0']==band][['f1','f2']][0]
    f_min = f_min*1e9
    f_max = f_max*1e9
    f_reference = 0.5*(f_min+f_max)


    url = f'{vlacal_url}/fluxscale/{source}'

    req = urllib.request.Request( url )
    with urllib.request.urlopen(req) as response:
        fluxscale_coeffs = json.loads(response.read().decode('utf-8'))

    target_date = int(qa.splitdate(date)['mjd'])
    new_flux = vlacal_interpolate_fluxscale( fluxscale_coeffs, f_reference, target_date, verbose=False )

    print('Obtained flux of {0:.3f} Jy on mjd {1} for source {2} at reference frequency {3} GHz'.format(new_flux, target_date, source, f_reference/1e9))

    pb17_coeffs = [fluxscale_coeffs[key]['coefficients'] for key in fluxscale_coeffs.keys() if fluxscale_coeffs[key]['date_observed'] == 57754][0]
    pb17_flux = vlacal_calc_flux_from_coeffs( pb17_coeffs, f_reference )

    excess_flux = new_flux - pb17_flux
    print('This interpolated flux differs from Perley-Butler 2017 by {0:0.2f}%'.format( 100.0*excess_flux/pb17_flux ))





def vlacal_setjy_predict_model_data(vis, field,  source, spw, band, date):

    vlacal_url = "http://vlacal.pythonanywhere.com"


    band_table_txt = '''\
    P, 0.23, 0.47
    L, 1, 2
    S, 2, 4
    C, 4, 8
    X, 8, 12
    U, 12, 18
    K, 18, 26.5
    A, 26.5, 40
    Q, 40, 50
    '''


    # get band center frequency from band table
    band_table = np.genfromtxt( band_table_txt.split('\n'), dtype=None, delimiter=',', encoding='utf-8' )
    f_min, f_max = band_table[band_table['f0']==band][['f1','f2']][0]
    f_min = f_min*1e9
    f_max = f_max*1e9
    f_reference = 0.5*(f_min+f_max)


    # retrieve fluxscale
    url = f'{vlacal_url}/fluxscale/{source}'

    req = urllib.request.Request( url )
    with urllib.request.urlopen(req) as response:
        fluxscale_coeffs = json.loads(response.read().decode('utf-8'))


    # interpolate flux
    target_date = int(qa.splitdate( date )['mjd'])
    new_flux = vlacal_interpolate_fluxscale( fluxscale_coeffs, f_reference, target_date, verbose=False )

    print(f'Interpolated {source} flux on mjd {target_date} of {new_flux:0.3f} Jy at reference frequency {f_reference/1e9} GHz')

    pb17_coeffs = [fluxscale_coeffs[key]['coefficients'] for key in fluxscale_coeffs.keys() if fluxscale_coeffs[key]['date_observed'] == 57754][0]
    pb17_flux = vlacal_calc_flux_from_coeffs( pb17_coeffs, f_reference )

    excess_flux = new_flux - pb17_flux
    pb17_low = vlacal_calc_flux_from_coeffs( pb17_coeffs, f_min )
    pb17_high = vlacal_calc_flux_from_coeffs( pb17_coeffs, f_max )

    flux_low = vlacal_interpolate_fluxscale( fluxscale_coeffs, f_min, target_date)
    flux_high = vlacal_interpolate_fluxscale( fluxscale_coeffs, f_max, target_date)

    excess_low = flux_low - pb17_low
    excess_high = flux_high - pb17_high

    if np.isfinite( np.log(excess_high/excess_low) ):
        core_index = np.log(excess_high/excess_low) / np.log(f_max/f_min)
    else:
        core_index = 0.0


    # retrieve_complist for requested source, band
    url = f'{vlacal_url}/components/{source}/{band}'

    req = urllib.request.Request( url )
    with urllib.request.urlopen(req) as response:
        components = json.loads(response.read().decode('utf-8'))

    direction_rad = components['csys']['direction0']['crval']
    core_direction = 'J2000 {0} {1}'.format( qa.time('{0}rad'.format(direction_rad[0]), prec=12)[0], qa.angle('{0}rad'.format(direction_rad[1]),prec=12 )[0] )

    cl_name = f'{source}_{band}_vlacal_mjd{target_date}.cl' 


    # modify cl table
    if os.path.isdir( cl_name ):
        rmtables( cl_name ) 

    cl.close(log=False)
    cl.fromrecord(components['complist'])
    cl.addcomponent( flux=excess_flux, fluxunit='Jy', dir=core_direction, shape='point', freq='{0}Hz'.format(f_reference), spectrumtype='spectral index', index=core_index  )
    clrec = cl.torecord()
    cl.rename( cl_name )
    cl.done()

#    print(f'Saving modified component list table: {cl_name}')
    print(f"Filling MODEL_DATA column using task_ft for data selection: field='{field}', spw='{spw}'")


    # run ft
    ft(vis=vis, field=field, spw=spw, complist=cl_name, usescratch=True)