"""
Functions for handling command-line input using argparse module.

TODO: add option to split pdist and plot once the pdist to txt feature is done
      so the pdist txt file could be: X_column, Y_column, Z_matrix_columns.
      This would be something like pdist.dap default and option to check and 
      read in a wedap or westpa pdist file for pdist function.
"""

import argparse
import sys

# import and use gooey conditionally
# adapted from https://github.com/chriskiehl/Gooey/issues/296
try:
    import gooey
    #from gooey import Gooey
    #from gooey import GooeyParser
except ImportError:
    gooey = None

def flex_add_argument(f):
    """Make the add_argument accept (and ignore) the widget option."""

    def f_decorated(*args, **kwargs):
        kwargs.pop('widget', None)
        return f(*args, **kwargs)

    return f_decorated

# Monkey-patching a private class…
argparse._ActionsContainer.add_argument = \
    flex_add_argument(argparse.ArgumentParser.add_argument)

# Do not run GUI if it is not available or if command-line arguments are given.
if gooey is None or len(sys.argv) > 1:
    ArgumentParser = argparse.ArgumentParser

    def gui_decorator(f):
        return f
else:
    ArgumentParser = gooey.GooeyParser
    gui_decorator = gooey.Gooey(
        program_name='WEKAP',
        #navigation='TABBED',
        #advanced=True,
        suppress_gooey_flag=True,
        optional_cols=4, 
        default_size=(1000, 600),
        #tabbed_groups=True,
    )

@gui_decorator
def create_cmd_arguments(): 
    """
    Use the `argparse` module to make the optional and required command-line
    arguments for the `wekap`. 

    Parameters 
    ----------

    Returns
    -------
    argparse.ArgumentParser: 
        An ArgumentParser that is used to retrieve command line arguments. 
    """
    wekap_desc = "================================================================ \n" + \
                 "=== Weighted Ensemble Kinetics Analysis and Plotting (WEKAP) === \n" + \
                 "================================================================ \n" + \
                 "\nPlot flux values from a direct.h5 file as rates"
                 # TODO: add examples to desc

    # create argument parser (gooey based if available)
    if gooey is None:
        parser = argparse.ArgumentParser(description=wekap_desc, 
                                        formatter_class=argparse.RawDescriptionHelpFormatter)
    else:
        parser = gooey.GooeyParser(description=wekap_desc, 
                                   formatter_class=argparse.RawDescriptionHelpFormatter)

    ###########################################################
    ######################## ARGUMENTS ########################
    ###########################################################
    # nargs = '?' "One argument will be consumed from the command line if possible, 
        # and produced as a single item. If no command-line argument is present, 
        # the value from default will be produced."

    main = parser.add_argument_group("Main Arguments")
    optional = parser.add_argument_group("Optional Extra Arguments")

    # TODO: allow for multiple h5 files and calc error using bootstrapping
    main.add_argument(
        "--direct", "--direct-h5", "-dh5",
        dest="direct",
        type=str,
        default="direct.h5",
        help="Name of output direct.h5 file from WESTPA w_direct or w_ipa.",
        widget="FileChooser"
    )
    main.add_argument(
        "--assign", "--assign-h5", "-ah5",
        dest="assign",
        type=str,
        default=None,
        help="Name of specific assign.h5 file. Needed for labeled population data.",
        widget="FileChooser"
    )
    main.add_argument(
        "--tau", "-t",
        dest="tau",
        type=float,
        default=100e-12,
        help="The resampling interval of the WE simulation in seconds. "
             "Default: 100ps = 100 * 10^-12 (s).",
    )
    main.add_argument(
        "--state",
        dest="state",
        type=int,
        default=1,
        help="State for flux calculations (flux into `state`), 0 = A and 1 = B, etc."
    )
    main.add_argument("-o", "--output", default=None,
                        dest="output_path",
                        help="The filename to which the plot will be saved. "
                             "Various image formats are available. You " 
                             "may choose one by specifying an extension. "
                             "\nLeave this empty if you don't want to save "
                             "the plot to a serparate file.",
                        type=str)
    optional.add_argument(
        "--label",
        dest="label",
        type=str,
        default=None,
        help="Data label."
    )
    optional.add_argument(
        "--statepop",
        dest="statepop",
        type=str,
        default="direct",
        choices=["direct", "assign"],
        help="'direct' for state_population_evolution from direct.h5 or 'assign' for labeled_populations from assign.h5."
    )
    optional.add_argument(
        "--units",
        dest="units",
        type=str,
        default="rates",
        choices=["rates", "mfpts"],
        help="Measurement units. Can be `rates` (default) or `mfpts`."
    )
    optional.add_argument(
        "--savefig",
        dest="savefig",
        type=str,
        default=None,
        help="Path to optionally save the figure."
    )
    optional.add_argument(
        "--cumulative-avg", "-cavg",
        dest="cumulative_avg",
        action="store_true",
        help="Set to True when kinetics were calculated with cumulative averaging."
    )
    # TODO: this isn't really needed for anyone else and can use hline instead
    optional.add_argument(
        "--exp-values", "-exp",
        dest="exp_values",
        action="store_true",
        help="Plot experimental D1-->D2 values."
    )
    # TODO: moltime arg for WE iteration vs moltime
