#!/usr/bin/env python2
# -*- coding: utf-8 -*-
"""
Created on Fri Sep 28 15:19:30 2018

@author:
Dr. Maximilian N. Günther
European Space Agency (ESA)
European Space Research and Technology Centre (ESTEC)
Keplerlaan 1, 2201 AZ Noordwijk, The Netherlands
Email: maximilian.guenther@esa.int
GitHub: mnguenther
Twitter: m_n_guenther
Web: www.mnguenther.com
"""

from __future__ import print_function, division, absolute_import

#::: plotting settings
import seaborn as sns
sns.set(context='paper', style='ticks', palette='deep', font='sans-serif', font_scale=1.5, color_codes=True)
sns.set_style({"xtick.direction": "in","ytick.direction": "in"})
sns.set_context(rc={'lines.markeredgewidth': 1})

#::: modules
import os
#import collections
import numpy as np
import matplotlib.pyplot as plt
from matplotlib.ticker import ScalarFormatter
import pickle
from corner import corner
from tqdm import tqdm
from astropy.constants import M_earth, M_jup, M_sun, R_earth, R_jup, R_sun, au
import copy
from multiprocessing import Pool
from contextlib import closing

#::: allesfitter modules
from . import config
from .utils.latex_printer import round_tex
from .general_output import logprint
from .priors.simulate_PDF import simulate_PDF
from .limb_darkening import LDC3
from .computer import update_params, calculate_model, flux_fct, flux_subfct_ellc, flux_subfct_sinusoidal_phase_curves
from .exoworlds_rdx.lightcurves.index_transits import index_transits
from .lightcurves import get_epoch_occ

###############################################################################
#::: constants (replaced with astropy.constants)
###############################################################################
#M_earth = 5.9742e+24 	#kg 	Earth mass
#M_jup   = 1.8987e+27 	#kg 	Jupiter mass
#M_sun   = 1.9891e+30 	#kg 	Solar mass
#R_earth = 6378136      #m 	Earth equatorial radius
#R_jup   = 71492000 	#m 	Jupiter equatorial radius
#R_sun   = 695508000 	#m 	Solar radius

###############################################################################
#::: globals
#::: sorry for that... it's multiprocessing, not me, I swear!
###############################################################################
# companion = None
# inst = None
# samples2 = None
# derived_samples = None

###############################################################################
#::: calculate values from model curves
###############################################################################
def calculate_values_from_model_curves(p, inst, companion):
    '''
    Parameters
    ----------
    p : dict
        parameters corresponding to one single sample
    inst : str
        instrument name
    companion : str
        companion name

    Returns
    -------
    list
        list containing the transit depth, occultation depth, and nightside flux
    '''

    #==========================================================================
    #::: init
    #==========================================================================
    depth_tr = np.nan
    depth_occ = np.nan
    nightside_flux = np.nan
    epoch_occ = get_epoch_occ(p[companion+'_epoch'], p[companion+'_period'], p[companion+'_f_s'], p[companion+'_f_c'])

    #==========================================================================
    #:: calculating
    #==========================================================================

    #::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::
    #::: without phase curve
    #::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::
    if (config.BASEMENT.settings['phase_curve'] is False):
        #::: compute transit / primary eclipse depth
        depth_tr = 1. - flux_subfct_ellc(p, inst, companion, xx=[p[companion+'_epoch']])[0]

        #::: compute occultation / secondary eclipse depth (if wished)
        if (config.BASEMENT.settings['secondary_eclipse'] is True):
            depth_occ = 1. - flux_subfct_ellc(p, inst, companion, xx=[epoch_occ])[0]

    #::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::
    #::: with phase curve sine_series or sine_physical
    #::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::
    elif (config.BASEMENT.settings['phase_curve'] is True) and (config.BASEMENT.settings['phase_curve_style'] in ['sine_series','sine_physical']):

        #::: 0: epoch, 1: epoch_occ
        xx0 = [p[companion+'_epoch'], epoch_occ]

        #::: for debugging
        # xx0 = np.linspace(p[companion+'_epoch']-0.25*p[companion+'_period'], p[companion+'_epoch']+0.75*p[companion+'_period'], 1001)

        #::: the full model flux with phase curve and dips
        phase_curve_dips = flux_fct(p, inst, companion, xx=xx0)

        #::: the phase curve without any dips
        ellc_flux, ellc_flux1, ellc_flux2 = flux_subfct_ellc(p, inst, companion, xx=xx0, return_fluxes=True)
        phase_curve_no_dips = flux_subfct_sinusoidal_phase_curves(p, inst, companion, np.ones_like(xx0), xx=xx0)

        #::: the phase curve with atmopsheric dips, but without nightside flux (sbratio=1e-12)
        p2 = copy.deepcopy(p)
        p2[companion+'_sbratio_'+inst] = 1e-12
        ellc_flux, ellc_flux1, ellc_flux2 = flux_subfct_ellc(p2, inst, companion, xx=xx0, return_fluxes=True)
        phase_curve_atmo_dips = flux_subfct_sinusoidal_phase_curves(p2, inst, companion, ellc_flux2, xx=xx0)

        #::: for debugging
        # fig = plt.figure()
        # plt.plot(xx0, phase_curve_dips, label='phase_curve_dips')
        # plt.plot(xx0, phase_curve_no_dips, label='phase_curve_no_dips')
        # plt.plot(xx0, phase_curve_atmo_dips, label='phase_curve_atmo_dips')
        # plt.legend()
        # plt.ylim([0.999,1.001])
        # plt.axhline(1,c='grey',ls='--')
        # fig.savefig(os.path.join(config.BASEMENT.outdir,'phase_curve_depths.pdf'), bbox_inches='tight')

        #::: compute transit / primary eclipse depth
        depth_tr = 1e3 * (phase_curve_no_dips[0] - phase_curve_dips[0]) #in ppt; 0: epoch

        #::: compute
        if (config.BASEMENT.settings['secondary_eclipse'] is True):

            #::: compute occultation / secondary eclipse depth
            depth_occ = 1e3 * (phase_curve_no_dips[1] - phase_curve_dips[1]) #in ppt; 1: epoch_occ

            #::: compute nightside flux
            nightside_flux = 1e3 * (phase_curve_atmo_dips[1] - phase_curve_dips[1]) #in ppt; 1: epoch_occ

    #::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::
    #::: with phase curve ellc_physical
    #::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::
    elif (config.BASEMENT.settings['phase_curve'] is True) and (config.BASEMENT.settings['phase_curve_style'] in ['ellc_physical']):
        pass #TODO: not yet implemented

    #==========================================================================
    #::: return
    #==========================================================================
    return [depth_tr, depth_occ, nightside_flux]

