#! /usr/bin/python3


import numpy as np
import libMCMCMCFitting as lMC3F
import argparse # command-line interface


# This script is for PPMCMCMCFitting, which is
# Pump-probe Multiple Channels Markov Chain Monte-Carlo Fitting :)


if __name__ == '__main__':

    # Define and parse possible arguments of the script
    parser = argparse.ArgumentParser(description='A script to fit the pump-probe curves.') 

    parser.add_argument('-d', '--DataDescription', type=str, default=None, 
                    help="Data description file")  

    parser.add_argument('-c', '--ChannelsDescription', type=str, default=None, 
                    help="Channels description file")      

    parser.add_argument('-v', '--VariablesDescription', type=str, default=None, 
                    help="Variables description file")    

    parser.add_argument('-r', '--RegularizationDescription', type=str, default=None, 
                    help="Applied regularization description file")  

    parser.add_argument('--UseLocalOptimization', action='store_true', 
                    help="Switch from using the Global Optimization to the Local Optimization") 

    parser.add_argument('--TurnOffOptimization', action='store_true', 
                    help="Do not do any optimization at all") 

    parser.add_argument('--TurnOffMCMC', action='store_true', 
                    help="Do not do any Markov Chain Monte-Carlo (MCMC) sampling") 

    LOptMethods=[ None,     'Nelder-Mead', 'Powell',      'CG',           'BFGS', 'Newton-CG', 'L-BFGS-B', 'TNC', 'COBYLA', 'SLSQP', 'trust-constr',
                 'dogleg', 'trust-ncg',   'trust-exact', 'trust-krylov']

    parser.add_argument('--LocalOptimizationMethod', type=str, default=None, choices=LOptMethods,
                    help="The explicit choice of the local optimization method (these are choices from the Scipy.optimize.minimize method)") 

    GOptMethods=['DiffEvo', 'DualAnn']

    parser.add_argument('--GlobalOptimizationMethod', type=str, default='DualAnn', choices=GOptMethods,
                    help="The explicit choice of the global optimization method (these are a few choices from the Scipy.optimize library)") 

    parser.add_argument('--NumOfTheorPoints', type=int, default=100, 
                    help="Number of points in the theoretical approximation of the fitted functions (resolution, for plotting them later)")    

    parser.add_argument('--MinMaxPart', type=float, default=0.1, 
                    help="A part by which the theoretical approximation will be extended outside the experimental data region (for plotting)") 

    parser.add_argument('--AddInfo', type=str, default=None, 
                    help="Additional information appended to the name of the produced files, to identify them")    

    parser.add_argument('--NumOfMCMCSimPoints', type=int, default=10000, 
                    help="Total number of points in the MCMC simulation")    

    parser.add_argument('--WhichPartOfMCMCTrjToIgnore', type=float, default=0.1, 
                    help="Which part of the MCMC trajectory to ignore as an equilibration phase (value = x should be in range 0.0<=x<1.0)")  

    parser.add_argument('--MetroProbScale', type=float, default=1.0, 
                    help="Metropolis probability scaling factor (should be >0.0)")  

    parser.add_argument('--NumOfBinsInHist', type=int, default=50, 
                    help="Total number of bins in the histograms for parameters from the MCMC simulations")   


    args,unknown = parser.parse_known_args()

    # Check if the files are present
    if args.DataDescription is None:
        print("The data description file is required, but it's missing. Cannot do anything.")
        exit()
    
    if args.VariablesDescription is None:
        print("The variables description file is required, but it's missing. Cannot do anything.")
        exit()

    if args.ChannelsDescription is None:
        print("The channel description file is required, but it's missing. Cannot do anything.")
        exit()

    if args.RegularizationDescription is None:
        print("Note: the regularization description file is abscent, and thus the regularization is not applied.")
    else:
        print("Note: regularization will be applied.")
    
    # read the data and store them in a comfy format
    ListOfData, VarHand, RegInstance =  lMC3F.ReadTheData(DataDescription      = args.DataDescription,      
                                                          ChannelsDescription  = args.ChannelsDescription, 
                                                          VariablesDescription = args.VariablesDescription, 
                                                          RegularizationDescription = args.RegularizationDescription, 
                                                          xmin=None, xmax=None,  # if You wanna have more control of the plotting through these parameters, please use this script as an API library
                                                          Npts=args.NumOfTheorPoints, xMinMaxPart=args.MinMaxPart # by default we allow only a simple access to these properties 
                                                         )

    # Step 1: do the optimization, if not requested otherwise
    if not args.TurnOffOptimization:
        if args.UseLocalOptimization: # use Local optimization, if requested
            print("Started local fitting of the nonlinear parameters using "+str(args.LocalOptimizationMethod)+" method")
            # Get optimized variables vector 
            OptVarVec = lMC3F.DoTheLocalFit(ListOfData = ListOfData, 
                                            VarHand    = VarHand, 
                                            AddConstr  = RegInstance, 
                                            Method     = args.LocalOptimizationMethod, 
                                            AdditionalArguments=dict(), # if You wanna have more control of the local optimization through these parameters, please use this script as an API library
                                            ReturnDictionary=False,
                                            AddInfoToFiles = args.AddInfo
                                           )

        else:                         # otherwise use Global optimization (default behavior)
            print("Started local fitting of the nonlinear parameters using "+str(args.GlobalOptimizationMethod)+" method")
            # Get optimized variables vector 
            OptVarVec = lMC3F.DoTheGlobalFit(ListOfData = ListOfData, 
                                             VarHand    = VarHand, 
                                             AddConstr  = RegInstance, 
                                             Method     = args.GlobalOptimizationMethod, 
                                             AdditionalArguments=dict(), # if You wanna have more control of the local optimization through these parameters, please use this script as an API library
                                             ReturnDictionary=False,
                                             AddInfoToFiles = args.AddInfo
                                            )

        VarHand.updateVarVec(VarVec = OptVarVec) # now update the variables with the new, optimized parameters
        print("Fitting is done!\n    These are the parameters:")
        # print the fitted values
        res = VarHand.getDictOfVars(x = None)
        for key in res.keys():
            print(" %10s = %15.10f " % (key, res[key])) 

    # Step 2: do the MCMC simulation, if not requested otherwise
    # If the Optimization was done, we will start the simulation from the optimized parameters
    # If Optimization was skipped, the simulation will start from the user-defined values
    if not args.TurnOffMCMC:
        print("Started MCMC sampling")
        MCMCRes = lMC3F.DoTheMCMCSampling(ListOfData = ListOfData, VarHand = VarHand, AddConstr = RegInstance, 
                                          NSteps=args.NumOfMCMCSimPoints,            
                                          NStepsToIgnore=None,     # if You wanna have more control of the MCMC, please use this script as an API library
                                          PartToIgnore=args.WhichPartOfMCMCTrjToIgnore,         
                                          AcceptanceProbabilityScale = args.MetroProbScale,  
                                          AddInfoToFiles = args.AddInfo, 
                                          NBinsInHist    = args.NumOfBinsInHist, 
                                          PrintStatus = True      # if You wanna skip all the annoying messages, please use this script as an API library
                                         )
        print("MCMC sampling is done!\n    These are the parameters:")        
        for key in MCMCRes.keys():
            print(" %10s = %15.10f +/- %15.10f " % (key, MCMCRes[key][0],MCMCRes[key][1])) 

    print("Our job is done.\nGood bye! And thank You for choosing our software :)")