#     optional.add_argument("-mt", "--molecular-time", "--moltime"
#                         dest = "moltime",
#                         help = "Include this argument to not output the plot to "
#                         "your display.", 
#                         action= "store_true") 
    optional.add_argument("-nots", "--no-output-to-screen",
                        dest = "no_output_to_screen",
                        help = "Include this argument to not output the plot to "
                        "your display.", 
                        action= "store_true") 
    optional.add_argument("-npb", "--no-progress-bar",
                        dest = "no_pbar",
                        help = "Include this argument to not output the tqdm progress bar.",
                        action= "store_true")

    ##########################################################
    ############### FORMATTING ARGUMENTS #####################
    ##########################################################

    formatting = parser.add_argument_group("Plot Formatting Arguments") 

    formatting.add_argument("--style", default="default", nargs="?",
                        dest="style",
                        help="mpl style, can leave blank to use default, "
                             "input `None` for basic mpl settings, can use a custom "
                             "path to a mpl.style text file, or could use a mpl included "
                             "named style, e.g. `ggplot`. "
                             "Edit the wedap/styles/default.mplstyle file to "
                             "change default wedap plotting style options.",
                        type=str)
    formatting.add_argument("--color",
                        dest="color", help="Color for 1D plots, contour lines, and trace plots.",
                        widget="ColourChooser")
    formatting.add_argument("--linewidth", "-lw", default=None, nargs="?",
                        dest="linewidth", help="Linewidth for 1D plots, contour lines, and trace plots.",
                        type=float)
    formatting.add_argument("--linestyle", "-ls", default="-", nargs="?",
                        dest="linestyle", help="Linestyle for 1D plots, contour lines, and trace plots.",
                        type=str)
    formatting.add_argument("--xlabel", dest="xlabel", type=str)
    formatting.add_argument("--xlim", help="LB UB", dest="xlim", nargs=2, type=float)
    formatting.add_argument("--ylabel", dest="ylabel", type=str)
    formatting.add_argument("--ylim", help="LB UB", dest="ylim", nargs=2, type=float)
    formatting.add_argument("--title", dest="title", type=str)
    formatting.add_argument("--suptitle", dest="suptitle", type=str)
    formatting.add_argument("--grid", dest="grid", default=False, action="store_true")
    formatting.add_argument("--axvline", "-vl", help="Can be a single value or a list of lines.", 
                            dest="axvline", nargs="*", type=float)
    formatting.add_argument("--axhline", "-hl", help="Can be a single value or a list of lines.",
                            dest="axhline", nargs="*", type=float)
#     formatting.add_argument("--text", help="3 args for ax.text: x, y, string",
#                             dest="text", nargs="3", type=float)
#     formatting.add_argument("--figsize", default=None, nargs=2,
#                             dest="figsize",
#                             help="Matplotlib figure size, e.g. (6,4)",
#                             type=float)
    formatting.add_argument('--postprocess', '--postprocess-function', '-ppf', default=None,
                            dest='postprocess_func',
                            help='After plotting data, load and execute the '
                                 'Python function specified by '
                                 'POSTPROCESS_FUNC. POSTPROCESS_FUNC should be '
                                 'a string of the form ``mymodule.myfunction``.'
                                 'Example: '
                                 '``--postprocess mymodule.myfunction``.')

    # return the argument parser
    return parser 

# TODO: adjust all to fix str/int/type auto
def handle_command_line(argument_parser): 
    """
    Take command line arguments, check for issues, return the arguments. 

    Parameters
    ----------
    argument_parser : argparse.ArgumentParser 
        The argument parser that is returned in `create_cmd_arguments()`.
    
    Returns
    -------
    argparse.NameSpace
        Contains all arguments passed into wedap.
    
    Raises
    ------  
    Prints specific issues to terminal.
    """
    # retrieve args
    args = argument_parser.parse_args() 

     # TODO: check make sure that -jp and -3d not being using at same time

    # h5 file and file exists
    # if not os.path.exists(args.file) or not ".h5" in args.file:  
    #     # print error message and exits
    #     sys.exit("Must input file that exists and is in .h5 file format.")

    # if not args.percentage.isdigit():  # not correct input   
    #     # print out any possible issues
    #     print("You must input a percentage digit. EXAMPLES: \
    #     \n '-p 40' \n You CANNOT add percent sign (eg. 50%) \n You \
    #     CANNOT add decimals (eg. 13.23)") 
    #     sys.exit(0) # exit program

    # # ignore if NoneType since user doesn't want --maxEnsembleSize parameter
    # if args.max_ensemble_size is None: 
    #     pass

    # elif not args.max_ensemble_size.isdigit(): # incorrect input 
    #     # needs to be whole number
    #     print("You must input a whole number with no special characters (eg. 4).")  
    #     sys.exit(0) # exit program 

    # elif args.max_ensemble_size is '0': # user gives 0 to --maxEnsembleSize flag 
    #     print("You cannot input '0' to --maxEnsembleSize flag.")
    #     sys.exit(0) # exit program 

    return args