###############################################################################
#::: the main derive function
###############################################################################
def derive(samples, mode):
    '''
    Derives parameter of the system using Winn 2010

    Input:
    ------
    samples : array
        samples from the mcmc or nested sampling
    mode : str
        'mcmc' or 'ns'

    Returns:
    --------
    derived_samples : dict
        with keys 'i', 'R1a', 'R2a', 'k', 'depth_undiluted', 'b_tra', 'b_occ', 'Ttot', 'Tfull'
        each key contains all the samples derived from the MCMC samples
        (not mean values, but pure samples!)
        i = inclination
        R1a = R1/a, radius companion over semiamplitude
        R2a = R2/a, radius star over semiamplitude
        Ttot = T_{1-4}, total transit width
        Tfull = T_{2-3}, full-transit width

    Output:
    -------
    latex table of results
    corner plot of derived values posteriors
    '''

    #::: using a global keyword
    #::: sorry for that... it's multiprocessing, not me, I swear!
    global companion
    global inst
    global samples2
    global derived_samples

    samples2 = samples #global variable
    N_samples = samples.shape[0]

    #==========================================================================
    #::: stellar 'posteriors'
    #==========================================================================
    if os.path.exists( os.path.join(config.BASEMENT.datadir,'params_star.csv') ):
        buf = np.genfromtxt( os.path.join(config.BASEMENT.datadir,'params_star.csv'), delimiter=',', names=True, dtype=None, encoding='utf-8', comments='#' )
        star = {}
        star['R_star'] = simulate_PDF(buf['R_star'], buf['R_star_lerr'], buf['R_star_uerr'], size=N_samples, plot=False)
        star['M_star'] = simulate_PDF(buf['M_star'], buf['M_star_lerr'], buf['M_star_uerr'], size=N_samples, plot=False)
        star['Teff_star'] = simulate_PDF(buf['Teff_star'], buf['Teff_star_lerr'], buf['Teff_star_uerr'], size=N_samples, plot=False)
    else:
        star = {'R_star':np.nan, 'M_star':np.nan, 'Teff_star':np.nan}

    #==========================================================================
    #::: derive all the params
    #==========================================================================
    companions = config.BASEMENT.settings['companions_all']

    def get_params(key):
        ind = np.where(config.BASEMENT.fitkeys==key)[0]
        if len(ind)==1:
            return samples[:,ind].flatten() #if it was fitted for
        else:
            try:
                if config.BASEMENT.params[key] is None:
                    return np.nan #if None, retun nan instead
                else:
                    return config.BASEMENT.params[key] #else take the input value
            except KeyError:
                return np.nan #if all fails, return nan

    def sin_d(alpha): return np.sin(np.deg2rad(alpha))
    def cos_d(alpha): return np.cos(np.deg2rad(alpha))
    def arcsin_d(x): return np.rad2deg(np.arcsin(x))
    def arccos_d(x): return np.rad2deg(np.arccos(x))

    derived_samples = {}
    for cc in companions:
        companion = cc

        #----------------------------------------------------------------------
        #::: radii
        #----------------------------------------------------------------------
        derived_samples[companion+'_R_star/a'] = get_params(companion+'_rsuma') / (1. + get_params(companion+'_rr'))
        derived_samples[companion+'_a/R_star'] = (1. + get_params(companion+'_rr')) / get_params(companion+'_rsuma')
        derived_samples[companion+'_R_companion/a'] = get_params(companion+'_rsuma') * get_params(companion+'_rr') / (1. + get_params(companion+'_rr'))
        derived_samples[companion+'_R_companion_(R_earth)'] = star['R_star'] * get_params(companion+'_rr') * R_sun.value / R_earth.value #in R_earth
        derived_samples[companion+'_R_companion_(R_jup)'] = star['R_star'] * get_params(companion+'_rr') * R_sun.value / R_jup.value #in R_jup

        #----------------------------------------------------------------------
        #::: orbit
        #----------------------------------------------------------------------
        derived_samples[companion+'_a_(R_sun)'] = star['R_star'] / derived_samples[companion+'_R_star/a']
        derived_samples[companion+'_a_(AU)'] = derived_samples[companion+'_a_(R_sun)'] * R_sun.value/au.value
        derived_samples[companion+'_i'] = arccos_d(get_params(companion+'_cosi')) #in deg
        derived_samples[companion+'_e'] = get_params(companion+'_f_s')**2 + get_params(companion+'_f_c')**2
        derived_samples[companion+'_e_sinw'] = get_params(companion+'_f_s') * np.sqrt(derived_samples[companion+'_e'])
        derived_samples[companion+'_e_cosw'] = get_params(companion+'_f_c') * np.sqrt(derived_samples[companion+'_e'])
        derived_samples[companion+'_w'] = np.rad2deg(np.mod( np.arctan2(get_params(companion+'_f_s'), get_params(companion+'_f_c')), 2*np.pi) ) #in deg, from 0 to 360
        if np.isnan(derived_samples[companion+'_w']).all():
            derived_samples[companion+'_w'] = 0.

        #----------------------------------------------------------------------
        #::: masses
        #----------------------------------------------------------------------
        #::: for detached binaries, where K and q were fitted:
        if (companion+'_K' in config.BASEMENT.params) and len(config.BASEMENT.settings['inst_rv2'])>0:
            derived_samples[companion+'_M_companion_(M_earth)'] = get_params(companion+'_q') * star['M_star'] * M_sun.value / M_earth.value #in M_earth
            derived_samples[companion+'_M_companion_(M_jup)'] = get_params(companion+'_q') * star['M_star'] * M_sun.value / M_jup.value #in M_jup
            derived_samples[companion+'_M_companion_(M_sun)'] = get_params(companion+'_q') * star['M_star'] #in M_sun

        #::: for exoplanets or single-lined binaries, where only K was fitted, approximate/best-guess q form K:
        elif companion+'_K' in config.BASEMENT.params:
            a_1 = 0.019771142 * get_params(companion+'_K') * get_params(companion+'_period') * np.sqrt(1. - derived_samples[companion+'_e']**2)/sin_d(derived_samples[companion+'_i'])
    #        derived_samples[companion+'_a_rv'] = (1.+1./ellc_params[companion+'_q'])*a_1
            derived_samples[companion+'_q'] = 1./(( derived_samples[companion+'_a_(R_sun)'] / a_1 ) - 1.)
            derived_samples[companion+'_M_companion_(M_earth)'] = derived_samples[companion+'_q'] * star['M_star'] * M_sun.value / M_earth.value #in M_earth
            derived_samples[companion+'_M_companion_(M_jup)'] = derived_samples[companion+'_q'] * star['M_star'] * M_sun.value / M_jup.value #in M_jup
            derived_samples[companion+'_M_companion_(M_sun)'] = derived_samples[companion+'_q'] * star['M_star'] #in M_sun

        #----------------------------------------------------------------------
        #::: time of secondary eclipse
        #----------------------------------------------------------------------
        if config.BASEMENT.settings['secondary_eclipse'] is True:
            derived_samples[companion+'_epoch_occ'] = get_params(companion+'_epoch') + get_params(companion+'_period')/2. * (1. + 4./np.pi * derived_samples[companion+'_e'] * cos_d(derived_samples[companion+'_w'])  ) #approximation from Winn2010

        #----------------------------------------------------------------------
        #::: impact params of primary eclipse with eccentricity corrections (from Winn 2010)
        #----------------------------------------------------------------------
        eccentricity_correction_b_tra = ( (1. - derived_samples[companion+'_e']**2) / ( 1. + derived_samples[companion+'_e']*sin_d(derived_samples[companion+'_w']) ) )

        derived_samples[companion+'_b_tra'] = (1./derived_samples[companion+'_R_star/a']) * get_params(companion+'_cosi') * eccentricity_correction_b_tra

        #----------------------------------------------------------------------
        #::: impact params of secondary eclipse with eccentricity corrections (from Winn 2010)
        #----------------------------------------------------------------------
        eccentricity_correction_b_occ = ( (1. - derived_samples[companion+'_e']**2) / ( 1. - derived_samples[companion+'_e']*sin_d(derived_samples[companion+'_w']) ) )

        if config.BASEMENT.settings['secondary_eclipse'] is True:
            derived_samples[companion+'_b_occ'] = (1./derived_samples[companion+'_R_star/a']) * get_params(companion+'_cosi') * eccentricity_correction_b_occ

        #----------------------------------------------------------------------
        #::: transit duration (in hours) with eccentricity corrections (from Winn 2010)
        #----------------------------------------------------------------------
        eccentricity_correction_T_tra = ( np.sqrt(1. - derived_samples[companion+'_e']**2) / ( 1. + derived_samples[companion+'_e']*sin_d(derived_samples[companion+'_w']) ) )

        derived_samples[companion+'_T_tra_tot'] = get_params(companion+'_period')/np.pi *24.  \
                                                  * np.arcsin( derived_samples[companion+'_R_star/a'] \
                                                               * np.sqrt( (1. + get_params(companion+'_rr'))**2 - derived_samples[companion+'_b_tra']**2 ) \
                                                               / sin_d(derived_samples[companion+'_i']) ) \
                                                  * eccentricity_correction_T_tra    #in h
        derived_samples[companion+'_T_tra_full'] = get_params(companion+'_period')/np.pi *24.  \
                                                   * np.arcsin( derived_samples[companion+'_R_star/a'] \
                                                                * np.sqrt( (1. - get_params(companion+'_rr'))**2 - derived_samples[companion+'_b_tra']**2  )\
                                                                / sin_d(derived_samples[companion+'_i']) ) \
                                                   * eccentricity_correction_T_tra    #in h

        #----------------------------------------------------------------------
        #::: primary and secondary eclipse depths (per inst)
        #::: / transit and occultation depths (per inst)
        #----------------------------------------------------------------------
        for ii in config.BASEMENT.settings['inst_phot']:

            #- - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -
            #::: setup
            #- - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -
            inst = ii
            N_less_samples = 1000
            derived_samples[companion+'_depth_tr_dil_'+inst] = np.nan*np.empty(N_less_samples)
            derived_samples[companion+'_depth_occ_dil_'+inst] = np.nan*np.empty(N_less_samples)
            derived_samples[companion+'_nightside_flux_dil_'+inst] = np.nan*np.empty(N_less_samples)

            #- - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -
            #::: iterate through all samples, draw different models and measure the depths
            #- - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -
            print('Deriving eclipse depths (and more) from the model curves for companion', companion, 'and instrument', inst+'...')
            for i in range(N_less_samples):
                s = samples[ np.random.randint(low=0,high=samples2.shape[0]) , : ]
                p = update_params(s)
                r = calculate_values_from_model_curves(p, inst, companion)
                derived_samples[companion+'_depth_tr_dil_'+inst][i] = r[0]
                derived_samples[companion+'_depth_occ_dil_'+inst][i] = r[1]
                derived_samples[companion+'_nightside_flux_dil_'+inst][i] = r[2]

            #- - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -
            #resize the arrays to match the true N_samples (by redrawing the 1000 values)
            #- - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -
            derived_samples[companion+'_depth_tr_dil_'+inst] = np.resize(derived_samples[companion+'_depth_tr_dil_'+inst], N_samples)
            derived_samples[companion+'_depth_occ_dil_'+inst] = np.resize(derived_samples[companion+'_depth_occ_dil_'+inst], N_samples)
            derived_samples[companion+'_nightside_flux_dil_'+inst] = np.resize(derived_samples[companion+'_nightside_flux_dil_'+inst], N_samples)

        #----------------------------------------------------------------------
        #::: undiluted (per companion; per inst)
        #----------------------------------------------------------------------
        for inst in config.BASEMENT.settings['inst_phot']:
            dil = get_params('dil_'+inst)
            if all(np.atleast_1d(np.isnan(dil))): dil = 0
            derived_samples[companion+'_depth_tr_undil_'+inst] = derived_samples[companion+'_depth_tr_dil_'+inst] / (1. - dil) #in ppt
            derived_samples[companion+'_depth_occ_undil_'+inst] = derived_samples[companion+'_depth_occ_dil_'+inst] / (1. - dil) #in ppt
            derived_samples[companion+'_nightside_flux_undil_'+inst] = derived_samples[companion+'_nightside_flux_dil_'+inst] / (1. - dil) #in ppt

        #----------------------------------------------------------------------
        #::: equilibirum temperature
        #::: currently assumes Albedo of 0.3 and Emissivity of 1
        #----------------------------------------------------------------------
        albedo = 0.3
        emissivity = 1.
        derived_samples[companion+'_Teq'] = star['Teff_star']  * ( (1.-albedo)/emissivity )**0.25 * np.sqrt(derived_samples[companion+'_R_star/a'] / 2.)

        #----------------------------------------------------------------------
        #::: stellar density from orbit
        #----------------------------------------------------------------------
        if companion in config.BASEMENT.settings['companions_phot']:
            if all(np.atleast_1d(get_params(companion+'_rr'))<0.215443469): #see computer.py; get_params could return np.nan (float) or array; all(np.atleast_1d(...)) takes care of that
                derived_samples[companion+'_host_density'] = 3. * np.pi * (1./derived_samples[companion+'_R_star/a'])**3. / (get_params(companion+'_period')*86400.)**2 / 6.67408e-8 #in cgs

        #----------------------------------------------------------------------
        #::: companion densities
        #----------------------------------------------------------------------
        derived_samples[companion+'_density'] = ( (derived_samples[companion+'_M_companion_(M_earth)'] * M_earth) / (4./3. * np.pi * (derived_samples[companion+'_R_companion_(R_earth)'] * R_earth)**3 ) ).cgs.value #in cgs

        #----------------------------------------------------------------------
        #::: the companion's surface gravity (individual posterior distribution for each companion; via Southworth+ 2007)
        #----------------------------------------------------------------------
        try:
            derived_samples[companion+'_surface_gravity'] = 2. * np.pi / (get_params(companion+'_period')*86400.) * np.sqrt((1.-derived_samples[companion+'_e']**2)) * (get_params(companion+'_K')*1e5) / (derived_samples[companion+'_R_companion/a'])**2 / sin_d(derived_samples[companion+'_i'])
        except:
            pass

        #----------------------------------------------------------------------
        #::: period ratios (for ressonance studies)
        #----------------------------------------------------------------------
        if len(companions)>1:
            for other_companion in companions:
                if other_companion is not companion:
                    derived_samples[companion+'_period/'+other_companion+'_period'] = get_params(companion+'_period') / get_params(other_companion+'_period')

        #----------------------------------------------------------------------
        #::: limb darkening
        #----------------------------------------------------------------------
        for inst in config.BASEMENT.settings['inst_all']:

            #- - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -
            #::: host
            #- - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -
            if config.BASEMENT.settings['host_ld_law_'+inst] is None:
                pass

            elif config.BASEMENT.settings['host_ld_law_'+inst] == 'lin':
                derived_samples['host_ldc_u1_'+inst] = get_params('host_ldc_q1_'+inst)

            elif config.BASEMENT.settings['host_ld_law_'+inst] == 'quad':
                derived_samples['host_ldc_u1_'+inst] = 2 * np.sqrt(get_params('host_ldc_q1_'+inst)) * get_params('host_ldc_q2_'+inst)
                derived_samples['host_ldc_u2_'+inst] = np.sqrt(get_params('host_ldc_q1_'+inst)) * (1. - 2. * get_params('host_ldc_q2_'+inst))

            elif config.BASEMENT.settings['host_ld_law_'+inst] == 'sing':
                derived_samples['host_ldc_u1_'+inst] = np.nan*np.empty(N_samples)
                derived_samples['host_ldc_u2_'+inst] = np.nan*np.empty(N_samples)
                derived_samples['host_ldc_u3_'+inst] = np.nan*np.empty(N_samples)
                for i in range(N_samples):
                    u1, u2, u3 = LDC3.forward([get_params('host_ldc_q1_'+inst)[i], get_params('host_ldc_q2_'+inst)[i], get_params('host_ldc_q3_'+inst)[i]])
                    derived_samples['host_ldc_u1_'+inst][i] = u1
                    derived_samples['host_ldc_u2_'+inst][i] = u2
                    derived_samples['host_ldc_u3_'+inst][i] = u3

            else:
                raise ValueError("Currently only 'none', 'lin', 'quad' and 'sing' limb darkening are supported.")

    #==========================================================================
    #::: median stellar density
    #==========================================================================
    derived_samples['combined_host_density'] = []
    for companion in config.BASEMENT.settings['companions_phot']:
        try: derived_samples['combined_host_density'] = np.append(derived_samples['combined_host_density'], derived_samples[companion+'_host_density'])
        except: pass

    ###############################################################################
    #::: write keys for output
    ###############################################################################
    names = []
    labels = []
    for companion in companions:

        names.append( companion+'_R_star/a' )
        labels.append( 'Host radius over semi-major axis '+companion+'; $R_\star/a_\mathrm{'+companion+'}$' )

        names.append( companion+'_a/R_star' )
        labels.append( 'Semi-major axis '+companion+' over host radius; $a_\mathrm{'+companion+'}/R_\star$' )

        names.append( companion+'_R_companion/a'  )
        labels.append( 'Companion radius '+companion+' over semi-major axis '+companion+'; $R_\mathrm{'+companion+'}/a_\mathrm{'+companion+'}$' )

        names.append( companion+'_R_companion_(R_earth)' )
        labels.append( 'Companion radius '+companion+'; $R_\mathrm{'+companion+'}$ ($\mathrm{R_{\oplus}}$)' )

        names.append( companion+'_R_companion_(R_jup)' )
        labels.append( 'Companion radius '+companion+'; $R_\mathrm{'+companion+'}$ ($\mathrm{R_{jup}}$)' )

        names.append( companion+'_a_(R_sun)' )
        labels.append( 'Semi-major axis '+companion+'; $a_\mathrm{'+companion+'}$ ($\mathrm{R_{\odot}}$)' )

        names.append( companion+'_a_(AU)' )
        labels.append( 'Semi-major axis '+companion+'; $a_\mathrm{'+companion+'}$ (AU)' )

        names.append( companion+'_i' )
        labels.append( 'Inclination '+companion+'; $i_\mathrm{'+companion+'}$ (deg)' )

        names.append( companion+'_e' )
        labels.append( 'Eccentricity '+companion+'; $e_\mathrm{'+companion+'}$' )

        names.append( companion+'_w' )
        labels.append( 'Argument of periastron '+companion+'; $w_\mathrm{'+companion+'}$ (deg)' )

        names.append( companion+'_q' )
        labels.append( 'Mass ratio '+companion+'; $q_\mathrm{'+companion+'}$' )

        names.append( companion+'_M_companion_(M_earth)' )
        labels.append( 'Companion mass '+companion+'; $M_\mathrm{'+companion+'}$ ($\mathrm{M_{\oplus}}$)' )

        names.append( companion+'_M_companion_(M_jup)' )
        labels.append( 'Companion mass '+companion+'; $M_\mathrm{'+companion+'}$ ($\mathrm{M_{jup}}$)' )

        names.append( companion+'_M_companion_(M_sun)' )
        labels.append( 'Companion mass '+companion+'; $M_\mathrm{'+companion+'}$ ($\mathrm{M_{\odot}}$)' )

        names.append( companion+'_b_tra' )
        labels.append( 'Impact parameter '+companion+'; $b_\mathrm{tra;'+companion+'}$' )

        names.append( companion+'_T_tra_tot'  )
        labels.append( 'Total transit duration '+companion+'; $T_\mathrm{tot;'+companion+'}$ (h)' )

        names.append( companion+'_T_tra_full' )
        labels.append( 'Full-transit duration '+companion+'; $T_\mathrm{full;'+companion+'}$ (h)' )

        names.append( companion+'_epoch_occ'  )
        labels.append( 'Epoch occultation '+companion+'; $T_\mathrm{0;occ;'+companion+'}$' )

        names.append( companion+'_b_occ'  )
        labels.append( 'Impact parameter occultation '+companion+'; $b_\mathrm{occ;'+companion+'}$' )

        names.append( companion+'_host_density' )
        labels.append( 'Host density from orbit '+companion+'; $\\rho_\mathrm{\star;'+companion+'}$ (cgs)' )

        names.append( companion+'_density' )
        labels.append( 'Companion density '+companion+'; $\\rho_\mathrm{'+companion+'}$ (cgs)' )

        names.append( companion+'_surface_gravity')
        labels.append( 'Companion surface gravity '+companion+'; $g_\mathrm{'+companion+'}$ (cgs)' )

        names.append( companion+'_Teq' )
        labels.append( 'Equilibrium temperature '+companion+'; $T_\mathrm{eq;'+companion+'}$ (K)' )

        for inst in config.BASEMENT.settings['inst_phot']:

            names.append( companion+'_depth_tr_undil_'+inst )
            labels.append( 'Transit depth (undil.) '+companion+'; $\delta_\mathrm{tr; undil; '+companion+'; '+inst+'}$ (ppt)' )

            names.append( companion+'_depth_tr_dil_'+inst )
            labels.append( 'Transit depth (dil.) '+companion+'; $\delta_\mathrm{tr; dil; '+companion+'; '+inst+'}$ (ppt)' )

            names.append( companion+'_depth_occ_undil_'+inst )
            labels.append( 'Occultation depth (undil.) '+companion+'; $\delta_\mathrm{occ; undil; '+companion+'; '+inst+'}$ (ppt)' )

            names.append( companion+'_depth_occ_dil_'+inst )
            labels.append( 'Occultation depth (dil.) '+companion+'; $\delta_\mathrm{occ; dil; '+companion+'; '+inst+'}$ (ppt)' )

            names.append( companion+'_nightside_flux_undil_'+inst )
            labels.append( 'Nightside flux (undil.)'+companion+'; $F_\mathrm{nightside; undil; '+companion+'; '+inst+'}$ (ppt)' )

            names.append( companion+'_nightside_flux_dil_'+inst )
            labels.append( 'Nightside flux (dil.)'+companion+'; $F_\mathrm{nightside; dil; '+companion+'; '+inst+'}$ (ppt)' )

        #::: period ratios (for ressonance studies)
        if len(companions)>1:
            for other_companion in companions:
                if other_companion is not companion:
                    names.append( companion+'_period/'+other_companion+'_period' )
                    labels.append( 'Period ratio; $P_\mathrm{'+companion+'} / P_\mathrm{'+other_companion+'}$' )

    #::: host
    for inst in config.BASEMENT.settings['inst_all']:
        if config.BASEMENT.settings['host_ld_law_'+inst] is None:
            pass

        elif config.BASEMENT.settings['host_ld_law_'+inst] == 'lin':
            names.append( 'host_ldc_u1_'+inst )
            labels.append( 'Limb darkening; $u_\mathrm{1; '+inst+'}$' )

        elif config.BASEMENT.settings['host_ld_law_'+inst] == 'quad':
            names.append( 'host_ldc_u1_'+inst )
            labels.append( 'Limb darkening; $u_\mathrm{1; '+inst+'}$' )
            names.append( 'host_ldc_u2_'+inst )
            labels.append( 'Limb darkening; $u_\mathrm{2; '+inst+'}$' )

        elif config.BASEMENT.settings['host_ld_law_'+inst] == 'sing':
            names.append( 'host_ldc_u1_'+inst )
            labels.append( 'Limb darkening; $u_\mathrm{1; '+inst+'}$' )
            names.append( 'host_ldc_u2_'+inst )
            labels.append( 'Limb darkening; $u_\mathrm{2; '+inst+'}$' )
            names.append( 'host_ldc_u3_'+inst )
            labels.append( 'Limb darkening; $u_\mathrm{3; '+inst+'}$' )

        else:
            raise ValueError("Currently only 'none', 'lin', 'quad' and 'sing' limb darkening are supported.")

    names.append( 'combined_host_density' )
    labels.append( 'Combined host density from all orbits; $rho_\mathrm{\star; combined}$ (cgs)' )

    ###############################################################################
    #::: delete pointless values
    ###############################################################################
    ind_good = []
    for i,name in enumerate(names):
        if (name in derived_samples) and isinstance(derived_samples[name], np.ndarray) and not all(np.isnan(derived_samples[name])) and not all(np.array(derived_samples[name])==0):
            ind_good.append(i)

    names = [ names[i] for i in ind_good ]
    labels = [ labels[i] for i in ind_good ]

    ###############################################################################
    #::: if any meaningful values are left, go output them
    ###############################################################################
    if len(names)>0:

        #=====================================================================
        #::: save all in pickle
        #=====================================================================
        pickle.dump(derived_samples, open(os.path.join(config.BASEMENT.outdir,mode+'_derived_samples.pickle'),'wb'))

        #=====================================================================
        #::: save txt & latex table & latex commands
        #=====================================================================
        with open(os.path.join(config.BASEMENT.outdir,mode+'_derived_table.csv'),'w') as outfile,\
             open(os.path.join(config.BASEMENT.outdir,mode+'_derived_latex_table.txt'),'w') as f,\
             open(os.path.join(config.BASEMENT.outdir,mode+'_derived_latex_cmd.txt'),'w') as f_cmd:

            outfile.write('#property,value,lower_error,upper_error,source\n')

            f.write('Parameter & Value & Source \\\\ \n')
            f.write('\\hline \n')
            f.write('\\multicolumn{3}{c}{\\textit{Derived parameters}} \\\\ \n')
            f.write('\\hline \n')

            for name,label in zip(names, labels):
                ll, median, ul = np.nanpercentile(derived_samples[name], [15.865, 50., 84.135])
                outfile.write( str(label)+','+str(median)+','+str(median-ll)+','+str(ul-median)+',derived\n' )

                value = round_tex(median, median-ll, ul-median)
                f.write( label + ' & $' + value + '$ & derived \\\\ \n' )

                simplename = name.replace("_", "").replace("/", "over").replace("(", "").replace(")", "").replace("1", "one").replace("2", "two")
                f_cmd.write('\\newcommand{\\'+simplename+'}{$'+value+'$} %'+label+' = $'+value+'$\n')

        logprint('\nSaved '+mode+'_derived_results.csv, '+mode+'_derived_latex_table.txt, and '+mode+'_derived_latex_cmd.txt')

        #=====================================================================
        #::: plot corner
        #=====================================================================
        if 'combined_host_density' in names: names.remove('combined_host_density') #has (N_companions x N_dims) dimensions, thus does not match the rest

        #::: clean up any isolated NaN's before calling corner
        for name in names:
            median = np.nanmedian(derived_samples[name])
            ind = np.where(np.isnan(derived_samples[name]))
            derived_samples[name][ind] = median

        #::: prep the matrix for corner
        x = np.column_stack([ derived_samples[name] for name in names ])
        fontsize = np.min(( 24. + 0.5*(len(names)), 40 ))

        fig = corner(x,
                     range = [0.999]*len(names),
                     labels = names,
                     quantiles=[0.15865, 0.5, 0.84135],
                     show_titles=True,
                     label_kwargs={"fontsize":fontsize, "rotation":45, "horizontalalignment":'right'},
                     max_n_ticks=3)
        caxes = np.reshape(np.array(fig.axes), (len(names),len(names)))

        #::: set allesfitter titles
        for i, name in enumerate(names):

            ll, median, ul = np.nanpercentile(derived_samples[name], [15.865, 50., 84.135])
            value = round_tex(median, median-ll, ul-median)
            ctitle = r'' + labels[i] + '\n' + r'$=' + value + '$'
            if len(names)>1:
                # caxes[i,i].set_title(ctitle)
                caxes[i,i].set_title(ctitle, fontsize=fontsize, rotation=45, horizontalalignment='left')
                for i in range(caxes.shape[0]):
                    for j in range(caxes.shape[1]):
                        caxes[i,j].xaxis.set_label_coords(0.5, -0.5)
                        caxes[i,j].yaxis.set_label_coords(-0.5, 0.5)

                        if i==(caxes.shape[0]-1):
                            fmt = ScalarFormatter(useOffset=False)
                            caxes[i,j].xaxis.set_major_formatter(fmt)
                        if (i>0) and (j==0):
                            fmt = ScalarFormatter(useOffset=False)
                            caxes[i,j].yaxis.set_major_formatter(fmt)

                        for tick in caxes[i,j].xaxis.get_major_ticks(): tick.label.set_fontsize(24)
                        for tick in caxes[i,j].yaxis.get_major_ticks(): tick.label.set_fontsize(24)
            else:
                caxes.set_title(ctitle)
                caxes.xaxis.set_label_coords(0.5, -0.5)
                caxes.yaxis.set_label_coords(-0.5, 0.5)

        dpi = np.max(( 100. - len(names), 50 ))
        try: #some matplitlib versions cannot handle jpg
            fig.savefig( os.path.join(config.BASEMENT.outdir,mode+'_derived_corner.jpg'), dpi=dpi, bbox_inches='tight' )
        except:
            fig.savefig( os.path.join(config.BASEMENT.outdir,mode+'_derived_corner.png'), bbox_inches='tight' )
        plt.close(fig)

        #=====================================================================
        #::: finish
        #=====================================================================
        logprint('\nSaved '+mode+'_derived_corner.pdf')

    else:
        logprint('\nNo values available to be derived.')


