#! /usr/bin/env python
# -*- coding: utf-8 -*-

import os
import numpy as np
import pesfit as pf
import time
from hdfio import dict_io as io
import argparse
import multiprocessing as mp

n_cpu = mp.cpu_count()

# Definitions of command line interaction
parser = argparse.ArgumentParser(description='Input arguments')
parser.add_argument('-nb', '--nband', metavar='nband', nargs='?', type=int, help='Number of bands in fitting model, needs an integer between 1 and 14')
parser.add_argument('-ns', '--nspectra', metavar='nspectra', nargs='?', type=int, help='Number of spectra to fit, needs an integer larger than 1')
parser.add_argument('-ds', '--datasource', metavar='datasource', nargs='?', type=str, help='Name of the data source for band fitting')
parser.add_argument('-ofs', '--eoffset', metavar='eoffset', nargs='*', help='Global energy offset for each energy band')
parser.add_argument('-op', '--operation', metavar='operation', nargs='?', type=str, help='What computing method to run the benchmark program with')
parser.add_argument('-bk', '--backend', metavar='backend', nargs='?', type=str, help='Backend software package used for execution (in paralleli or in sequence)')
parser.add_argument('-nw', '--nworker', metavar='nworker', nargs='?', type=int, help='Number of workers to spawn')
parser.add_argument('-cs', '--chunksize', metavar='chunksize', nargs='?', type=int, help='Chunk size of tasks assigned to each worker')
parser.add_argument('-tc', '--timecount', metavar='timcount', nargs='?', type=bool, help='Whether to include time profiling')
parser.add_argument('-pkl', '--pickle', metavar='pickle', nargs='?', type=bool, help='Whether to pickle the detailed results to file (may be large)')
parser.add_argument('-persin', '--persistent_init', metavar='persistent_init', nargs='?', type=bool, help='Initialization include persistent settings')
parser.add_argument('-varin', '--varying_init', metavar='varying_init', nargs='?', type=str, help='Initialization including varying settings')
parser.add_argument('-jittin', '--jitter_init', metavar='jitter_init', nargs='?', type=bool, help='Add jitter to initialization for better fits')
parser.set_defaults(nband=2, nspectra=10, datasource='preprocessed', eoffset=[0.], operation='sequential', backend='multiprocessing', nworker=n_cpu, chunksize=1, timecount=True, pickle=False, persistent_init=True, varying_init='recon', jitter_init=False)
cli_args = parser.parse_args()

# Sequential fitting of photoemission data patch around the M point of WSe2
## Number of bands to fit
NBAND = cli_args.nband
if NBAND > 14 or NBAND < 1:
    raise ValueError('The number of bands to reconstruct is within [1, 14] for WSe2.')
## Data source used for band fitting, ('symmetrized', 'preprocessed')
DATASOURCE = cli_args.datasource
## Number of line spectra to fit
NSPECTRA = cli_args.nspectra
## Global energy offset for fitting initialization of each band
argofs = list(map(float, cli_args.eoffset))
nargofs = len(argofs)
if nargofs == 1:
    argofs = argofs*NBAND
    allofs = argofs + [0]*(14-NBAND)
elif nargofs == NBAND:
    allofs = argofs + [0]*(14-NBAND)
EOFFSET = np.asarray([allofs]).T
## Option for computing method ('sequential' or 'parallel')
OPERATION = cli_args.operation
## Backend software package for execution
BACKEND = cli_args.backend
## Number of workers to use in running the benchmark
NWORKER = cli_args.nworker
## Number of tasks assigned to each worker
CHUNKSIZE = cli_args.chunksize
## Option to enable code profiling
TIMECOUNT = cli_args.timecount
## Option to pickle all detailed fitting results to a file (may be large)
PICKLE = cli_args.pickle
## Option to introduce persistent initial conditions
PERSISTENT_INIT = cli_args.persistent_init
## Specification of spectrum-dependent initial conditions ('theory' or 'recon')
VARYING_INIT = cli_args.varying_init
## Option to apply jittering to initializations to obtain better fitting results
JITTER_INIT = cli_args.jitter_init
n=2

# Photoemission band mapping data
data_dir = r'../data/WSe2'
pes_fname = r'/pes/mpoint/mpoint_{}.h5'.format(DATASOURCE)
pes_path = data_dir + pes_fname
pes_data = io.h5_to_dict(pes_path)

