import argparse
import logging
import multiprocessing
import os
import os.path
import sys

from logging import debug, getLogger
from threading import Thread

import platformdirs

from . import version
from .core import custom_logging
from .core.custom_logging import activate_excepthook, activate_logging, iprint, setup_system

config_path = os.path.abspath(platformdirs.user_data_dir("GenX3", "ArturGlavic"))
os.environ["NUMBA_CACHE_DIR"] = os.path.join(config_path, "numba_cache")


def start_interactive(args):
    """
    Start genx in interactive mode (with the gui)

    :param args: command line arguments evaluated with argparse.
    """
    debug("enter start_interactive")
    activate_excepthook()
    # Fix blurry text on Windows 10
    import ctypes

    try:
        ctypes.windll.shcore.SetProcessDpiAwareness(True)
    except:
        pass

    from .gui import main_window

    if args.infile != "":
        debug("start GUI setup with file to load")
        filename = os.path.abspath(args.infile)
        # if getattr(sys, 'frozen', False):
        # make sure the path is set to the program executable in a frozen binary
        # this fixes unnecessery re-compiles of numba functions on Windows
        #    os.chdir(os.path.dirname(os.path.abspath(sys.executable)))
        #    debug(f'changed curdir to {os.path.abspath(os.path.curdir)}')
    else:
        filename = None
    debug("start GUI setup")
    if args.debug:
        debug("initialize wx Inspection App")
        from wx.lib.mixins.inspection import InspectionMixin

        class InspApp(main_window.GenxApp, InspectionMixin):
            def OnInit(self):
                # Initialize the inspection tool.
                self.Init()
                return main_window.GenxApp.OnInit(self)

        app = InspApp(filename=filename, dpi_overwrite=args.dpi_overwrite)
    else:
        app = main_window.GenxApp(filename=filename, dpi_overwrite=args.dpi_overwrite)
    debug("setup complete, start WX MainLoop")
    app.MainLoop()
    debug("leave start_interactive")


def create_simulated_data(args):
    """Function to create simulated data from the model and add it to data.y"""
    from scipy.stats import poisson

    from .core import config as io
    from .diffev import DiffEv
    from .model_control import ModelController

    io.config.load_default(os.path.join(os.path.dirname(os.path.abspath(__file__)), "profiles", "default.profile"))
    ctrl = ModelController(DiffEv())
    mod = ctrl.model

    iprint("Loading file: %s " % args.infile)
    ctrl.load_file(args.infile)
    # io.load_opt_config(opt, config)
    iprint("File loaded")

    iprint("Simualting..")
    mod.simulate()
    iprint("Storing simualted data")
    for data_set in mod.data:
        data_set.y_raw = poisson.rvs(data_set.y_sim)
        data_set.y_command = "y"
        data_set.run_y_command()
        data_set.error_command = "sqrt(where(y > 0, y, 1.0))"
        data_set.run_error_command()

    iprint("Saving the model to %s" % args.outfile)
    ctrl.save_file(args.outfile)


def extract_parameters(args):
    """Extracts the parameters to outfile"""

    from .core import config as io
    from .diffev import DiffEv
    from .model_control import ModelController

    # Open the genx file
    io.config.load_default(os.path.join(os.path.dirname(os.path.abspath(__file__)), "profiles", "default.profile"))
    ctrl = ModelController(DiffEv())
    mod = ctrl.model

    iprint("Loading file: %s " % args.infile)
    ctrl.load_file(args.infile)
    iprint("File loaded")

    names, values = mod.parameters.get_sim_pars()
    index_names = mod.parameters.get_names()
    errors = [
        mod.parameters[index_names.index(ni)].error if mod.parameters.get_fit_state_by_name(ni) == 1 else "-"
        for ni in names
    ]
    names = [ni if mod.parameters.get_fit_state_by_name(ni) != 1 else ni + "*" for ni in names]

    paramstr = "\t".join([name for name in names]) + "\n" + "\t".join(["%f" % val for val in values]) + "\n"
    if args.error:
        paramstr += "\t".join([err for err in errors]) + "\n"
    if args.outfile != "":
        iprint(f"Appending parameter values to file {args.outfile}")
        with open(args.outfile, "a") as fout:
            # Add header
            fout.write(paramstr)
    iprint("\n" + paramstr)