#   Copyright 2020 The PyMC Developers
#
#   Licensed under the Apache License, Version 2.0 (the "License");
#   you may not use this file except in compliance with the License.
#   You may obtain a copy of the License at
#
#       http://www.apache.org/licenses/LICENSE-2.0
#
#   Unless required by applicable law or agreed to in writing, software
#   distributed under the License is distributed on an "AS IS" BASIS,
#   WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
#   See the License for the specific language governing permissions and
#   limitations under the License.

import warnings

from collections import OrderedDict

import theano
import theano.tensor as tt

from pymc3.model import inputvars, modelcontext
from pymc3.step_methods.arraystep import ArrayStepShared
from pymc3.theanof import make_shared_replacements, tt_rng

__all__ = []

EXPERIMENTAL_WARNING = (
    "Warning: Stochastic Gradient based sampling methods are experimental step methods and not yet"
    " recommended for use in PyMC3!"
)

def _value_error(cond, str):
    """Throws ValueError if cond is False"""
    if not cond:
        raise ValueError(str)

def _check_minibatches(minibatch_tensors, minibatches):
    _value_error(isinstance(minibatch_tensors, list), "minibatch_tensors must be a list.")

    _value_error(hasattr(minibatches, "__iter__"), "minibatches must be an iterator.")

