#!/usr/bin/env python
"""
This script runs an FPP calculation for a provided
transit signal.

All parameters are defined in the ``.ini`` files,
which are by default assumed to be
in the current path (or in the folder defined as the
argument to this script).

The ``fpp.ini`` file should be of the following form::

    name = k2oi #anything
    ra = 11:30:14.510 #can be decimal form too
    dec = +07:35:18.21

    period = 32.988 #days
    rprs = 0.0534   #Rp/Rstar
    photfile = lc_k2oi.csv #contains transit photometry

    #provide Teff, feh, [logg optional] if spectrum available
    #Teff = 3503, 80  #value, uncertainty
    #feh = 0.09, 0.09
    #logg = 4.89, 0.1

and ``star.ini`` should look something like this::

    #observed magnitudes of target star
    # If uncertainty provided, will be used to fit StarModel
    [mags]
    J = 9.763, 0.03
    H = 9.135, 0.03
    K = 8.899, 0.02
    Kepler = 12.473

Running this script will create the following files,
in the same directory as the ``.ini`` file:

  * ``trsig.pkl``: the pickled :class:`vespa.TransitSignal` object.
  * ``starfield.h5``: the TRILEGAL field star simulation
  * ``starmodel.h5``: the :class:`isochrones.StarModel` fit
  * ``starmodel_corner_physical.png,
     starmodel_corner_observed.png``:
     :class:`isochrones.StarModel` corner plots.
  * ``popset.h5``: the :class:`vespa.PopulationSet` object
    representing the model population simulations.
  * ``eb.png, heb.png, beb.png, pl.png``: likelihood plots
    for each model.
  * ``trsig.png``: plot of the transit signal
  * ``FPPsummary.png``: FPP summary plot.
  * ``calcfpp.log``: logfile of calculation.

This script is not yet set up to incorporate observational constraints,
but this will be easily doable in the near future.

"""
from __future__ import print_function, division

import matplotlib
matplotlib.use('Agg')

import sys, os, re, time, os.path, glob
import argparse
import logging

from vespa import FPPCalculation
from vespa.populations import BoxyModel, LongModel
from isochrones import StarModel, BinaryStarModel, TripleStarModel

import warnings
warnings.simplefilter("error")
warnings.simplefilter("ignore", DeprecationWarning)
warnings.simplefilter("ignore", RuntimeWarning)
warnings.simplefilter("ignore", UserWarning)

import logging

#utility function to initialize logging
def initLogging(filename, logger):
    if logger == None:
        logger = logging.getLogger()
    else:  # wish there was a logger.close()
        for handler in logger.handlers[:]:  # make a copy of the list
            logger.removeHandler(handler)

    logger.setLevel(logging.INFO)
    formatter = logging.Formatter(fmt='%(asctime)s: %(message)s')

    fh = logging.FileHandler(filename)
    fh.setFormatter(formatter)
    logger.addHandler(fh)

    sh = logging.StreamHandler(sys.stdout)
    logger.addHandler(sh)
    return logger

