import os
import numpy as np
import subprocess
import mmap
import warnings
import h5py
import re

from mcstasscript.helper.formatting import bcolors
from mcstasscript.data.data import McStasMetaData
from mcstasscript.data.data import McStasDataBinned
from mcstasscript.data.data import McStasDataEvent


class ManagedMcrun:
    """
    A class for performing a mcstas simulation and organizing the data
    into python objects

    ManagedMcrun is usually called by the instrument class of
    McStasScript but can be used independently.  It runs the mcrun
    command using the system command, and if this is not in the path,
    the absolute path can be given in a keyword argument executable_path.

    Attributes
    ----------
    name_of_instrumentfile : str
        Name of instrument file to be executed

    data_folder_name : str
        Name of datafolder mcrun writes to disk

    ncount : int
        Number of rays to simulate

    mpi : int
        Number of mpi threads to run

    parameters : dict
        Dictionary of parameter names and values for this simulation

    custom_flags : string
        Custom flags that are passed to the mcrun command

    executable_path : string
        Path to the mcrun command (can be empty if already in path)

    Methods
    -------
    run_simulation()
        Runs simulation, returns list of McStasData instances

    """

    def __init__(self, instr_name, **kwargs):
        """
        Performs call to McStas with given options

        Uses subprocess to call mcrun / mxrun to perform simulation of given
        instrument file.

        Parameters
        ----------
        instr_name : str
            Name of instrument file to be simulated

        kwargs : keyword arguments
            output_path : str, required
                Sets data_folder_name
            ncount : int, default 1E6
                Sets ncount
            mpi : int, default None
                Sets thread count, None to disable mpi
            gravity : bool, default False
                Enables gravity if True
            parameters : dict
                Sets parameters
            custom_flags : str, default ""
                Sets custom_flags passed to mcrun
            executable_path : str
                Path to mcrun command, "" if already in path
            increment_folder_name : bool, default True
                If True, automatically appends output_path to make it unique
            force_compile : bool, default True
                If True, forces compile. If False no new instrument is written
            run_folder : str
                Path to folder in which to run McStas
            openacc : bool, default False
                If True, adds the --openacc flag to mcrun call
            NeXus : bool, default False
                If True, adds the --format=NeXus to mcrun call

        """

        self.name_of_instrumentfile = instr_name

        self.data_folder_name = ""
        self.ncount = int(1E6)
        self.mpi = None
        self.gravity = False
        self.openacc = False
        self.NeXus = False
        self.parameters = {}
        self.custom_flags = ""
        self.executable_path = ""
        self.executable = ""
        self.increment_folder_name = True
        self.compile = True
        self.run_path = "."
        self.seed = None
        self.suppress_output = False

        # executable_path always in kwargs
        if "executable_path" in kwargs:
            self.executable_path = kwargs["executable_path"]

        if "executable" in kwargs:
            self.executable = kwargs["executable"]

        if "output_path" in kwargs:
            self.data_folder_name = kwargs["output_path"]
        else:
            raise NameError(
                "ManagedMcrun needs output_path to load data, add "
                + "with keyword argument.")

        if "ncount" in kwargs:
            self.ncount = int(kwargs["ncount"])

            if self.ncount < 1:
                raise ValueError("ncount should be a positive integer, was "
                                 + str(self.ncount))

        if "mpi" in kwargs:
            self.mpi = kwargs["mpi"]
            try:
                self.mpi = int(self.mpi)
            except (TypeError, ValueError) as e:
                if self.mpi is not None:
                    raise RuntimeError("MPI should be an integer, was "
                                       + str(self.mpi))

            if self.mpi is not None:
                if self.mpi < 1:
                    raise ValueError("MPI should be an integer larger than"
                                     + " 0, was " + str(self.mpi))

        if "gravity" in kwargs:
            self.gravity = kwargs["gravity"]

        if "openacc" in kwargs:
            self.openacc = kwargs["openacc"]

        if "NeXus" in kwargs:
            self.NeXus = kwargs["NeXus"]

        if "seed" in kwargs:
            self.seed = kwargs["seed"]

        if "parameters" in kwargs:
            self.parameters = kwargs["parameters"]

            if not isinstance(self.parameters, dict):
                raise RuntimeError("Parameters should be given as dict.")

        if "custom_flags" in kwargs:
            self.custom_flags = kwargs["custom_flags"]

            if not isinstance(self.custom_flags, str):
                raise RuntimeError("ManagedMcrun detected given customf_flags"
                                   + " was not a string.")

        if "increment_folder_name" in kwargs:
            self.increment_folder_name = kwargs["increment_folder_name"]

        if "force_compile" in kwargs:
            self.compile = kwargs["force_compile"]

        if "run_path" in kwargs:
            self.run_path = kwargs["run_path"]

        if "suppress_output" in kwargs:
            self.suppress_output = bool(kwargs["suppress_output"])


        # get relevant paths and check their validity
        current_directory = os.getcwd()

        if not os.path.isabs(self.data_folder_name):
            self.data_folder_name = os.path.join(current_directory,
                                                 self.data_folder_name)
        else:
            split_data_path = os.path.split(self.data_folder_name)
            if not os.path.isdir(split_data_path[0]):
                raise RuntimeError("Parent folder for datafolder invalid: "
                                   + str(split_data_path[0]))

        if not os.path.isabs(self.run_path):
            self.run_path = os.path.join(current_directory, self.run_path)
        else:
            split_run_path = os.path.split(self.run_path)
            if not os.path.isdir(split_run_path[0]):
                raise RuntimeError("Parent folder for run_path invalid: "
                                   + str(split_run_path[0]))

        if not os.path.isdir(self.run_path):
            raise RuntimeError("ManagedMcrun found run_path to "
                               + "be invalid: " + str(self.run_path))

        if not os.path.isdir(self.executable_path):
            raise RuntimeError("ManagedMcrun found executable_path to "
                               + "be invalid: " + str(self.executable_path))

    def run_simulation(self):
        """
        Runs McStas simulation described by initializing the object
        """

        # construct command to run
        option_string = ""
        if self.compile:
            option_string += "-c "

        if self.gravity:
            option_string += "-g "

        if self.NeXus:
            option_string += "--format=NeXus "

        if self.openacc:
            option_string += "--openacc "

        if self.mpi is not None:
            mpi_string = "--mpi=" + str(self.mpi) + " "  # Set mpi
        else:
            mpi_string = ""

        if self.seed is not None:
            seed_string = "--seed=" + str(self.seed) + " " # Set seed
        else:
            seed_string = ""

        option_string = (option_string
                         + "-n " + str(self.ncount) + " " # Set ncount
                         + mpi_string
                         + seed_string)

        if os.path.exists(self.data_folder_name):
            if self.increment_folder_name:
                counter = 0
                new_name = self.data_folder_name + "_" + str(counter)
                while os.path.isdir(new_name):
                    counter = counter + 1
                    new_name = self.data_folder_name + "_" + str(counter)

                self.data_folder_name = new_name
            else:
                raise NameError("output_path already exists and "
                                + "increment_folder_name was set to False.")

        if len(self.data_folder_name) > 0:
            option_string = (option_string
                             + "-d "
                             + self.data_folder_name)

        # add parameters to command
        parameter_string = ""
        for key, val in self.parameters.items():
            parameter_string = (parameter_string + " "
                                + str(key)  # parameter name
                                + "="
                                + str(val))  # parameter value

        mcrun_full_path = os.path.join(self.executable_path, self.executable)
        if len(self.executable_path) > 1:
            if not (self.executable_path[-1] == "\\"
                    or self.executable_path[-1] == "/"):
                mcrun_full_path = os.path.join(self.executable_path,
                                               self.executable)
        mcrun_full_path = '"' + mcrun_full_path + '"' # Path in quotes to allow spaces

        # Run the mcrun command on the system
        full_command = (mcrun_full_path + " "
                        + option_string + " "
                        + self.custom_flags + " "
                        + self.name_of_instrumentfile
                        + parameter_string)

        process = subprocess.run(full_command, shell=True,
                                 stdout=subprocess.PIPE,
                                 stderr=subprocess.STDOUT,
                                 universal_newlines=True,
                                 cwd=self.run_path)

        if self.suppress_output is False:
            print_sim_output(process.stdout)

        if not os.path.isdir(self.data_folder_name):
            warnings.warn("Simulation did not create data folder, most likely failed.")

    def load_results(self, *args):
        """
        Method for loading data from a mcstas simulation

        Loads data on all monitors in a McStas data folder, and returns these
        as a list of McStasData objects.

        Parameters
        ----------

        optional first argument : str
            path to folder from which data should be loaded

        """

        if len(args) == 0:
            data_folder_name = self.data_folder_name
        elif len(args) == 1:
            data_folder_name = args[0]
        else:
            raise RuntimeError("load_results can be called "
                               + "with 0 or 1 arguments")

        if os.path.isdir(data_folder_name):
            return load_results(data_folder_name)
        else:
            warnings.warn("No data available to load.")
            return None