def prior_dlogp(vars, model, flat_view):
    """Returns the gradient of the prior on the parameters as a vector of size D x 1"""
    terms = tt.concatenate([theano.grad(var.logpt, var).flatten() for var in vars], axis=0)
    dlogp = theano.clone(terms, flat_view.replacements, strict=False)

    return dlogp

def elemwise_dlogL(vars, model, flat_view):
    """
    Returns Jacobian of the log likelihood for each training datum wrt vars
    as a matrix of size N x D
    """
    # select one observed random variable
    obs_var = model.observed_RVs[0]
    # tensor of shape (batch_size,)
    logL = obs_var.logp_elemwiset.sum(axis=tuple(range(1, obs_var.logp_elemwiset.ndim)))
    # calculate fisher information
    terms = []
    for var in vars:
        output, _ = theano.scan(
            lambda i, logX=logL, v=var: theano.grad(logX[i], v).flatten(),
            sequences=[tt.arange(logL.shape[0])],
        )
        terms.append(output)
    dlogL = theano.clone(tt.concatenate(terms, axis=1), flat_view.replacements, strict=False)
    return dlogL

class BaseStochasticGradient(ArrayStepShared):
    R"""
    BaseStochasticGradient Object

    For working with BaseStochasticGradient Object
    we need to supply the probabilistic model
    (:code:`model`) with the data supplied to observed
    variables of type `GeneratorOp`

    Parameters
    ----------
    vars: list
        List of variables for sampler
    batch_size`: int
        Batch Size for each step
    total_size: int
        Total size of the training data
    step_size: float
        Step size for the parameter update
    model: PyMC Model
        Optional model for sampling step. Defaults to None (taken from context)
    random_seed: int
        The seed to initialize the Random Stream
    minibatches: iterator
        If the ObservedRV.observed is not a GeneratorOp then this parameter must not be None
    minibatch_tensor: list of tensors
        If the ObservedRV.observed is not a GeneratorOp then this parameter must not be None
        The length of this tensor should be the same as the next(minibatches)

    Notes
    -----
    Defining a BaseStochasticGradient needs
    custom implementation of the following methods:
        - :code: `.mk_training_fn()`
            Returns a theano function which is called for each sampling step
        - :code: `._initialize_values()`
            Returns None it creates class variables which are required for the training fn
    """

    def __init__(
        self,
        vars=None,
        batch_size=None,
        total_size=None,
        step_size=1.0,
        model=None,
        random_seed=None,
        minibatches=None,
        minibatch_tensors=None,
        **kwargs
    ):
        warnings.warn(EXPERIMENTAL_WARNING)

        model = modelcontext(model)

        if vars is None:
            vars = model.vars

        vars = inputvars(vars)

        self.model = model
        self.vars = vars
        self.batch_size = batch_size
        self.total_size = total_size
        _value_error(
            total_size != None or batch_size != None,
            "total_size and batch_size of training data have to be specified",
        )
        self.expected_iter = int(total_size / batch_size)

        # set random stream
        self.random = None
        if random_seed is None:
            self.random = tt_rng()
        else:
            self.random = tt_rng(random_seed)

        self.step_size = step_size

        shared = make_shared_replacements(vars, model)

        self.updates = OrderedDict()
        self.q_size = int(sum(v.dsize for v in self.vars))

        flat_view = model.flatten(vars)
        self.inarray = [flat_view.input]

        self.dlog_prior = prior_dlogp(vars, model, flat_view)
        self.dlogp_elemwise = elemwise_dlogL(vars, model, flat_view)
        self.q_size = int(sum(v.dsize for v in self.vars))

        if minibatch_tensors != None:
            _check_minibatches(minibatch_tensors, minibatches)
            self.minibatches = minibatches

            # Replace input shared variables with tensors
            def is_shared(t):
                return isinstance(t, theano.compile.sharedvalue.SharedVariable)

            tensors = [(t.type() if is_shared(t) else t) for t in minibatch_tensors]
            updates = OrderedDict(
                {t: t_ for t, t_ in zip(minibatch_tensors, tensors) if is_shared(t)}
            )
            self.minibatch_tensors = tensors
            self.inarray += self.minibatch_tensors
            self.updates.update(updates)

        self._initialize_values()
        super().__init__(vars, shared)

    def _initialize_values(self):
        """Initializes the parameters for the stochastic gradient minibatch
        algorithm"""
        raise NotImplementedError

    def mk_training_fn(self):
        raise NotImplementedError

    def training_complete(self):
        """Returns boolean if astep has been called expected iter number of times"""
        return self.expected_iter == self.t

    def astep(self, q0):
        """Perform a single update in the stochastic gradient method.

        Returns new shared values and values sampled
        The size and ordering of q0 and q must be the same
        Parameters
        -------
        q0: list
            List of shared values and values sampled from last estimate

        Returns
        -------
        q
        """
        if hasattr(self, "minibatch_tensors"):
            return q0 + self.training_fn(q0, *next(self.minibatches))
        else:
            return q0 + self.training_fn(q0)


