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

All parameters are defined in the ``.ini`` file,
which is by default assumed to be named ``fpp.ini``
in the current path, but can in principle be anything.

The ``.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

    #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_triangle_physical.png,
     starmodel_triangle_observed.png``:
     :class:`isochrones.StarModel` triangle 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.
  * ``run.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 isochrones import StarModel

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

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=['.'])
    parser.add_argument('-i','--inifile', type=str, default='fpp.ini')
    parser.add_argument('-n','--n', type=int, default=20000)
    parser.add_argument('--recalc', action='store_true')

    args = parser.parse_args()

    logger = None #dummy
            
    for folder in args.folders:
        #initialize logger for this folder:

        logfile = os.path.join(folder, 'run.log')
        logger = initLogging(logfile, logger)

        try:
            inifile = os.path.join(folder, args.inifile)

            f = FPPCalculation.from_ini(inifile, recalc=args.recalc,
                                        n=args.n)
            f.FPPplots()

            starmodel_file = os.path.join(folder, 'starmodel.h5')
            triangle_file1 = os.path.join(folder, 'starmodel_triangle_physical.png')
            triangle_file2 = os.path.join(folder, 'starmodel_triangle_observed.png')
            if not os.path.exists(triangle_file1) or not os.path.exists(triangle_file2):
                logger.info('Making StarModel triangle plots...')
                starmodel = StarModel.load_hdf(starmodel_file)
                triangle_base = os.path.join(folder, 'starmodel_triangle')
                starmodel.triangle_plots(triangle_base)

            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)