def load_results(data_folder_name):
    """
    Function for loading data from a mcstas simulation

    Loads data on all monitors in a McStas data folder, and returns these
    as a list of McStasData objects.

    Parameters
    ----------

    data_folder_name : str
        path to folder from which data should be loaded

    """

    if not os.path.isdir(data_folder_name):
        raise NameError("Given data directory does not exist.")

    # Find all data files in generated folder
    files_in_folder = os.listdir(data_folder_name)

    # Raise an error if mccode.sim is not available
    if "mccode.sim" in files_in_folder:
        NeXus = False
    elif "mccode.h5" in files_in_folder:
        NeXus = True
    else:
        raise NameError("No mccode.sim or mccode.h5 in data folder.")

    if NeXus:
        # Open mccode to read metadata for all datasets written to disk
        with h5py.File(os.path.join(data_folder_name, "mccode.h5"), "r", swmr=True) as f:

            # Pass file object to all functions to avoid multiple open / close
            metadata_list = load_metadata_nexus(f)

            results = []
            for metadata in metadata_list:
                result = load_monitor_nexus(metadata, f)
                result.set_data_location(data_folder_name)
                results.append(result)

    else:
        # Older workflow, still handles both text and NeXus
        metadata_list = load_metadata(data_folder_name)

        results = []
        for metadata in metadata_list:
            result = load_monitor(metadata, data_folder_name)
            result.set_data_location(data_folder_name)
            results.append(result)

    return results