#This software is Copyright 2012 The Regents of the University of California. All Rights Reserved.
#Permission to use, copy, modify, and distribute this software and its documentation for educational, research and non-profit purposes for non-profit institutions, without fee, and without a written agreement is hereby granted, provided that the above copyright notice, this paragraph and the following three paragraphs appear in all copies.
#Permission to make commercial use of this software may be obtained by contacting:
#Technology Transfer Office
#9500 Gilman Drive, Mail Code 0910
#University of California
#La Jolla, CA 92093-0910
#(858) 534-5815
#invent@ucsd.edu
#This software program and documentation are copyrighted by The Regents of the University of California. The software program and documentation are supplied "as is", without any accompanying services from The Regents. The Regents does not warrant that the operation of the program will be uninterrupted or error-free. The end-user understands that the program was developed for research purposes and is advised not to rely exclusively on the program for any reason.
#IN NO EVENT SHALL THE UNIVERSITY OF CALIFORNIA BE LIABLE TO
#ANY PARTY FOR DIRECT, INDIRECT, SPECIAL, INCIDENTAL, OR
#CONSEQUENTIAL DAMAGES, INCLUDING LOST PROFITS, ARISING
#OUT OF THE USE OF THIS SOFTWARE AND ITS DOCUMENTATION,
#EVEN IF THE UNIVERSITY OF CALIFORNIA HAS BEEN ADVISED OF
#THE POSSIBILITY OF SUCH DAMAGE. THE UNIVERSITY OF
#CALIFORNIA SPECIFICALLY DISCLAIMS ANY WARRANTIES,
#INCLUDING, BUT NOT LIMITED TO, THE IMPLIED WARRANTIES OF
#MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE.
#THE SOFTWARE PROVIDED HEREUNDER IS ON AN "AS IS" BASIS, AND THE UNIVERSITY OF CALIFORNIA HAS NO OBLIGATIONS TO
#PROVIDE MAINTENANCE, SUPPORT, UPDATES, ENHANCEMENTS, OR
#MODIFICATIONS.