def modify_file(args):
    """Modify a GenX file given command line arguments"""
    from .core import config as io
    from .diffev import DiffEv
    from .model_control import ModelController

    # Open the genx file
    io.config.load_default(os.path.join(os.path.dirname(os.path.abspath(__file__)), "profiles", "default.profile"))
    ctrl = ModelController(DiffEv())
    mod = ctrl.model

    iprint("Loading file: %s " % args.infile)
    ctrl.load_file(args.infile)
    iprint("File loaded")

    if args.datafile:
        datafile = os.path.abspath(args.datafile)
        if 0 > args.data_set or args.data_set >= len(mod.data):
            iprint("The selected data set does not exist - select one between 0 and %d" % (len(mod.data) - 1))
            return
        iprint("Loading dataset %s into data set %d" % (datafile, args.data_set))
        mod.data[args.data_set].loadfile(datafile)

        if args.outfile:
            iprint("Saving the fit to %s" % args.outfile)
            ctrl.save_file(args.outfile)

    elif args.save_datafile:
        save_datafile = os.path.abspath(args.save_datafile)
        iprint("Simualting..")
        mod.simulate()
        if 0 > args.data_set or args.data_set >= len(mod.data):
            iprint("The selected data set does not exist - select one between 0 and %d" % (len(mod.data) - 1))
            return
        iprint("Exporting data set %d into ASCII file %s" % (args.data_set, save_datafile))
        mod.data[args.data_set].save_file(save_datafile)


def set_numba_single():
    debug("Setting numba JIT compilation to single CPU")
    os.environ["NUMBA_CACHE_DIR"] = os.path.join(config_path, "single_cpu_numba_cache")

    import numba

    old_jit = numba.jit

    def jit(*args, **opts):
        opts["parallel"] = False
        return old_jit(*args, **opts)

    numba.jit = jit
    numba.GENX_OVERWRITE_SINGLE = True
    try:
        numba.set_num_threads(1)
    except AttributeError:
        pass


class InputThread(Thread):
    def __init__(self):
        Thread.__init__(self, daemon=True)
        self.stop_fit = False

    def run(self):
        # concole input has to be queried differently for non-UNIX systems
        if sys.platform.startswith("win"):
            self.run_windows()
        else:
            self.run_unix()

    def run_windows(self):
        import msvcrt
        import time

        while not self.stop_fit:
            if msvcrt.kbhit():
                debug("Key event received")
                key = msvcrt.getwch()
                if "q" in key:
                    self.stop_fit = True
            time.sleep(0.1)

    def run_unix(self):
        import fcntl
        import selectors
        import time

        # set sys.stdin non-blocking
        orig_fl = fcntl.fcntl(sys.stdin, fcntl.F_GETFL)
        fcntl.fcntl(sys.stdin, fcntl.F_SETFL, orig_fl | os.O_NONBLOCK)

        m_selector = selectors.DefaultSelector()
        m_selector.register(sys.stdin, selectors.EVENT_READ)
        while True:
            (k, evt) = m_selector.select()[0]
            res = k.fileobj.read()
            if "q" in res:
                self.stop_fit = True
                break
        m_selector.close()


def start_fitting(args, rank=0):
    """
    Function to start fitting from the command line.
    """
    import time

    from .core import config as io
    from .core.console import calc_errorbars, setup_console
    from .model_control import ModelController

    # Open the genx file
    io.config.load_default(os.path.join(os.path.dirname(os.path.abspath(__file__)), "profiles", "default.profile"))
    if not args.bumps:
        from .diffev import DiffEv

        ctrl = ModelController(DiffEv())
        mod = ctrl.model
        opt: DiffEv = ctrl.optimizer
        set_optimiser_pars = set_diffev_pars
    else:
        from .bumps_optimizer import BumpsOptimizer

        ctrl = ModelController(BumpsOptimizer())
        mod = ctrl.model
        opt: BumpsOptimizer = ctrl.optimizer
        set_optimiser_pars = set_bumps_pars

    if rank == 0:
        setup_console(ctrl, args.error, args.outfile, use_curses=(args.use_curses and not args.mpi))

    if (args.mpi or args.pr > 0) and not args.disable_numba:
        try:
            set_numba_single()
        except ImportError:
            pass

    if rank == 0:
        iprint("Loading model %s..." % args.infile)
    ctrl.load_file(args.infile)
    # has to be used in order to save everything....
    if args.esave:
        io.config.set("solver", "save all evals", True)
    # Simulate, this will also compile the model script
    if rank == 0:
        iprint("Simulating model...")
    mod.simulate()

    # Sets up the fitting ...
    set_optimiser_pars(opt, args)
    if rank == 0:
        iprint("Setting up the optimizer...")
        iprint(opt)

    if args.outfile and rank == 0:
        iprint("Saving the initial model to %s" % args.outfile)
        ctrl.save_file(args.outfile)

    # To start the fitting
    if rank == 0:
        iprint("Fitting starting...")
        t1 = time.time()
    # print opt.use_mpi, opt.use_parallel_processing
    opt.start_fit(mod)
    if rank == 0:
        inp = InputThread()
        inp.start()
    else:

        class Inp:
            stop_fit = False

        inp = Inp()
    while opt.is_running():
        try:
            time.sleep(0.1)
            if inp.stop_fit:
                opt.stop = True
        except KeyboardInterrupt:
            iprint("KeyboardInterrupt, trying to stop fit.")
            opt.stop = True

    if rank == 0:
        t2 = time.time()
        iprint("Fitting finished!")
        iprint("Time to fit: ", (t2 - t1) / 60.0, " min")

    if rank == 0:
        iprint("Updating the parameters")
        mod.parameters.set_value_pars(opt.best_vec)

    if args.outfile and rank == 0:
        if args.error:
            iprint("Calculating errorbars")
            calc_errorbars(mod, opt)
        iprint("Saving the fit to %s" % args.outfile)
        # opt.set_use_mpi(False)
        ctrl.save_file(args.outfile)

    if rank == 0:
        iprint("Fitting successfully completed")