def load_metadata(data_folder_name):
    """
    Function that loads metadata from a mcstas simulation

    Returns list of metadata objects corresponding to each monitor, all
    information is taken from mccode.sim file.

    Parameters
    ----------

    first argument : str
        path to folder from which metadata should be loaded
    """

    if not os.path.isdir(data_folder_name):
        raise NameError("Given data directory does not exist.")

    # Find all data files in generated folder
    files_in_folder = os.listdir(data_folder_name)

    # Raise an error if mccode.sim is not available
    if "mccode.sim" in files_in_folder:
        return load_metadata_text(data_folder_name)
    elif "mccode.h5" in files_in_folder:
        with h5py.File(os.path.join(data_folder_name, "mccode.h5"), "r", swmr=True) as f:
            return load_metadata_nexus(f)
    else:
        raise NameError("No mccode.sim or mccode.h5 in data folder.")


def load_metadata_text(data_folder_name):
    instrument_parameters = {}

    # Open mccode to read metadata for all datasets written to disk
    with open(os.path.join(data_folder_name, "mccode.sim"), "r") as f:

        # Loop that reads mccode.sim sections
        metadata_list = []
        current_object = None
        in_data = False
        in_sim = False
        for lines in f:
            # Could read other details about run

            if lines == "end data\n":
                # No more data for this metadata object
                # Add parameter information
                current_object.add_info("Parameters", instrument_parameters)
                # Extract the information
                current_object.extract_info()
                # Add to metadata list
                if current_object.filename != "":
                    metadata_list.append(current_object)
                # Stop reading data
                in_data = False

            if in_sim:
                if "Param" in lines:
                    parm_lst = lines.split(':')[1].split('=')
                    try:
                        value = float(parm_lst[1].strip())
                    except ValueError:
                        value = parm_lst[1].strip()

                    instrument_parameters[parm_lst[0].strip()] = value

            if in_data:
                # This line contains info to be added to metadata
                colon_index = lines.index(":")
                key = lines[2:colon_index]
                value = lines[colon_index + 2:].strip()
                current_object.add_info(key, value)

            if lines == "begin data\n":
                # Found data section, create new metadata object
                current_object = McStasMetaData()
                # Start recording data to metadata object
                in_data = True

            if 'begin simulation:' in lines:
                in_sim = True
            if 'end simulation:' in lines:
                in_sim = False

        # Close mccode.sim
        f.close()

        """
        # Create a list for McStasData instances to return
        results = []


        # Load datasets described in metadata list individually
        for metadata in metadata_list:

            # Load data with numpy
            data = np.loadtxt(data_folder_name
                              + "/"
                              + metadata.filename.rstrip())

            # Split data into intensity, error and ncount
            if type(metadata.dimension) == int and metadata.dimension == 0:
                Intensity = data.T

            if type(metadata.dimension) == int and metadata.dimension != 0:

                xaxis = data.T[0, :]
                Intensity = data.T[1, :]
                Error = data.T[2, :]
                Ncount = data.T[3, :]

            elif len(metadata.dimension) == 2:
                xaxis = []  # Assume evenly binned in 2d
                data_lines = metadata.dimension[1]
                Intensity = data.T[:, 0:data_lines - 1]
                Error = data.T[:, data_lines:2*data_lines - 1]
                Ncount = data.T[:, 2*data_lines:3*data_lines - 1]
        """

    return metadata_list