# Functions for comparing contours

from numpy import *
from geometry import *
from point_set import *

def biggestGap(contour1, contour2):
    """See Hausdorff distance."""

    biggestGapValue = 0

    locations1 = contour1.locations()
    locations2 = contour2.locations()

    for index1 in range(len(locations1)):
        point1 = array(locations1[index1])

        # the minimum distance represents the gap from point1 to point2

        # initialize gap
        gap = distance(point1, array(locations2[0]))

        # find the minimum
        #print "find the minimum"
        for coordList2 in locations2:
            point2 = array(coordList2)
            dist = distance(point1, point2)
            #print dist
            if dist < gap:
                gap = dist

        # if the gap is the largest found so far, keep it
        if gap > biggestGapValue:
            biggestGapValue = gap

    return biggestGapValue

def overlap_old(contour1, contour2):
    """Deprecated"""

    temp1 = zeros((1000, 1000), dtype=int8)
    temp2 = zeros((1000, 1000), dtype=int8)

    binaryImage1 = contour1.binaryImage
    binaryImage2 = contour2.binaryImage

    boundingBox1 = contour1.get2DBoundingBox()
    boundingBox2 = contour2.get2DBoundingBox()

    temp1[boundingBox1[0][0]:boundingBox1[1][0]+1,
          boundingBox1[0][1]:boundingBox1[1][1]+1] = binaryImage1

    temp2[boundingBox2[0][0]:boundingBox2[1][0]+1,
          boundingBox2[0][1]:boundingBox2[1][1]+1] = binaryImage2

    andImage = logical_and(temp1, temp2) * 1
    orImage = logical_or(temp1, temp2) * 1
    fraction = float(sum(andImage)) / float(sum(orImage))

    return fraction

def overlap(contour1, contour2):
    """Contour overlap"""

    #temp1 = zeros((1000, 1000), dtype=int8)
    #temp2 = zeros((1000, 1000), dtype=int8)

    binaryImage1 = contour1.binaryImage
    binaryImage2 = contour2.binaryImage

    boundingBox1 = contour1.get2DBoundingBox()
    boundingBox2 = contour2.get2DBoundingBox()

    minX = min(boundingBox1[0][0], boundingBox2[0][0])
    minY = min(boundingBox1[0][1], boundingBox2[0][1])
    maxX = max(boundingBox1[1][0], boundingBox2[1][0])
    maxY = max(boundingBox1[1][1], boundingBox2[1][1])

    temp1 = zeros((maxX-minX+1, maxY-minY+1))
    temp2 = zeros((maxX-minX+1, maxY-minY+1))

    temp1[boundingBox1[0][0]-minX:boundingBox1[1][0]-minX+1,
          boundingBox1[0][1]-minY:boundingBox1[1][1]-minY+1] = binaryImage1

    temp2[boundingBox2[0][0]-minX:boundingBox2[1][0]-minX+1,
          boundingBox2[0][1]-minY:boundingBox2[1][1]-minY+1] = binaryImage2

    andImage = logical_and(temp1, temp2) * 1
    orImage = logical_or(temp1, temp2) * 1
    fraction = float(sum(andImage)) / float(sum(orImage))

    return fraction


# Copyright (c) Anand Patil, 2007

__docformat__ = 'reStructuredText'

import pymc as pm
from . import linalg_utils
import copy
import types
import numpy as np
from .gp_submodel import *
import warnings

from pymc import six
from pymc.six import print_
xrange = six.moves.xrange

from .Realization import Realization
from .Mean import Mean
from .Covariance import Covariance
from .GPutils import observe, regularize_array

__all__ = [
    'wrap_metropolis_for_gp_parents',
    'GPEvaluationGibbs',
    'GPParentAdaptiveMetropolis',
    'GPStepMethod',
    'GPEvaluationMetropolis',
    'MeshlessGPMetropolis']

class GPStepMethod(pm.NoStepper):

    @staticmethod
    def competence(stochastic):
        if isinstance(stochastic, GaussianProcess):
            return 1
        else:
            return 0

    def tune(self, verbose=0):
        return False

def wrap_metropolis_for_gp_parents(metro_class):
    """
    Wraps Metropolis step methods so they can handle extended parents of
    Gaussian processes.
    """
    class wrapper(metro_class):
        __doc__ = """A modified version of class %s that handles parents of Gaussian processes.
Docstring of class %s: \n\n%s""" % (metro_class.__name__, metro_class.__name__, metro_class.__doc__)

        def __init__(self, stochastic, *args, **kwds):

            self.metro_class.__init__(self, stochastic, *args, **kwds)

            mb = set(self.markov_blanket)
            for c in list(self.children):
                if isinstance(c, GaussianProcess):
                    self.children |= c.extended_children
                    mb |= c.extended_children
            self.markov_blanket = list(mb)

            # Remove f from the set that will be used to compute
            # logp_plus_loglike.
            self.markov_blanket_no_f = set(
                filter(
                    lambda x: not isinstance(
                        x,
                        GaussianProcess),
                    self.markov_blanket))
            self.fs = filter(
                lambda x: isinstance(
                    x,
                    GaussianProcess),
                self.markov_blanket)
            self.fr_checks = [f.submodel.fr_check for f in self.fs]

        def get_logp_plus_loglike(self):
            return pm.logp_of_set(self.markov_blanket_no_f)
        logp_plus_loglike = property(get_logp_plus_loglike)

        def propose(self):
            self.metro_class.propose(self)
            try:
                # First make sure none of the stochastics handled by
                # metro_method forbid their current values.
                for s in self.stochastics:
                    s.logp
                # Then make sure the covariances are all still full-rank on the
                # observation locations.
                for frc in self.fr_checks:
                    frc.logp
                for f in self.fs:
                    f.rand()
                self.f_proposed = True
            except pm.ZeroProbability:
                self.f_proposed = False

        def reject(self):
            self.metro_class.reject(self)
            if self.f_proposed:
                for f in self.fs:
                    f.revert()

        @staticmethod
        def competence(stochastic, metro_class=metro_class):
            if any([isinstance(child, GaussianProcess)
                    for child in stochastic.extended_children]):
                return metro_class.competence(stochastic) + .01
            else:
                return 0

    wrapper.__name__ = 'GPParent%s' % metro_class.__name__
    wrapper.metro_class = metro_class

    return wrapper

# Wrap all registered Metropolis step methods to use GP parents.
new_sm_dict = {}
filtered_registry = [
    sm for sm in pm.StepMethodRegistry if issubclass(
        sm,
        pm.Metropolis)]
for sm in filtered_registry:
    wrapped_method = wrap_metropolis_for_gp_parents(sm)
    new_sm_dict[wrapped_method.__name__] = wrapped_method
GPParentAdaptiveMetropolis = wrap_metropolis_for_gp_parents(
    pm.AdaptiveMetropolis)
__all__ += new_sm_dict.keys()
locals().update(new_sm_dict)

class MeshlessGPMetropolis(pm.Metropolis):

    def __init__(self, gp):
        pm.Metropolis.__init__(
            self,
            gp,
            proposal_distribution='Prior',
            check_before_accepting=False)

    def propose(self):
        self.stochastic.rand()

    @staticmethod
    def competence(stochastic):
        if isinstance(stochastic, GaussianProcess):
            if len(stochastic.submodel.mesh) == 0:
                return 3
            else:
                return 0
        else:
            return 0

class _GPEvaluationMetropolis(pm.Metropolis):

    """
    Updates a GP evaluation, the 'f_eval' attribute of a GP submodel.
    The stationary distribution of the assymetric proposal is equal
    to the prior distribution, an attempt to minimize jumps to values
    forbidden by the prior.
    """

    def __init__(self, stochastic, proposal_sd=1, **kwds):
        pm.Metropolis.__init__(
            self,
            stochastic,
            proposal_sd=proposal_sd,
            **kwds)

    def propose(self):
        sig = pm.utils.value(self.stochastic.parents['sig'])
        mu = pm.utils.value(self.stochastic.parents['mu'])

        delta = pm.rmv_normal_chol(0 * mu, sig)

        beta = np.minimum(1, self.proposal_sd * self.adaptive_scale_factor)
        bsig = beta * sig
        sb2 = np.sqrt(1 - beta ** 2)
        self.stochastic.value = (
            self.stochastic.value - mu) * sb2 + beta * delta + mu
        xp, x = self.stochastic.value, self.stochastic.last_value
        self._hastings_factor = pm.mv_normal_chol_like(
            x,
            (xp - mu) * sb2 + mu,
            bsig) - pm.mv_normal_chol_like(xp,
                                           (x - mu) * sb2 + mu,
                                           bsig)

        # self.stochastic.value = self.stochastic.value + self.adaptive_scale_factor*self.proposal_sd*delta
        # self._hastings_factor = 0

    def hastings_factor(self):
        return self._hastings_factor

    @staticmethod
    def competence(stochastic):
        if isinstance(stochastic, GPEvaluation):
            return 3
        else:
            return 0