if VARYING_INIT == 'theory':
    # Theoretical calculations interpolated to the same momentum grid (as one type of initialization)
    theo_fname = r'/theory/mpoint/mpoint_LDA_50.h5'
    theo_path = data_dir + theo_fname
    theo_data = io.h5_to_dict(theo_path)['bands']
    print(theo_data.shape)
    inits_vary = theo_data[:,::n,::n]

elif VARYING_INIT == 'recon':
    # Reconstructed photoemission band dispersion (as another type of initialization)
    recon_fname = r'/recon/mpoint/mpoint_LDA_recon.h5'
    recon_path = data_dir + recon_fname
    recon_data = io.h5_to_dict(recon_path)['recon']
    inits_vary = recon_data

# Setting initial conditions persistent throughout the fitting. These fixed constraints are tested for fitting the band dispersion nearby the M' point of WSe2 and are required to make the fitting relatively stable!
if PERSISTENT_INIT:
    vardict = {}
    preftext = 'lp'
    lp_prefixes = [preftext+str(i)+'_' for i in range(1, NBAND+1)]

    ## Case of 2 bands near M point
    vardict['02'] = [{'lp1_':{'amplitude':dict(value=0.2, min=0, max=2, vary=True),
                            'sigma':dict(value=0.1, min=0.05, max=0.5, vary=False),
                            'gamma':dict(value=0.05, min=0, max=1, vary=True)}},
                
                    {'lp2_':{'amplitude':dict(value=0.5, min=0, max=2, vary=True),
                            'sigma':dict(value=0.1, min=0.05, max=0.5, vary=False),
                            'gamma':dict(value=0.05, min=0, max=1, vary=True)}}]

    ## Case of 4 bands near M point
    vardict['04'] = [{'lp3_':{'amplitude':dict(value=0.2, min=0, max=2, vary=1.),
                    'sigma':dict(value=0.1, min=0.05, max=2, vary=False),
                    'gamma':dict(value=0.05, min=0, max=2, vary=True)}},
            
                    {'lp4_':{'amplitude':dict(value=0.2, min=0, max=2, vary=True),
                        'sigma':dict(value=0.1, min=0.05, max=2, vary=False),
                        'gamma':dict(value=0.05, min=0, max=2, vary=True)}},
                        
                    {'bg_':{'amplitude':dict(value=0.8, min=0, max=2, vary=True),
                            'sigma':dict(value=2, min=0, max=10, vary=True),
                            'center':dict(value=-5, min=-20, max=-3, vary=True)}}]
    vardict['04'] = vardict['02'] + vardict['04']

    ## Case of 8 bands near M point
    amplitudes = pf.fitter.init_generator(lpnames=lp_prefixes[4:8], parname='amplitude',
                                        varkeys=['value', 'min', 'max', 'vary'],
                                        parvals=[[0.5, 0, 2, True] for i in range(8-4)])
    sigmas = pf.fitter.init_generator(lpnames=lp_prefixes[4:8], parname='sigma',
                                    varkeys=['value', 'min', 'max', 'vary'],
                                    parvals=[[0.1, 0.05, 0.2, False] for i in range(8-4)])
    gammas = pf.fitter.init_generator(lpnames=lp_prefixes[4:8], parname='gamma',
                                    varkeys=['value', 'min', 'max', 'vary'],
                                    parvals=[[0.05, 0, 2, True] for i in range(8-4)])
    vardict['08'] = amplitudes + sigmas + gammas
    vardict['08'] = vardict['04'] + vardict['08']

    ## Case of 14 bands near M point
    amplitudes = pf.fitter.init_generator(lpnames=lp_prefixes[4:14], parname='amplitude',
                                        varkeys=['value', 'min', 'max', 'vary'],
                                        parvals=[[0.5, 0, 2, True] for i in range(14-4)])
    sigmas = pf.fitter.init_generator(lpnames=lp_prefixes[4:14], parname='sigma',
                                    varkeys=['value', 'min', 'max', 'vary'],
                                    parvals=[[0.1, 0.05, 0.2, False] for i in range(14-4)])
    gammas = pf.fitter.init_generator(lpnames=lp_prefixes[4:14], parname='gamma', 
                                    varkeys=['value', 'min', 'max', 'vary'],
                                    parvals=[[0.05, 0, 2, True] for i in range(14-4)])
    vardict['14'] = amplitudes + sigmas + gammas
    vardict['14'] = vardict['08'] + vardict['14']

    ## Other number of bands
    if NBAND not in [2, 4, 8, 14]:
        vals = vardict['14']
        vardict_other = {}
        lp_exclude = [preftext+str(i)+'_' for i in range(NBAND+1, 15)]
        for inits in vardict['14']:
            for lpn, lpv in inits.items():
                if lpn in lp_exclude:
                    vals.pop(lpn)
        vardict_other[str(NBAND).zfill(2)] = vals