def load_metadata_nexus(file_object):
    instrument_parameters = {}

    f = file_object

    if "entry1" not in list(f.keys()):
        raise ValueError("h5 file not formatted as expected.")

    if "data" not in list(f["entry1"].keys()):
        raise ValueError("h5 file not formatted as expected.")

    if "simulation" not in list(f["entry1"].keys()):
        raise ValueError("h5 file not formatted as expected.")

    if "Param" not in list(f["entry1"]["simulation"].keys()):
        raise ValueError("h5 file not formatted as expected.")

    # Common information

    # Instrument parameters
    instrument_parameters = {}
    loaded_par_dict = f["entry1"]["simulation"]["Param"].attrs
    for par_name in loaded_par_dict:
        if par_name == "NX_class":
            continue

        try:
            value = float(loaded_par_dict[par_name])
        except:
            value = str(loaded_par_dict[par_name])

        instrument_parameters[par_name] = value

    metadata_list = []

    # For each entry in data, make a metadata object
    for key in f["entry1"]["data"].keys():

        # Make the metadata object and add instrument parameters
        metadata = McStasMetaData()
        metadata.add_info("Parameters", instrument_parameters)

        # Add NeXus field name
        metadata.add_info("NeXus_field", key)

        # Add all the read info from attribute section
        info = dict(f["entry1"]["data"][key].attrs)
        info = decode_dict(info)
        for name, value in info.items():
            if isinstance(value, bytes):
                value = value.decode('utf-8')

            metadata.add_info(name, value)

        metadata_list.append(metadata)

        # Now all info is added, extract info loads it into nice attributes
        metadata.extract_info()

    return metadata_list


def decode_dict(dictionary):
    for key, value in dictionary.items():
        if isinstance(value, bytes):
            try:
                dictionary[key] = value.decode('utf-8')
            except:
                # Investigate cases where this fail when reading from nexus
                dictionary[key] = value.decode('utf-8', errors='replace')  # Replaces invalid bytes with '?'

    return dictionary


def load_monitor(metadata, data_folder_name):
    """
    Switches to appropriate loader function
    """

    if "NeXus_field" in metadata.info:
        with h5py.File(os.path.join(data_folder_name, "mccode.h5"), "r", swmr=True) as f:
            return load_monitor_nexus(metadata, f)
    else:
        return load_monitor_text(metadata, data_folder_name)