GPEvaluationMetropolis = wrap_metropolis_for_gp_parents(
    _GPEvaluationMetropolis)

class GPEvaluationGibbs(pm.Metropolis):

    """
    Updates a GP evaluation f_eval. Assumes the only children of f_eval
    are as distributed follows:

    eps_p_f ~ Normal(f_eval[ti], 1./V)

    or

    eps_p_f ~ Normal(f_eval, 1./V)

    if ti is None.
    """

    def __init__(self, submod, V, eps_p_f, ti=None, tally=True, verbose=0):

        self.f_eval = submod.f_eval
        self.f = submod.f
        pm.StepMethod.__init__(self, [self.f, self.f_eval], tally=tally)

        self.children_no_data = copy.copy(self.children)
        if isinstance(eps_p_f, pm.Variable):
            self.children_no_data.remove(eps_p_f)
            self.eps_p_f = eps_p_f
        else:
            for epf in eps_p_f:
                self.children_no_data.remove(epf)
            self.eps_p_f = pm.Lambda(
                'eps_p_f',
                lambda e=eps_p_f: np.hstack(
                    e),
                trace=False)

        self.V = pm.Lambda(
            '%s_vect' % V.__name__,
            lambda V=V: np.resize(V,
                                  len(submod.mesh)))
        self.C_eval = submod.C_eval
        self.M_eval = submod.M_eval
        self.S_eval = submod.S_eval

        M_eval_shape = pm.utils.value(self.M_eval).shape
        C_eval_shape = pm.utils.value(self.C_eval).shape
        self.ti = ti or np.arange(M_eval_shape[0])

        # Work arrays
        self.scratch1 = np.asmatrix(np.empty(C_eval_shape, order='F'))
        self.scratch2 = np.asmatrix(np.empty(C_eval_shape, order='F'))
        self.scratch3 = np.empty(M_eval_shape)

        # Initialize hidden attributes
        self.accepted = 0.
        self.rejected = 0.
        self._state = ['rejected', 'accepted', 'proposal_distribution']
        self._tuning_info = []
        self.proposal_distribution = None

        self.verbose = verbose

    def get_logp(self):
        return 0.
    logp = property(get_logp)

    def get_loglike(self):
        return pm.utils.logp_of_set(self.children_no_data)
    loglike = property(get_loglike)

    def get_logp_plus_loglike(self):
        return self.get_loglike()
    logp_plus_loglike = property(get_logp_plus_loglike)

    def reject(self):
        self.rejected += 1
        if self.verbose:
            print_(self._id + ' rejecting')
        # Revert the field evaluation and the rest of the field.
        self.f_eval.revert()
        self.f.revert()

    def tune(self, verbose=0):
        return False

    def propose(self):

        if self.verbose:
            print_(self._id + ' proposing')

        fc = pm.gp.fast_matrix_copy

        eps_p_f = pm.utils.value(self.eps_p_f)
        f = pm.utils.value(self.f_eval)
        for i in xrange(len(self.scratch3)):
            self.scratch3[i] = np.sum(eps_p_f[self.ti[i]] - f[i])

        # Compute Cholesky factor of covariance of eps_p_f, C(x,x) + V
        C_eval_value = pm.utils.value(self.C_eval)
        C_eval_shape = C_eval_value.shape

        # Get the Cholesky factor of C_eval, plus the nugget.
        # I don't think you can use S_eval for speed, unfortunately.
        in_chol = fc(C_eval_value, self.scratch1)

        v_val = pm.utils.value(self.V)
        for i in xrange(pm.utils.value(C_eval_shape)[0]):
            in_chol[i, i] += v_val[i] / np.alen(self.ti[i])

        info = pm.gp.linalg_utils.dpotrf_wrap(in_chol)
        if info > 0:
            raise np.linalg.LinAlgError

        # Compute covariance of f conditional on eps_p_f.
        offdiag = fc(C_eval_value, self.scratch2)
        offdiag = pm.gp.trisolve(
            in_chol,
            offdiag,
            uplo='U',
            transa='T',
            inplace=True)

        C_step = offdiag.T * offdiag
        C_step *= -1
        C_step += C_eval_value

        # Compute mean of f conditional on eps_p_f.
        for i in xrange(len(self.scratch3)):
            self.scratch3[i] = np.mean(eps_p_f[self.ti[i]])
        m_step = pm.utils.value(
            self.M_eval) + np.dot(offdiag.T,
                                  pm.gp.trisolve(in_chol,
                                                 (self.scratch3 -
                                                  self.M_eval.value),
                                                 uplo='U',
                                                 transa='T')).view(np.ndarray).ravel()

        sig_step = C_step
        info = pm.gp.linalg_utils.dpotrf_wrap(C_step.T)
        if info > 0:
            warnings.warn(
                'Full conditional covariance was not positive definite.')
            return

        # Update value of f.
        self.f_eval.value = m_step + np.dot(
            sig_step,
            np.random.normal(
                size=sig_step.shape[
                    1])).view(
                        np.ndarray).ravel(
        )
        # Propose the rest of the field from its conditional prior.
        self.f.rand()


import numpy as np
import os
import copy

from Replicates import *
from utils import *

import matplotlib as mat
import matplotlib.pyplot as plt

try:
    from mpi4py import MPI      # Import MPI if available
except:
    pass