if __name__=='__main__':
    parser = argparse.ArgumentParser(description='Calculate FPP for a transit signal')

    parser.add_argument('folders', nargs='*', default=['.'],
                        help='Directory (or directories) for which to calculate FPP.  ' +
                        '`starfit --all <directory>` must have been run already for each directory.')
    parser.add_argument('-i','--inifile', type=str, default='fpp.ini',
                        help='Name of ini file (within directory) containing configuration (default `fpp.ini`).')
    parser.add_argument('-n','--n', type=int, default=20000,
                        help='Size of simulated populations (default=20000).')
    parser.add_argument('--recalc', action='store_true',
                        help='Delete all existing population simulations and recompute.')
    parser.add_argument('--debug', action='store_true',
                        help='Set logging level to DEBUG.')
    parser.add_argument('--newlog', action='store_true',
                        help='Start a fresh fpp.log file.')
    parser.add_argument('--recalc_lhood', action='store_true',
                        help='Recaclulate likelihoods instead of reading from cache file.')
    parser.add_argument('--refit_trap', action='store_true',
                        help='Refit trapezoidal models of simulated populations ' +
                             '(keeping existing population simulations).')
    parser.add_argument('--refit_trsig', action='store_true',
                        help='Redo MCMC fit trapezoidal model to transit signal.')
    parser.add_argument('--bootstrap', type=int, default=0,
                        help='Number of bootstrap resamplings to do to estimate uncertainty of FPP. ' +
                             'Set to >1 to do bootstrap error estimates.')
    parser.add_argument('--include_artificial', action='store_true',
                        help='Whether to include artificial "boxy" and "long" models in computing FPP.')
    parser.add_argument('--artificial_prior', type=float, default=10./200000,
                        help='Prior to use for artificial models. Default is 5e-5.')
    parser.add_argument('--ichrone', default='mist', type=str,
                        help='Stellar models to use.  Use "mist" (default) or "dartmouth".')


    args = parser.parse_args()

    logger = None #dummy

    for folder in args.folders:
        #initialize logger for this folder:

        logfile = os.path.join(folder, 'calcfpp.log')
        if args.newlog:
            os.remove(logfile)
        logger = initLogging(logfile, logger)
        if args.debug:
            logger.setLevel(logging.DEBUG)

        try:

            f = FPPCalculation.from_ini(folder, ini_file=args.inifile,
                                        recalc=args.recalc,
                                        refit_trap=args.refit_trap,
                                        ichrone=args.ichrone,
                                        n=args.n)

            if args.refit_trsig:
                f.trsig.MCMC(refit=True)
                f.trsig.save(os.path.join(folder,'trsig.pkl'))

            trap_corner_file = os.path.join(folder, 'trap_corner.png')
            if not os.path.exists(trap_corner_file) or args.refit_trsig:
                f.trsig.corner(outfile=trap_corner_file)

            if args.include_artificial:
                boxmodel = BoxyModel(args.artificial_prior, f['pl'].stars.slope.max())
                longmodel = LongModel(args.artificial_prior, f['pl'].stars.duration.quantile(0.99))
                f.add_population(boxmodel)
                f.add_population(longmodel)

            f.FPPplots(recalc_lhood=args.recalc_lhood)
            if args.bootstrap > 0:
                logger.info('Re-fitting trapezoid MCMC model...')
                f.bootstrap_FPP(args.bootstrap)

            for mult,Model in zip(['single','binary','triple'],
                                  [StarModel, BinaryStarModel, TripleStarModel]):
                starmodel_file = os.path.join(folder, '{}_starmodel_{}.h5'.format(args.ichrone, mult))
                corner_file1 = os.path.join(folder,
                                              '{}_corner_{}_physical.png'.format(args.ichrone, mult))
                corner_file2 = os.path.join(folder,
                                              '{}_corner_{}_observed.png'.format(args.ichrone, mult))
                if not os.path.exists(corner_file1) or not os.path.exists(corner_file2):
                    logger.info('Making StarModel corner plots...')
                    starmodel = Model.load_hdf(starmodel_file)
                    corner_base = os.path.join(folder,
                                                 '{}_corner_{}'.format(args.ichrone, mult))
                    starmodel.corner_plots(corner_base)

            if args.bootstrap > 0:
                logger.info('Bootstrap results ({}) written to {}.'.format(args.bootstrap,
                                                                          os.path.join(os.path.abspath(folder),
                                                                                   'results_bootstrap.txt')))
            logger.info('FPP calculation successful. ' +
                        'Results/plots written to {}.'.format(os.path.abspath(folder)))

            print('FPP for {}: {}'.format(f.name,f.FPP()))
        except KeyboardInterrupt:
            raise
        except:
            logger.error('FPP calculation failed for {}.'.format(folder), exc_info=True)