def set_diffev_pars(optimiser, args):
    """
    Sets the optimiser parameters from args
    """
    if args.pr:
        optimiser.set_processes(args.pr)
        optimiser.set_use_parallel_processing(True)

    if args.mpi:
        optimiser.set_use_mpi(True)
    else:
        optimiser.set_use_mpi(False)

    if args.cs:
        optimiser.set_chunksize(args.cs)

    if args.mgen:
        optimiser.set_max_generations(args.mgen)
        optimiser.set_use_max_generations(True)

    if args.pops:
        optimiser.set_pop_size(args.pops)
        optimiser.set_use_pop_mult(False)

    if args.asi:
        optimiser.set_autosave_interval(args.asi)
        optimiser.set_use_autosave(True)

    if args.var >= 0:
        optimiser.opt.min_parameter_spread = args.var

    if args.km >= 0:
        optimiser.set_km(args.km)

    if args.kr >= 0:
        optimiser.set_kr(args.kr)


def set_bumps_pars(optimiser, args):
    if args.pr:
        optimiser.opt.use_parallel_processing = True
        optimiser.opt.parallel_processes = args.pr


def compile_numba(cache_dir=None):
    try:
        # perform a compilation of numba functions with console feedback
        import numba

        from .models.lib.numba_integration import configure_numba

        configure_numba()

        import inspect

        real_jit = numba.jit

        class UpdateJit:
            update_counter = 1

            def __call__(self, *args, **opts):
                if inspect.stack()[1][3] != "<lambda>":
                    print(f"compiling numba functions {self.update_counter}/21")
                    self.update_counter += 1
                return real_jit(*args, **opts)

        print("Starting to compile numba functions..")
        numba.jit = UpdateJit()
        print("paratt")
        from .models.lib import paratt_numba

        print("neutron")
        from .models.lib import neutron_numba

        print("instrument")
        from .models.lib import instrument_numba

        print("offspec")
        from .models.lib import offspec

        print("surface")
        from .models.lib import surface_scattering

        numba.jit = real_jit
    except Exception as e:
        print("An exception occured when trying to compile the numba functions:")
        print(e)
        return 1
    return 0


