#!/usr/bin/env python

import sys
import argparse
import configparser
import os
import numpy as np 
import time
import pylab as plt
from multiprocessing import Pool, cpu_count
from cf_tools import Multipoles
from galaxy_bao_fitter import Data, Chi2

parser = argparse.ArgumentParser()
parser.add_argument('ini', help='Configuration file for fit')
parser.add_argument('--data-name', help='Input correlation function')
parser.add_argument('--cov-name', help='Input covariance matrix')
parser.add_argument('--output', help='Root for outputs')
parser.add_argument('--cosmology', help='Name of fiducial cosmology')
parser.add_argument('--rmin', help='Minimum separation', 
                    type=float)
parser.add_argument('--rmax', help='Maximum separation', 
                    type=float)
parser.add_argument('--cov-scalefactor', type=float, help='Covariance matrix scale factor')

args = parser.parse_args()
config = configparser.ConfigParser()
config.read(args.ini)

#-- Read several options from config file (or from args if provided)
options = {}
for k in config['options']:
    options[k] = config['options'][k] == 'True'
options['ell_max'] = int(config['options']['ell_max'])
if 'broadband' in config:
    options['fit_broadband'] = True
    options['bb_min'] = int(config['broadband']['bb_min'])
    options['bb_max'] = int(config['broadband']['bb_max'])
else:
    options['fit_broadband'] = False
options['rmin'] = float(config['cuts']['rmin']) if args.rmin is None else args.rmin
options['rmax'] = float(config['cuts']['rmax']) if args.rmax is None else args.rmax
options['data_name'] = config['data']['name'] if args.data_name is None else args.data_name
options['cov_name'] = config['covariance']['name'] if args.cov_name is None else args.cov_name
options['cov_scale_factor'] = float(config['covariance']['scalefactor']) if args.cov_scalefactor is None else args.cov_scalefactor
options['output_root'] = config['data']['output'] if args.output is None else args.output


print('== Galaxy BAO fitter (by Julian Bautista) ==') 
print('')


##-- Read input multipoles 
print('Reading data and covariance matrix')
m = Multipoles(options['data_name'])
if options['ell_max'] == 2:
    m.hexa = None
elif options['ell_max'] == 0:
    m.quad = None

#-- Read covariance matrix
m.read_cov(options['cov_name'])
m.coss *= options['cov_scale_factor']
nmocks = int(config['covariance']['nmocks'])

#-- Define output name
rmin = options['rmin']
rmax = options['rmax']
output = ''.join([
        options['output_root'],
        f'-rmin{rmin:.1f}-rmax{rmax:.1f}',
        '-bb'*(options['fit_broadband']),
        f'-ellmax{options["ell_max"]}',
        f'-nopeak'*(options['no_peak'])  
        ])
output_base = os.path.basename(output)

#-- Read parameters 
parameters = {}
for par in config['parameters']:
    entries = config['parameters'][par].split()
    parameters[par] = {}
    parameters[par]['value'] = float(entries[0])
    parameters[par]['error'] = float(entries[1])
    parameters[par]['limit_low'] = float(entries[2]) if entries[2] != 'None' else None
    parameters[par]['limit_upp'] = float(entries[3]) if entries[3] != 'None' else None
    parameters[par]['fixed'] = (entries[4] == 'fixed')

#-- Define priors 
if config['priors']:
    for par in config['priors']:
        if par not in parameters: 
            print(f'Warning: parameter {par} not in parameters')
            continue
        entries = config['priors'][k].split()
        parameters[par]['prior_mean'] = float(entries[0]) if entries[0] != 'None' else None
        parameters[par]['prior_sigma'] = float(entries[1]) if entries[1] != 'None' else None

#-- Define data, cutting over rmin,rmax 
data = Data(m.r, m.mono, m.coss, 
            quad=m.quad, hexa=m.hexa, 
            nmocks=nmocks, rmax=rmax, rmin=rmin)

#-- Load model: BAO or RSD?
mod = config['model']
if mod['fit_type'] == 'BAO':
    from galaxy_bao_fitter import PowerSpectrum
    if 'pk_file' in mod:
        pk_file = mod['pk_file'] 
        model   = PowerSpectrum(pk_file=pk_file)
    else: 
        cosmology  = mod['cosmology'] if args.cosmology is None else args.cosmology
        z_eff      = float(mod['z_eff'])
        non_lin_pk = (mod['non_linear_pk'] == 'True')
        model = PowerSpectrum(z=z_eff, name=cosmology, non_linear=non_lin_pk)
elif mod['fit_type'] == 'RSD':
    from galaxy_rsd_fitter import CLPTGS
    model = CLPTGS(mod['xi_file'], mod['v12_file'], mod['s12_file'])


#-- Perform fit and export
chi2 = None
print('\nStarting chi2 minimization with iMinuit')
sec0 = time.time()
chi2 = Chi2(data=data, model=model, parameters=parameters, options=options)
chi2.fit() 
chi2.export_bestfit_parameters(output+'.pars')
#chi2.export_covariance(output+'.covpars')
sec1 = time.time()
tmin = (sec1-sec0)/60
print(f'Time elapsed during fit: {tmin:.2f} minutes')

#-- Get assymetric error bars
if 'minos' in config:
    sec2 = time.time()
    for par in config['minos']['parameters'].split():
        print('\nPerforming minos for', par)
        minos = chi2.mig.minos(var=par)
        fout = open(output+'.%s.minos'%par, 'w')
        val = minos[par]['min']
        low = minos[par]['lower']
        upp = minos[par]['upper']
        err = 0.5*(-low+upp)
        print(f'{par} = {val:.4f} +/- {err:.4f}') 
        print(par, val, low, upp, minos[par]['is_valid'], file=fout)
        fout.close()

    sec3 = time.time()
    tmin = (sec3-sec2)/60
    print(f'Time elapsed during minos: {tmin:.2f} minutes')

