#===========================================================#
# Run peppar PSF photometry/astrometry on Sgr B2 JWST NIRCam
# cal.fits images, one (filter, detector) per invocation.
#
# Usage:
#   python run_peppar_sgrb2.py <FILT_RUN> <DET_RUN>
# e.g.
#   python run_peppar_sgrb2.py F212N NRCA1
#   python run_peppar_sgrb2.py F480M NRCA5
#
# Data (cal.fits symlink farms, uncal excluded):
#   F212N -> /orange/adamginsburg/jwst/sgrb2/peppar/data_F212N
#   F480M -> /orange/adamginsburg/jwst/sgrb2/peppar/data_F480M
# Output root (stf_dir): /orange/adamginsburg/jwst/sgrb2/peppar
#   outputs land in stf_dir/<FILT>/<DET>/
#
# NB: flystar is NOT imported (not installed; only needed for some
# optional plotting/matching functions we do not call).
#===========================================================#
import os
import sys
import glob

from peppar import peppar

BASE = '/orange/adamginsburg/jwst/sgrb2/peppar'
DATA_DIRS = {
    'F212N': os.path.join(BASE, 'data_F212N'),
    'F480M': os.path.join(BASE, 'data_F480M'),
}


def define_variables(filt_run, det_run):
    """Set user variables for a single (filter, detector) run."""
    out_suffix = 'sgrb2'

    data_dir = DATA_DIRS[filt_run]
    stf_dir = BASE  # outputs written under stf_dir/<filt>/<det>/

    #---PSF Extraction parameters---#
    Niter = 2
    test_extraction = False
    bkg_method = '2d'
    threshold = 5
    peakmax = 5000
    group_dist_coeff = 2.
    shift_tol_coeff = 1.

    # Trimfake parameters
    threshold_trimfake = 2.
    fitshape_trimfake = 17
    subshape_trimfake = 31
    brite_percentile_trimfake = 20

    #---WebbPSF Parameters---#
    num_psf_model = 4
    oversample = 4
    add_detector_effects = True
    fov_model = 55
    recomp_psf_model = False

    #---Saturated star PSF-fit params---#
    mask_saturated = True
    fov_sat = fov_model
    fitshape_sat = 31
    threshold_sat = threshold
    update_sat_psf = True

    #---Background Parameters---#
    box_size = 100
    filter_size = 3

    #---PSF Fitting Parameters---#
    niter_epsf_construct = 3
    fov = 13
    fitshape = 9
    subshape = 13
    grid_points = 2
    min_npsf = 3
    aper_radius = 3
    use_combo_psf = False
    combo_r0 = fitshape / 2.0
    combo_k = 3

    #---Finding PSF stars---#
    peakmax_psf = peakmax
    threshold_psf = 100
    brite_percentile = 10
    min_dist = 3
    mag_diff = 3
    sharp_sig_lim = 3

    #---Validity checks (do not edit)---#
    assert fov_sat <= fov_model
    assert fov <= fov_model
    assert fitshape_trimfake <= fov_model
    assert subshape <= fov
    assert fitshape <= subshape

    if oversample == 1:
        detsample = True
    else:
        detsample = False

    return locals()