def main():
    multiprocessing.set_start_method("spawn")
    multiprocessing.freeze_support()
    if os.path.abspath(__file__).startswith("/snap"):
        # try fix multiprocessing in SNAP
        snap_sem = "/snap.genx.mp"
        multiprocessing.current_process()._config["semprefix"] = snap_sem
    else:
        snap_sem = None
    # Attempt to load mpi:
    try:
        from mpi4py import MPI
    except ImportError:
        rank = 0
        __mpi__ = False
    else:
        __mpi__ = bool(MPI.COMM_WORLD.Get_size() > 1)
        rank = MPI.COMM_WORLD.Get_rank()

    parser = argparse.ArgumentParser(
        prog="genx",
        description="GenX %s, fits data to a model." % version.__version__,
        epilog="For support, manuals and bug reporting see http://genx.sf.net",
    )
    run_group = parser.add_mutually_exclusive_group()
    run_group.add_argument("-r", "--run", action="store_true", help="run GenX fit (no gui)")
    if __mpi__:
        run_group.add_argument("--mpi", action="store_true", help="run GenX fit with mpi (no gui)")
    run_group.add_argument("-g", "--gen", action="store_true", help="generate data.y with poisson noise added")
    run_group.add_argument("--pars", action="store_true", help="extract the parameters from the infile")
    run_group.add_argument("--mod", action="store_true", help="modify the GenX file")
    opt_group = parser.add_argument_group("optimization arguments")
    opt_group.add_argument("--pr", type=int, default=0, help="Number of processes used in parallel fitting.")
    opt_group.add_argument("--cs", type=int, default=0, help="Chunk size used for parallel processing.")
    opt_group.add_argument("--mgen", type=int, default=0, help="Maximum number of generations that is used in a fit")
    opt_group.add_argument("--pops", type=int, default=0, help="Population size - number of individuals.")
    opt_group.add_argument("--asi", type=int, default=0, help="Auto save interval (generations).")
    opt_group.add_argument("--km", type=float, default=-1, help="Mutation constant (float 0 < km < 1)")
    opt_group.add_argument("--kr", type=float, default=-1, help="Cross over constant (float 0 < kr < 1)")
    opt_group.add_argument("-s", "--esave", action="store_true", help="Force save evals to gx file.")
    opt_group.add_argument("-e", "--error", action="store_true", help="Calculate/export error bars.")
    opt_group.add_argument(
        "--var", type=float, default=-1, help="Minimum relative parameter variation to stop the fit (%%)"
    )
    opt_group.add_argument(
        "--bumps", action="store_true", help="Use Bumps DREAM optimizer instead of GenX Differential Evolution"
    )
    data_group = parser.add_argument_group("data arguments")
    data_group.add_argument(
        "-d", dest="data_set", type=int, default=0, help="Active data set to act upon. Index starting at 0."
    )
    data_group.add_argument("--load", dest="datafile", help="Load file into active data set. Index starting at 0.")
    data_group.add_argument("--export", dest="save_datafile", help="Save active data set to file. Index starting at 0.")
    data_group = parser.add_argument_group("startup options")
    data_group.add_argument(
        "-l", "--logfile", dest="logfile", default=None, type=str, help="Output debug information to logfile."
    )
    data_group.add_argument(
        "--debug",
        dest="debug",
        default=False,
        action="store_true",
        help="Show additional debug information on console/logfile",
    )
    data_group.add_argument(
        "--dpi-scale",
        dest="dpi_overwrite",
        default=None,
        type=float,
        help="Overwrite the detection of screen dpi scaling factor (=72/dpi)",
    )
    data_group.add_argument(
        "--no-curses",
        dest="use_curses",
        default=True,
        action="store_false",
        help="Disable Curses interactive console interface for command line " "fitting on UNIX systems.",
    )
    data_group.add_argument(
        "--disable-nb",
        dest="disable_numba",
        default=False,
        action="store_true",
        help="Disable the use of numba JIT compiler",
    )
    data_group.add_argument(
        "--nb1",
        dest="numba_single",
        default=False,
        action="store_true",
        help="Compile numba JIT functions without parallel computing support (use one core only). "
        "Caching in this case is done in a different user directory.",
    )
    data_group.add_argument(
        "--compile-nb",
        dest="compile_nb",
        default=False,
        action="store_true",
        help="Perform a first-/recompilation of the numba modules and exit",
    )

    parser.add_argument(
        "infile", nargs="?", default="", help="The .gx or .hgx file to load or .ort file to use as basis for model"
    )
    parser.add_argument("outfile", nargs="?", default="", help="The .gx  or hgx file to save into")

    args = parser.parse_args()
    if not __mpi__:
        args.mpi = False

    if args.run or args.mpi or args.pars or args.mod:
        # make sure at least info-messages are shown (default is warning)
        custom_logging.CONSOLE_LEVEL = min(logging.INFO, custom_logging.CONSOLE_LEVEL)
    if rank > 0:
        custom_logging.CONSOLE_LEVEL = logging.WARNING
    setup_system()
    if snap_sem:
        logging.debug(f"Multiprocessing was setup with semprefix for SNAP to {snap_sem}")

    if args.logfile:
        if __mpi__:
            lfbase, lfend = args.logfile.rsplit(".", 1)
            activate_logging(f"{lfbase}_{rank:02}.{lfend}")
        else:
            activate_logging(args.logfile)
    debug("Arguments from parser: %s" % args)
    if not args.outfile:
        args.outfile = args.infile
    args.outfile = os.path.abspath(args.outfile)
    if args.infile:
        args.infile = os.path.abspath(args.infile)

    if args.disable_numba:
        debug("disable numba")
        # set numba flag
        from genx.models import lib as modellib

        modellib.USE_NUMBA = False
    elif args.numba_single:
        set_numba_single()
    if args.compile_nb:
        sys.exit(compile_numba())

    if args.run:
        start_fitting(args)
    elif args.mpi:
        start_fitting(args, rank)
    elif args.gen:
        create_simulated_data(args)
    elif args.pars:
        extract_parameters(args)
    elif args.mod:
        modify_file(args)
    elif not args.run and not args.mpi:
        # Check if the application has been frozen
        if getattr(sys, "frozen", False):
            # Redirect all the output to log files
            log_file_path = platformdirs.user_log_dir("GenX3", "ArturGlavic")
            # Create dir if not found
            if not os.path.exists(log_file_path):
                os.makedirs(log_file_path)
            # print log_file_path
            # log_file_path = genx_gui._path + 'app_data/'
            sys.stdout = open(log_file_path + "/genx_out.log", "w")
            sys.stderr = open(log_file_path + "/genx_err.log", "w")
        start_interactive(args)


if __name__ == "__main__":
    main()