#-- Perform one dimensional scans of chi2
if 'scan1d' in config:
    sec3 = time.time()
    scan = config['scan1d']
    for par in scan: 
        print('\nPerforming scan of chi2 for', par)
        opt = scan[par].split()
        low = float(opt[0])
        upp = float(opt[1])
        nbin =  int(opt[2])
        par_scan, chi2_scan = \
            chi2.scan1d(par_name=par, par_nsteps=nbin, par_min=low, par_max=upp)

        #-- export chi2 scans
        fout = open(output+f'.{par}.scan1d', 'w')
        print(chi2.best_pars[par]['value'], chi2.chi2min, file=fout)
        for i in range(par_scan.size):
            print( par_scan[i], chi2_scan[i], file=fout)
        fout.close()
        print('chi2 scan exported to:', output+f'.{par}.scan1d')

    sec4 = time.time()
    tmin = (sec4-sec3)/60
    print(f'Time elapsed during scan1d: {tmin:.2f} minutes')


#-- Perform two dimensional scans of chi2
if 'scan2d' in config:
    sec4 = time.time()
    scan = config['scan2d']
    for pars in scan:
        par_names = pars.split() 
        print('\nPerforming scan of chi2 for', par_names[0], par_names[1])
        values = scan[pars].split()
        low = [float(values[0]), float(values[3])]
        upp = [float(values[1]), float(values[4])]
        nbin =  [int(values[2]),   int(values[5])]
        par_scan0, par_scan1, chi2_scan = \
            chi2.scan_2d(par_names=par_names, par_nsteps=nbin, par_min=low, par_max=upp)

        filename = output+'.%s.%s.scan2d'%(par_names[0], par_names[1])
        fout = open(filename, 'w')
        print( chi2.best_pars[par_names[0]]['value'], \
               chi2.best_pars[par_names[1]]['value'], \
               chi2.chi2min, file=fout)
        for i in range(par_scan0.size):
            for j in range(par_scan1.size):
                print(par_scan0[i], par_scan1[j], chi2_scan[i, j], file=fout)
        fout.close()
        print('chi2 scan exported to:', filename)

    sec5 = time.time()
    tmin = (sec5-sec4)/60
    print(f'Time elapsed during scan2d: {tmin:.2f} minutes')

#-- Plot best-fit model
if 'plot' in config:
    scale_r = int(config['plot']['scale_r'])
    label = r'$\chi^2/{\rm dof}$ = '
    label += f'{chi2.chi2min:.1f}/({chi2.ndata}-{chi2.npars}) = {chi2.rchi2min:.2f}' 
    fig = chi2.plot_bestfit(label=label, scale_r=scale_r)
    plt.legend(loc=0)
    plt.tight_layout()
    if 'save' in config['plot']:
        plt.savefig(config['plot']['save'], bbox_inches='tight')
    plt.show()

#-- Export best-fit model
if 'minuit' in config:
    chi2.export_model(output+'.model')

#- Perform MCMC
if 'mcmc' in config:
    print('\nPerforming MCMC chain')
    sec_0 = time.time()

    if config['mcmc']['sampler'] == 'zeus':
        import zeus
        sampling_function = zeus.sampler
        name_sampler = 'zeus'
    elif config['mcmc']['sampler'] == 'emcee':
        import emcee
        sampling_function = emcee.EnsembleSampler
        name_sampler = 'emcee'
    else:
        print("Need to setup sampler 'zeus' or 'emcee'")

    nsteps = int(config['mcmc']['nsteps'])
    nwalkers = int(config['mcmc']['nwalkers'])

    if chi2 is None:
        chi2 = Chi2(data=data, model=model, parameters=parameters, options=options)
        chi2.fit()

    def log_prob(p):
        return chi2.log_prob(p)

    pars_free = [par for par in parameters if not parameters[par]['fixed']]
    npars = len(pars_free)

    #-- Use limits to set the start point for random walkers
    print('\nSetting up starting point of walkers:')
    start = np.zeros((nwalkers, npars))
    for j, par in enumerate(pars_free):
        if par in chi2.best_pars and chi2.best_pars[par]['error']>0.:
            limit_low = parameters[par]['limit_low']
            limit_upp = parameters[par]['limit_upp']
            value = chi2.best_pars[par]['value']
            error = chi2.best_pars[par]['error']
            limit_low = np.max([limit_low, value-10*error])
            limit_upp = np.min([limit_upp, value+10*error])
            print('Randomly sampling for', par, 'between', limit_low, limit_upp )
            start[:, j] = np.random.rand(nwalkers)*(limit_upp-limit_low)+limit_low

    if config['mcmc']['use_pool'] == 'True':
        print(f'Using multiprocessing with {cpu_count()} cores')
        with Pool() as pool:
            sampler = sampling_function(nwalkers, npars, log_prob, pool=pool)
            sampler.run_mcmc(start, nsteps, progress=True)
    else:
        print('Using a single core')
        sampler = sampling_function(nwalkers, npars, log_prob)
        sampler.run_mcmc(start, nsteps, progress=True)

    if config['mcmc']['sampler'] == 'zeus':
        sampler.summary

    chain = sampler.get_chain(flat=True)
    np.savetxt(output+f".chain_{name_sampler}", chain)

    sec_1 = time.time()
    tmin = (sec_1-sec_0)/60
    print(f'Time elapsed during MCMC: {tmin:.2f} minutes')