def run_peppar_script(filt_run, det_run):
    kwargs = define_variables(filt_run, det_run)

    stf_dir = kwargs['stf_dir']
    data_dir = kwargs['data_dir']

    if not os.path.exists(stf_dir):
        os.makedirs(stf_dir)
    print('Filter: {0}  Detector: {1}'.format(filt_run, det_run))
    print('Data dir: {0}'.format(data_dir))
    print('Reduction dir: {0}'.format(stf_dir))

    if det_run == 'i2d':
        images = sorted(glob.glob(data_dir + '/*nircam*_i2d.fits'))
        print(f'Found {len(images)} i2d images.')
    else:
        images = sorted(glob.glob(data_dir + '/*cal.fits'))
        print(f'Found {len(images)} cal images.')

    dict_images = peppar.organize_images_by_detector_filter(images)
    dict_utils = peppar.setup_filter_props()

    dict_image_props = peppar.setup_dict_image_props(filt_run, det_run)
    dict_images_run = peppar.setup_dict_images_for_run(dict_images, filt_run, det_run)

    print('=================')
    print('Run images dictionary keys: {0} / {1} -- {2} images'.format(
        filt_run, det_run, len(dict_images_run[filt_run][det_run]['images'])))
    print('=================')

    #---Setup PSFs---#
    fov_model = kwargs['fov_model']
    oversample = kwargs['oversample']
    num_psf_model = kwargs['num_psf_model']
    add_detector_effects = kwargs['add_detector_effects']
    detsample = kwargs['detsample']
    recomp_psf_model = kwargs['recomp_psf_model']

    psf_dir = '{0}/psf_default'.format(stf_dir)
    if not os.path.exists(psf_dir):
        os.makedirs(psf_dir)

    dict_psfs = peppar.get_psfs(dict_images_run, outdir=psf_dir, fov=fov_model,
                                oversample=oversample, num=num_psf_model,
                                add_detector_effects=add_detector_effects,
                                detsample=detsample, recomp_psf_model=recomp_psf_model)

    #---Run extract_catalogs---#
    make_plots = False
    peppar.extract_catalogs(
        dict_images_run, dict_psfs, dict_image_props, dict_utils,
        Niter=kwargs['Niter'],
        psf_model='webbpsf', psf_model_sat='webbpsf',
        peakmax=kwargs['peakmax'], bkg_method=kwargs['bkg_method'],
        threshold=kwargs['threshold'], threshold_sat=kwargs['threshold_sat'],
        update_sat_psf=kwargs['update_sat_psf'],
        aper_radius=kwargs['aper_radius'],
        fitshape=kwargs['fitshape'], subshape=kwargs['subshape'],
        fitshape_sat=kwargs['fitshape_sat'],
        fov=kwargs['fov'], fov_sat=kwargs['fov_sat'],
        niter_epsf_construct=kwargs['niter_epsf_construct'],
        grid_points=kwargs['grid_points'], threshold_psf=kwargs['threshold_psf'],
        min_npsf=kwargs['min_npsf'], brite_percentile=kwargs['brite_percentile'],
        mag_diff=kwargs['mag_diff'], min_dist=kwargs['min_dist'],
        sharp_sig_lim=kwargs['sharp_sig_lim'], oversample=kwargs['oversample'],
        group_dist_coeff=kwargs['group_dist_coeff'],
        shift_tol_coeff=kwargs['shift_tol_coeff'],
        threshold_trimfake=kwargs['threshold_trimfake'],
        fitshape_trimfake=kwargs['fitshape_trimfake'],
        subshape_trimfake=kwargs['subshape_trimfake'],
        brite_percentile_trimfake=kwargs['brite_percentile_trimfake'],
        box_size=kwargs['box_size'], filter_size=kwargs['filter_size'],
        use_combo_psf=kwargs['use_combo_psf'], combo_r0=kwargs['combo_r0'],
        combo_k=kwargs['combo_k'],
        out_suffix=kwargs['out_suffix'],
        out_dir=stf_dir, mask_saturated=kwargs['mask_saturated'],
        test_extraction=kwargs['test_extraction'], make_plots=make_plots)


if __name__ == '__main__':
    if len(sys.argv) != 3:
        print('Usage: python run_peppar_sgrb2.py <FILT_RUN> <DET_RUN>')
        sys.exit(1)
    filt_run = sys.argv[1]
    det_run = sys.argv[2]
    if filt_run not in DATA_DIRS:
        print('Unknown filter {0}; expected one of {1}'.format(filt_run, list(DATA_DIRS)))
        sys.exit(1)
    run_peppar_script(filt_run, det_run)