def load_monitor_nexus(metadata, file_object):
    """
    Function that loads data given metadata and name of data folder
    This version is for a nexus file

    Loads data for single monitor and returns a McStasData object

    Parameters
    ----------

    metadata : McStasMetaData object
        McStasMetaData object corresponding to the monitor to be loaded

    file_object : h5py file object in read mode
    """

    f = file_object

    if "entry1" not in list(f.keys()):
        raise ValueError("h5 file not formatted as expected.")

    if "data" not in list(f["entry1"].keys()):
        raise ValueError("h5 file not formatted as expected.")

    if "simulation" not in list(f["entry1"].keys()):
        raise ValueError("h5 file not formatted as expected.")

    if "Param" not in list(f["entry1"]["simulation"].keys()):
        raise ValueError("h5 file not formatted as expected.")

    NeXus_field = metadata.info["NeXus_field"]

    available_fields = f["entry1"]["data"][NeXus_field].keys()
    if not metadata.dimension == 0 and "events" not in available_fields:
        if "data" not in available_fields:
            raise ValueError("NeXus reading: data not found! \n"
                             + "Monitor metadata:\n" + str(metadata))

        if "errors" not in available_fields:
            raise ValueError("NeXus reading: errors not found! \n"
                             + "Monitor metadata:\n" + str(metadata))

        if "ncount" not in available_fields:
            raise ValueError("NeXus reading: ncount not found! \n"
                             + "Monitor metadata:\n" + str(metadata))

    # Need to check if it is binned data or event data
    if "events" in available_fields:
        Events = np.array(f["entry1"]["data"][NeXus_field]["events"])
        return McStasDataEvent(metadata, Events)

    # Split data into intensity, error and ncount
    if type(metadata.dimension) == int and metadata.dimension == 0:

        if "data" in f["entry1"]["data"][NeXus_field].keys():
            raise ValueError("Found array data in 0D dataset?")

        values = None
        if "values" in f["entry1"]["data"][NeXus_field].keys():
            values = np.array(f["entry1"]["data"][NeXus_field]["values"])

        if metadata.total_I is None:
            if values is not None:
                Intensity = np.array(values[0])
            else:
                raise ValueError("No info on this monitor can be found "
                                 + "in reading of NeXus file "
                                 + str(metadata))
        else:
            Intensity = np.array(metadata.total_I)

        if metadata.total_E is None:
            if values is not None:
                Error = np.array(values[2])
            else:
                Error = np.zeros(1)
        else:
            Error = np.array(metadata.total_E)

        if metadata.total_N is None:
            if values is not None:
                Ncount = np.array(values[3])
            else:
                Ncount = np.zeros(1)
        else:
            Ncount = np.array(metadata.total_N)

        return McStasDataBinned(metadata, Intensity, Error, Ncount)

    elif type(metadata.dimension) == int and metadata.dimension != 0:

        original_xlabel = metadata.info["xlabel"]

        # All special characters are substituted with _ in McStas NeXus file
        x_field = re.sub(r'[^a-zA-Z]', "_", original_xlabel)

        if x_field not in f["entry1"]["data"][NeXus_field].keys():
            error_text = ("Didn't find xaxis in NeXus file. \n"
                          + "Expected this field for x axis: "
                          + str(x_field) + "\n"
                          + "Existing fields: "
                          + str(f["entry1"]["data"][NeXus_field].keys()))

            raise ValueError(error_text)

        xaxis = np.array(f["entry1"]["data"][NeXus_field][x_field])
        Intensity = np.array(f["entry1"]["data"][NeXus_field]["data"])
        Error = np.array(f["entry1"]["data"][NeXus_field]["errors"])
        Ncount = np.array(f["entry1"]["data"][NeXus_field]["ncount"])

        # The data is saved as a McStasDataBinned object
        return McStasDataBinned(metadata, Intensity, Error, Ncount, xaxis=xaxis)

    elif len(metadata.dimension) == 2:
        xaxis = []  # Assume evenly binned in 2d
        Intensity = np.array(f["entry1"]["data"][NeXus_field]["data"]).T
        Error = np.array(f["entry1"]["data"][NeXus_field]["errors"]).T
        Ncount = np.array(f["entry1"]["data"][NeXus_field]["ncount"]).T

        # The data is saved as a McStasDataBinned object
        return McStasDataBinned(metadata, Intensity, Error, Ncount, xaxis=xaxis)
    else:
        raise NameError(
            "Dimension not read correctly in data set "
            + "connected to monitor named "
            + metadata.component_name)


