#!/usr/bin/env python
"""Fit the morphology of a number of sources.

Uses initial parameters from a JSON file (for now only Gaussians).
"""

# Parse command line arguments

from gammapy.utils.scripts import argparse, GammapyFormatter
parser = argparse.ArgumentParser(description=__doc__,
                                 formatter_class=GammapyFormatter)
parser.add_argument('--counts', type=str, default='counts.fits',
                    help='Counts FITS file name')
parser.add_argument('--exposure', type=str, default='exposure.fits',
                    help='Exposure FITS file name')
parser.add_argument('--background', type=str, default='background.fits',
                    help='Background FITS file name')
parser.add_argument('--psf', type=str, default='psf.json',
                    help='PSF JSON file name')
parser.add_argument('--sources', type=str, default='sources.json',
                    help='Sources JSON file name (contains start '
                    'values for fit of Gaussians)')
parser.add_argument('--roi', type=str, default='roi.reg',
                    help='Region of interest (ROI) file name (ds9 reg format)')
parser.add_argument('fit_result', type=str, default=None,
                    help='Output JSON file with fit results')
args = parser.parse_args()
args = vars(args)

# Execute script

import logging
logging.basicConfig(level=logging.DEBUG, format='%(levelname)s - %(message)s')
import sherpa.astro.ui
from gammapy.morphology.utils import read_json, write_all
from gammapy.morphology.psf import Sherpa

# ---------------------------------------------------------
# Load images, PSF and sources
# ---------------------------------------------------------
logging.info('Clearing the sherpa session')
sherpa.astro.ui.clean()

logging.info('Reading counts: {0}'.format(args['counts']))
sherpa.astro.ui.load_data(args['counts'])

logging.info('Reading exposure: {0}'.format(args['exposure']))
sherpa.astro.ui.load_table_model('exposure', args['exposure'])

logging.info('Reading background: {0}'.format(args['background']))
sherpa.astro.ui.load_table_model('background', args['background'])

logging.info('Reading PSF: {0}'.format(args['psf']))
Sherpa(args['psf']).set()

if options.roi:
    logging.info('Reading ROI: {0}'.format(args['roi']))
    sherpa.astro.ui.notice2d(args['roi'])
else:
    logging.info('No ROI selected.')

logging.info('Reading sources: {0}'.format(args['sources']))
read_json(args['sources'], sherpa.astro.ui.set_source)

# ---------------------------------------------------------
# Set up the full model and freeze PSF, exposure, background
# ---------------------------------------------------------
# Scale exposure by 1e-10 to get ampl or order unity and avoid some fitting problems
sherpa.astro.ui.set_full_model('background + 1e-10 * exposure * psf (' + sherpa.astro.ui.get_source().name + ')')
sherpa.astro.ui.freeze(background, exposure, psf)

# ---------------------------------------------------------
# Set up the fit
# ---------------------------------------------------------
sherpa.astro.ui.set_coord('physical')
sherpa.astro.ui.set_stat('cash')
sherpa.astro.ui.set_method('levmar')  # levmar, neldermead, moncar
sherpa.astro.ui.set_method_opt('maxfev', int(1e3))
sherpa.astro.ui.set_method_opt('verbose', 10)

# ---------------------------------------------------------
# Fit and save information we care about
# ---------------------------------------------------------

# show_all() # Prints info about data and model
sherpa.astro.ui.fit()  # Does the fit
sherpa.astro.ui.covar()  # Computes symmetric errors (fast)
# conf() # Computes asymmetric errors (slow)
# image_fit() # Shows data, model, residuals in ds9
logging.info('Writing %s' % outfile)
write_all(outfile)