# Set up and run the fitting routine
## Specify the sets of initialization to apply to the fitting
if PERSISTENT_INIT:
    bandkey = str(int(NBAND)).zfill(2)
    try:
        inits_persist = vardict[bandkey]
    except:
        inits_persist = vardict_other[bandkey]
else:
    inits_persist = None

## Select energy axis data range used for fitting
if NBAND == 2:
    en_range = slice(20, 120)
elif NBAND == 4:
    en_range = slice(20, 180)
elif NBAND == 8:
    en_range = slice(20, 250)
elif NBAND == 14:
    en_range = slice(20, 490)
else:
    en_range = slice(20, 490)

## Run the fitting benchmark
pesdata_shape = pes_data['V'][::n,::n,:].shape
maxspectra = pesdata_shape[0] * pesdata_shape[1]
nspec = min([NSPECTRA, maxspectra])

if OPERATION == 'sequential':    
    kfit = pf.fitter.PatchFitter(peaks={'Voigt':NBAND}, xdata=pes_data['E'], ydata=pes_data['V'][::n,::n,:], preftext=preftext)

    kfit.set_inits(inits_dict=inits_persist, band_inits=inits_vary, drange=en_range, offset=EOFFSET)

    tstart = time.perf_counter()
    kfit.sequential_fit(pbar=True, pbenv='classic', jitter_init=JITTER_INIT, shifts=np.arange(-0.08, 0.09, 0.01), nspec=nspec, ynorm=True)
    tstop = time.perf_counter()

    if TIMECOUNT:
        tdiff =  tstop - tstart
        print('Fitting took {} seconds'.format(tdiff))

    ## Save the fitting outcome
    out_dir = r'../data/WSe2/benchmark'
    if not os.path.exists(out_dir):
        os.makedirs(out_dir)

    fn_text = '/mpoint_{}_seqfit_nband={}_ofs={}_{}'.format(DATASOURCE, NBAND, argofs, VARYING_INIT)
    kfit.save_data(fdir=out_dir, fname=fn_text + '.h5')
    if PICKLE:
        pf.utils.pickle_obj(out_dir + fn_text + '.pickle', kfit.fitres)

elif OPERATION == 'parallel':
    if __name__ == '__main__':
        kfit = pf.fitter.DistributedFitter(peaks={'Voigt':NBAND}, xdata=pes_data['E'], ydata=pes_data['V'][::n,::n,:], drange=en_range, nfitter=nspec)

        kfit.set_inits(inits_dict=inits_persist, band_inits=inits_vary, offset=EOFFSET)

        if CHUNKSIZE > 0:
            tstart = time.perf_counter()
            kfit.parallel_fit(jitter_init=JITTER_INIT, shifts=np.arange(-0.08, 0.09, 0.01), nfitter=nspec, backend=BACKEND, scheduler='processes', num_workers=NWORKER, chunksize=CHUNKSIZE, pbar=True)
        else:
            tstart = time.perf_counter()
            kfit.parallel_fit(jitter_init=JITTER_INIT, shifts=np.arange(-0.08, 0.09, 0.01), nfitter=nspec, backend=BACKEND, scheduler='processes', num_workers=NWORKER, pbar=True)
        tstop = time.perf_counter()

        if TIMECOUNT:
            tdiff =  tstop - tstart
            print('Fitting took {} seconds'.format(tdiff))

        ## Save the fitting outcome
        out_dir = r'../data/WSe2/benchmark'
        if not os.path.exists(out_dir):
            os.makedirs(out_dir)

        fn_text = '/mpoint_{}_parafit_nband={}_ofs={}_{}'.format(DATASOURCE, NBAND, argofs, VARYING_INIT)
        kfit.save_data(fdir=out_dir, fname=fn_text + '.h5')
        if PICKLE:
            pf.utils.pickle_obj(out_dir + fn_text + '.pickle', kfit.fitres)

else:
    raise NotImplementedError