def ReachSteadyStateAndRescale(kmc_template, scale_parent_fldr, n_runs = 16, n_batches = 1000,
                                acf_cut = 0.05, include_stiff_reduc = True, max_events = int(1e3),
                                max_iterations = 30, ss_inc = 1.0, n_samples = 100, parallel_mode = 'Squidward',
                                ACF_tol = 0.08, rate_tol = 0.05):

    '''
    Handles rate rescaling and continuation of KMC runs

    :param kmc_template:            kmc_traj object with information about the physical system

    :param scale_parent_fldr:       Working directory

    :param n_runs:                  Number of trajectories to run, also the number of processors

    :param include_stiff_reduc:     True to allow for scaledown, False to turn this feature off

    :param max_events:              Maximum number of events for the first iteration

    :param max_iterations:          Maximum number of iterations to use

    :param ss_inc:                  Factor to scale the final time by if you have not yet reached steady state

    :param n_samples:               Number of time points to sample for each trajectory

    :param parallel_mode:           MPI, Squidward, or Farber
    '''

    if parallel_mode == 'MPI':
        try:
            COMM = MPI.COMM_WORLD
            COMM.Barrier()
        except:
            raise NameError('mpi4py dependency has not been imported.')

    prev_batch = Replicates()       # Set this if the starting iteration is not 1
    initial_states = None

    # Placeholder variables
    if not os.path.exists(scale_parent_fldr):
        os.makedirs(scale_parent_fldr)
    ClearFolderContents(scale_parent_fldr)
    SDF_vec = None        # scaledown factors for each iteration

    # Convergence variables
    is_steady_state = False
    iteration = 1

    scale_final_time = ss_inc

    while not is_steady_state and iteration <= max_iterations:

        # Make folder for iteration
        iter_fldr = os.path.join(scale_parent_fldr, 'Iteration_' + str(iteration))
        if not os.path.exists(iter_fldr):
            os.makedirs(iter_fldr)

        # Create object for batch
        cur_batch = Replicates()
        cur_batch.ParentFolder = iter_fldr
        cur_batch.n_trajectories = n_runs
        cur_batch.N_batches = n_batches
        cur_batch.Set_kmc_template(kmc_template)        # Set template KMC trajectory

        if iteration == 1:              # Sample on events, because we do not know the time scales

            # Set sampling parameters
            cur_batch.runtemplate.simin.MaxStep = max_events
            cur_batch.runtemplate.simin.SimTime_Max = 'inf'
            cur_batch.runtemplate.simin.WallTime_Max = 'inf'
            cur_batch.runtemplate.simin.restart = False

            cur_batch.runtemplate.simin.procstat = ['event', np.max( [max_events / n_samples, 1] ) ]
            cur_batch.runtemplate.simin.specnum = ['event', np.max( [max_events / n_samples, 1] ) ]
            cur_batch.runtemplate.simin.hist = ['event', np.max( [max_events * (n_samples-1) / n_samples, 1] )]       # only record the initial and final states

            SDF_vec = np.ones( cur_batch.runtemplate.mechin.get_num_rxns() )         # Initialize scaledown factors

        elif iteration > 1:             # Time sampling

            # Change sampling
            cur_batch.runtemplate.simin.MaxStep = 'inf'
            cur_batch.runtemplate.simin.WallTime_Max = 'inf'
            cur_batch.runtemplate.simin.restart = False
            cur_batch.runtemplate.simin.SimTime_Max = prev_batch.t_vec[-1] * scale_final_time
            cur_batch.runtemplate.simin.SimTime_Max = float('{0:.3E} \t'.format( cur_batch.runtemplate.simin.SimTime_Max ))     # round to 4 significant figures
            cur_batch.runtemplate.simin.procstat = ['time', cur_batch.runtemplate.simin.SimTime_Max / n_samples]
            cur_batch.runtemplate.simin.specnum = ['time', cur_batch.runtemplate.simin.SimTime_Max / n_samples]
            cur_batch.runtemplate.simin.hist = ['time', cur_batch.runtemplate.simin.SimTime_Max ]

            # Adjust pre-exponential factors based on the stiffness assessment of the previous iteration
            if include_stiff_reduc:
                cur_batch.runtemplate.AdjustPreExponentials(SDF_vec)

            # Use continuation
            initial_states = prev_batch.History_final_snaps

        # Run jobs and read output
        if parallel_mode == 'MPI':
            if COMM.rank == 0:
                cur_batch.BuildJobFiles(init_states = initial_states)

            # Collect whatever has to be done in a list. Here we'll just collect a list of
            # numbers. Only the first rank has to do this.
            if COMM.rank == 0:
                jobs = cur_batch.run_dirs
                jobs = [jobs[_i::COMM.size] for _i in range(COMM.size)]             # Split into however many cores are available.
            else:
                jobs = None

            jobs = COMM.scatter(jobs, root=0)           # Scatter jobs across cores.

            # Now each rank just does its jobs and collects everything in a results list.
            # Make sure to not use super big objects in there as they will be pickled to be
            # exchanged over MPI.
            for job in jobs:
                cur_batch.runtemplate.Path = job
                cur_batch.runtemplate.Run_sim()

        else:
            cur_batch.BuildJobFiles(init_states = initial_states)
            cur_batch.RunAllTrajectories_JobArray(server = parallel_mode, job_name = 'Iteration_' + str(iteration) )

        cur_batch.ReadMultipleRuns()

        if iteration == 1:
            cum_batch = copy.deepcopy(cur_batch)
        else:
            cum_batch = append_replicates(prev_batch, cur_batch)         # combine with previous data

        # Test steady-state
        cum_batch.AverageRuns()
        acf_data = cum_batch.Compute_rate()

        print '\nIteration ' + str(iteration)
        print 'Batches per trajectory: ' + str(cum_batch.Nbpt)
        print 'Batch length (s): ' + str(cum_batch.batch_length)
        print 'Rate: ' + str(cum_batch.rate)
        print 'Rate confidence interval: ' + str(cum_batch.rate_CI)
        print 'Autocorrelation: ' + str(cum_batch.ACF)
        print 'Autocorrelation confidence: ' + str(cum_batch.ACF_CI)

        # Test if autocorrelation function has converged

        if cum_batch.ACF is None:
            decorrelated = False
        else:
            decorrelated = ( cum_batch.ACF + cum_batch.ACF_CI < ACF_tol)

        # Test if rate is computed with sufficient accuracy
        if cum_batch.rate == 0:
            rate_accurate = False
        else:
            rate_accurate = (cum_batch.rate_CI / cum_batch.rate < rate_tol)

        print 'Decorrelated? ' + str(decorrelated)
        print 'Rate accurate? ' + str(rate_accurate)
        print '\n'

        is_steady_state = decorrelated and rate_accurate

        # Record information about the iteration
        cum_batch.runAvg.PlotGasSpecVsTime()
        cum_batch.runAvg.PlotSurfSpecVsTime()

        cur_batch.AverageRuns()
        cur_batch.runAvg.PlotElemStepFreqs()
        scaledown_data = ProcessStepFreqs(cur_batch.runAvg)         # compute change in scaledown factors based on simulation result
        delta_sdf = scaledown_data['delta_sdf']

        # Update scaledown factors
        for ind in range(len(SDF_vec)):
            SDF_vec[ind] = SDF_vec[ind] * delta_sdf[ind]

        scale_final_time = np.max( [1.0/np.min(delta_sdf), ss_inc] )

        prev_batch = copy.deepcopy(cum_batch)
        iteration += 1

    return cum_batch

def ProcessStepFreqs(run, stiff_cut = 100.0, delta = 0.05, equilib_cut = 0.1):        # Change to allow for irreversible reactions

    '''
    Takes an average KMC trajectory and assesses the reaction frequencies to identify fast reactions
    Process KMC output and determine how to further scale down reactions
    Uses algorithm from A. Chatterjee, A.F. Voter, Accurate acceleration of kinetic Monte Carlo simulations through the modification of rate constants, J. Chem. Phys. 132 (2010) 194101.
    '''

    delta_sdf = np.ones( run.mechin.get_num_rxns() )    # initialize the marginal scaledown factors
    rxn_speeds = []

    # data analysis
    freqs = run.procstatout.events[-1,:]
    fwd_freqs = freqs[0::2]
    bwd_freqs = freqs[1::2]
    net_freqs = fwd_freqs - bwd_freqs
    tot_freqs = fwd_freqs + bwd_freqs

    fast_rxns = []
    slow_rxns = []
    for i in range(len(tot_freqs)):
        if tot_freqs[i] == 0:
            slow_rxns.append(i)
            rxn_speeds.append('slow')
        else:
            PE = float(net_freqs[i]) / tot_freqs[i]
            if np.abs(PE) < equilib_cut:
                fast_rxns.append(i)
                rxn_speeds.append('fast')
            else:
                slow_rxns.append(i)
                rxn_speeds.append('slow')

    # Find slow scale rate
    slow_freqs = [1.0]      # put an extra 1 in case no slow reactions occur
    for i in slow_rxns:
        slow_freqs.append(tot_freqs[i])
    slow_scale = np.max(slow_freqs)

    # Adjust fast reactions closer to the slow scale
    for i in fast_rxns:
        N_f = tot_freqs[i] / float(slow_scale)              # number of fast events per rare event
        #alpha_UB = N_f * delta / np.log(1 / delta) + 1             # Chatterjee formula

        #delta_sdf[i] = np.min([1.0, np.max([stiff_cut / N_f, 1. / alpha_UB ]) ])
        delta_sdf[i] = np.min([1.0, stiff_cut / N_f ])

    return {'delta_sdf': delta_sdf, 'rxn_speeds': rxn_speeds, 'tot': tot_freqs, 'net': net_freqs}

def ReadScaledown(RunPath, fldrs_cut = None, product = None, n_batches = 1000):

    '''
    Read a scaledown that has already been run
    '''

    # Prepare data to graph
    batch_lengths = []
    rates_vec = []
    rates_vec_ci = []
    acf_vec = []
    acf_vec_ci = []

    # Count the iterations
    n_folders = len(os.listdir(RunPath))
    if not fldrs_cut is None:

        n_folders = fldrs_cut

    print str(n_folders) + ' iterations found'

    cum_batch = None
    for ind in range(1,n_folders+1):

        x = Replicates()
        x.ParentFolder = os.path.join(RunPath, 'Iteration_' + str(ind))
        x.ReadMultipleRuns()

        if ind == 1:
            cum_batch = x
        else:
            cum_batch = append_replicates(cum_batch, x)

        cum_batch.N_batches = n_batches
        cum_batch.gas_product = product
        acf_data = cum_batch.Compute_rate()
        print 'Iteration ' + str(ind)
        print 'Batches per trajectory: ' + str(cum_batch.Nbpt)
        print 'Batch length (s): ' + str(cum_batch.batch_length)
        print 'Rate: ' + str(cum_batch.rate)
        print 'Rate confidence interval: ' + str(cum_batch.rate_CI)
        print 'Autocorrelation: ' + str(cum_batch.ACF)
        print 'Autocorrelation confidence: ' + str(cum_batch.ACF_CI)

        print '\n'

        batch_lengths.append(cum_batch.batch_length)
        rates_vec.append(cum_batch.rate)
        rates_vec_ci.append(cum_batch.rate_CI)
        acf_vec.append(cum_batch.ACF)
        acf_vec_ci.append(cum_batch.ACF_CI)

    print batch_lengths
    print rates_vec
    print rates_vec_ci
    print acf_vec
    print acf_vec_ci

    return cum_batch