#!/usr/bin/env python

import argparse
import yaml

from astropy.coordinates import SkyCoord
from astropy import units as u
import galsim
from romanisim import log, wcs, persistence, parameters
from romanisim import ris_make_utils as ris
from copy import copy


def go(args):
    if args.config is not None:
        # Open and parse overrides file
        with open(args.config, "r") as config_file:
            config = yaml.safe_load(config_file)
            combo_dict = parameters.__dict__
            ris.merge_nested_dicts(combo_dict, config)
            parameters.__dict__.update(combo_dict)
    elif args.usecrds:
        # don't use default values
        for k in parameters.reference_data:
            parameters.reference_data[k] = None

    if args.sca == -1:
        # simulate all 18 SCAs sequentially
        origfilename = copy(args.filename)
        origprevious = copy(args.previous)
        for i in range(1, 19):
            args.sca = i
            args.filename = origfilename.format(f'WFI{args.sca:02d}')
            if origprevious is not None:
                args.previous = origprevious.format(f'WFI{args.sca:02d}')
            go(args)
        return

    metadata = ris.set_metadata(
        date=args.date, bandpass=args.bandpass,
        sca=args.sca, ma_table_number=args.ma_table_number,
        truncate=args.truncate)

    if args.radec is not None:
        coord = SkyCoord(ra=args.radec[0] * u.deg, dec=args.radec[1] * u.deg,
                         frame='icrs')
        wcs.fill_in_parameters(metadata, coord, boresight=args.boresight, pa_aper=args.roll)

    rng = galsim.UniformDeviate(args.rng_seed)

    # Create catalog
    cat = ris.create_catalog(metadata=metadata, catalog_name=args.catalog,
                             bandpasses=[args.bandpass], rng=rng,
                             nobj=args.nobj, usecrds=args.usecrds)

    # Create persistence object
    if args.previous is not None:
        persist = persistence.Persistence.read(args.previous)
    else:
        persist = persistence.Persistence()

    # Simulate image and write to file
    ris.simulate_image_file(args, metadata, cat, rng, persist)


if __name__ == '__main__':
    parser = argparse.ArgumentParser(
        description='Make a demo image.',
        epilog='EXAMPLE: %(prog)s output_image.asdf',
        formatter_class=argparse.ArgumentDefaultsHelpFormatter)
    parser.add_argument('filename', type=str, help='output image (asdf)')
    parser.add_argument('--bandpass', type=str, help='bandpass to simulate',
                        default='F087')
    parser.add_argument('--boresight', action='store_true', default=False,
                        help=('radec specifies location of boresight, not '
                              'center of WFI.'))
    parser.add_argument('--catalog', type=str, help='input catalog (ecsv)',
                        default=None)
    parser.add_argument('--config', type=str, help='input parameter override file (yaml)',
                        default=None)
    parser.add_argument('--date', type=str, default=None,
                        help=('UTC Date and Time of observation to simulate in ISOT format.'))
    parser.add_argument('--level', type=int, default=2,
                        help='1 or 2, for L1 or L2 output')
    parser.add_argument('--ma_table_number', type=int, default=1)
    parser.add_argument('--nobj', type=int, default=1000)
    parser.add_argument('--previous', default=None, type=str,
                        help=('previous simulated file in chronological order '
                              'used for persistence modeling.'))
    parser.add_argument('--radec', type=float, nargs=2,
                        help='ra and dec (deg)', default=None)
    parser.add_argument('--rng_seed', type=int, default=None)
    parser.add_argument('--roll', type=float, default=0,
                        help='Position angle (North towards YIdl) measured at the V2Ref/V3Ref of the aperture used.')
    parser.add_argument('--sca', type=int, default=7, help='SCA to simulate')
    parser.add_argument('--usecrds', action='store_true',
                        help='Use CRDS for distortion map')
    parser.add_argument('--webbpsf', action='store_true',
                        help='Use webbpsf for PSF')
    parser.add_argument('--truncate', type=int, default=None, help=(
        'If set, truncate the MA table at given number of resultants.'))
    parser.add_argument('--pretend-spectral', type=str, default=None, help=(
        'Pretend the image is spectral.  exposure.type and instrument.element '
        'are updated to be grism / prism.'))
    parser.add_argument('--drop-extra-dq', default=False, action='store_true',
                        help=('Do not store the optional simulated dq array.'))

    args = parser.parse_args()

    log.info('Starting simulation...')
    log.warning("romanisim is under active development.  Its output has "
                "not been formally validated; only limited testing has been "
                "performed.  For this reason, use of romanisim for "
                "preparation of ROSES proposals is not advised.  Other "
                "packages like galsim's roman package or STIPS may better "
                "serve such purposes.")

    go(args)