def load_monitor_text(metadata, data_folder_name):
    """
    Function that loads data given metadata and name of data folder
    This version is for a text file

    Loads data for single monitor and returns a McStasData object

    Parameters
    ----------

    metadata : McStasMetaData object
        McStasMetaData object corresponding to the monitor to be loaded

    data_folder_name : str
        path to folder from which metadata should be loaded
    """
    # Load data with numpy
    filename = os.path.join(data_folder_name, metadata.filename.rstrip())
    data = np.loadtxt(filename)

    # Split data into intensity, error and ncount
    if type(metadata.dimension) == int and metadata.dimension == 0:
        Intensity = data.T
        if metadata.total_E is None:
            Error = np.zeros(1)
        else:
            Error = np.array(metadata.total_E)

        if metadata.total_N is None:
            Ncount = np.zeros(1)
        else:
            Ncount = np.array(metadata.total_N)

        return McStasDataBinned(metadata, Intensity, Error, Ncount)

    elif type(metadata.dimension) == int and metadata.dimension != 0:
        xaxis = data.T[0, :]
        Intensity = data.T[1, :]
        Error = data.T[2, :]
        Ncount = data.T[3, :]

        # The data is saved as a McStasDataBinned object
        return McStasDataBinned(metadata, Intensity, Error, Ncount, xaxis=xaxis)

    elif len(metadata.dimension) == 2:
        # Need to check if it is binned data or event data

        with open(filename, 'rb', 0) as file, \
                mmap.mmap(file.fileno(), 0, access=mmap.ACCESS_READ) as s:
            if s.find(b'# Errors') != -1:
                data_type = "Binned"
            else:
                data_type = "Events"

        if data_type == "Events":
            Events = data

            return McStasDataEvent(metadata, Events)

        elif data_type == "Binned":
            # Binned 2D data
            xaxis = []  # Assume evenly binned in 2d
            data_lines = metadata.dimension[1]
            Intensity = data[0:data_lines, :]
            Error = data[data_lines:2 * data_lines, :]
            Ncount = data[2 * data_lines:3 * data_lines, :]

            # The data is saved as a McStasDataBinned object
            return McStasDataBinned(metadata, Intensity, Error, Ncount, xaxis=xaxis)
    else:
        raise NameError(
            "Dimension not read correctly in data set "
            + "connected to monitor named "
            + metadata.component_name)


def print_sim_output(sim_output):
    print(highlight(sim_output, "error", return_section=True, after_lines=10, highlight_type="FAIL"))
    print(highlight(sim_output, "error", return_section=False, highlight_type="FAIL"))


def highlight(string, search_term, return_section=False, highlight_type=None, after_lines=5):
    """
    Highlights search term in string and returns it, if return_section only sections with term is returned
    """

    search_term = search_term.lower()

    if not isinstance(string, str):
        return None

    # Early exit if search term is not in string
    output = string.lower().find(search_term)
    if output == -1:
        if return_section:
            return ""
        else:
            return string

    if return_section:
        instances = list(findall(string, search_term))
        n_instances = len(instances)
        print(f"---- Found {n_instances} places in McStas output with "
              f"keyword '{search_term}'. \n")

    if highlight_type is None:
        highlight_start = ""
        highlight_end = ""
    else:
        if not hasattr(bcolors, highlight_type):
            raise RuntimeError(f"Used highlight_type {highlight_type} "
                               f"in highlight not found in bcolors.")
        else:
            highlight_start = getattr(bcolors, highlight_type)
            highlight_end = bcolors.ENDC

    return_string = ""

    lines = string.split("\n")
    total_lines = len(lines)
    for index, line in enumerate(lines):
        output = line.lower().find(search_term)
        if output == -1:
            if not return_section:
                return_string += line + "\n"
        else:
            replaced_string = line[:output]
            replaced_string += highlight_start
            replaced_string += line[output:output + len(search_term)]
            replaced_string += highlight_end
            replaced_string += line[output + len(search_term):]
            replaced_string += "\n"

            return_string += replaced_string
            if return_section:
                extra_lines = min(total_lines - index, after_lines)
                for line_index in range(1, extra_lines):
                    line_to_include = lines[index + line_index]
                    if line_to_include.lower().find(search_term) != -1:
                        break
                    return_string += line_to_include + "\n"
                return_string += "-"*70 + "\n"

    return return_string


def findall(s, p):
    """
    Yields all the positions of the pattern p in the string s.
    """
    i = s.lower().find(p)
    while i != -1:
        yield i
        i = s.lower().find(p, i+1)