from mcstasscript.helper.mcstas_objects import DeclareVariable
from mcstasscript.helper.mcstas_objects import Component
from mcstasscript.instrument_diagnostics.diagnostics_instrument import DiagnosticsInstrument
from mcstasscript.instrument_diagnostics.view import View
from mcstasscript.instrument_diagnostics.plot_overview import PlotOverview
from mcstasscript.instrument_diagnostics import event_plotter
from mcstasscript.interface.functions import name_search

def sanitise_comp_name(comp_name):
    if isinstance(comp_name, Component):
        comp_name = Component.name

    if not isinstance(comp_name, str):
        raise ValueError("comp_name should be a string or component object!")

    return comp_name

class DiagnosticsPoint:
    def __init__(self, instr, comp_name, before=True, rays=50000):
        comp_name = sanitise_comp_name(comp_name)
        try:
            instr.get_component(comp_name)
        except KeyError:
            raise KeyError("Component '" + comp_name + "' not found in instrument.")

        self.component = comp_name
        self.before = before # False for after

        if isinstance(rays, str):
            if not rays == "all":
                raise ValueError("The ray keyword for a point should be all or an integer.")
        else:
            rays = int(rays)

        self.rays = rays

        # Attributes set by add_monitors
        self.filename = None
        # Attributes set by read_data
        self.recorded_rays = None

    def set_filename(self, filename):
        self.filename = filename

    def set_recorded_rays(self, rays):
        self.recorded_rays = rays

    def __eq__(self, other):
        return self.component == other.component and self.before == other.before

    def __repr__(self):
        string = "Diagnostics point "
        if self.before:
            string += "before: "
        else:
            string += "after:  "

        string += self.component.ljust(25)
        string += " - rays: "
        if self.recorded_rays is not None:
            string += str(self.recorded_rays)
            string += " / "
        string += str(self.rays)

        return string

class BeamDiagnostics(DiagnosticsInstrument):
    def __init__(self, instr):
        super().__init__(instr)

        # points to investigate with options
        self.points = []
        self.ordered_point_list = []

        # flags for monitor_nD user variables
        self.flags = []

        self.monitor_names = []

        self.views = []

        self.data = None

        self.event_plotters = []

    def __repr__(self):
        string = f"Instrument diagnostics for: {self.instr.name}\n"
        string += "Diagnostics points: \n"

        if len(self.points) == 0:
            string += "  No diagnostics points yet\n"
        for point in self.points:
            string += "  " + point.__repr__() + "\n"

        string += "Views: \n"
        if len(self.views) == 0:
            string += "  No views yet\n"

        for view in self.views:
            string += "  " + view.__repr__() + "\n"

        if len(self.flags) != 0:
            string += "Recording following user variables:\n"
        for flag in self.flags:
            string += "  " + str(flag) + "\n"

        if self.data is None:
            string += "Does not yet contain simulated data"
        else:
            string += "Does contain simulated data"

        return string

    def add_flag(self, flag):
        """
        Adds flag to be monitored in event list, maximum of 3 allowed

        The flag must be available in declared variables or user vars.

        Parameters:

        flag: str
            str with variable name to be recorded to event dataset
        """
        if not isinstance(flag, str):
            raise ValueError("flag has to be a string.")

        # find declare and user_var names in instrument
        names = [x.name for x in self.instr.declare_list
                 if isinstance(x, DeclareVariable)]
        names += [x.name for x in self.instr.user_var_list
                  if isinstance(x, DeclareVariable)]

        if flag not in names:
            raise ValueError(f"flag {flag} was not found in declared "
                             f"variables or uservars of the instrument.")

        if len(self.flags) >= 3:
            raise ValueError("Can't add more than three flags")

        self.flags.append(flag)

    def clear_flags(self):
        """
        Clears flags
        """
        self.flags = []

    def show_flags(self):
        """
        Shows list of current flags
        """
        if self.flags is not None:
            for index, flag in enumerate(self.flags):
                print(f" user{index+1}={flag}")

    def add_point(self, before=None, after=None, rays=50000):
        """
        Adds point in which the beam is investigated with given number of rays

        Can be before and/or after components, this is chosen by using the
        appropriate keyword argument.

        Parameters:

        before : str or component instance
            Reference to component beam should be monitored just before

        after : str or component instance
            Reference to component beam should be monitored just after

        rays : int or "all"
            Number of rays to record, either integer or the string "all"
        """
        if before is None and after is None:
            raise ValueError("Specify either before or after with component name.")

        if before is not None:
            self._add_diagnostics_point(comp_name=before, before=True, rays=rays)

        if after is not None:
            self._add_diagnostics_point(comp_name=after, before=False, rays=rays)

    def _add_diagnostics_point(self, comp_name, before, rays):
        """
        Internal function to add the DiagnosticsPoint to list
        """
        point = DiagnosticsPoint(self.instr, comp_name=comp_name, before=before, rays=rays)
        if point in self.points:
            # overwrite that point
            index = self.points.index(point)
            self.points[index] = point
        else:
            # add point to list
            self.points.append(point)

    def remove_point(self, before=None, after=None):
        """
        Removes a point from the list

        Specify whether the point in question is before or after the
        given component name or instance by using the appropriate
        keyword argument

        Parameter:

        before : str or component instance
            Name or component instance for which point before should be removed

        after : str or component instance
            Name or component instance for which point after should be removed
        """
        if before is None and after is None:
            raise ValueError("Specify either before or after with component name.")

        if before is not None:
            before = sanitise_comp_name(before)
            for index, point in enumerate(self.points):
                if point.component == before and point.before:
                    del self.points[index]

        if after is not None:
            after = sanitise_comp_name(after)
            for index, point in enumerate(self.points):
                if point.component == after and not point.before:
                    del self.points[index]

    def clear_points(self):
        """
        Remove all points
        """
        self.points = []

    def show_points(self):
        """
        Shows current diagnostics points
        """
        for point in self.points:
            print(point)

    def add_view(self, axis1, axis2=None, bins=100, same_scale=False, **kwargs):
        """
        Add a view with one or two axes that will be shown at all diagnostics points

        Parameters:
        axis1: str
            Name of parameter for first axis

        axis2: str
            Name of parameter for second axis

        bins : int or list of length 2
            Number of bins for histogram (can be list of length 2 for 2D)

        same_scale : bool
            Select if the same scale should be ensured for all points with this view
        """
        view = View(axis1=axis1, axis2=axis2, bins=bins, same_scale=same_scale, **kwargs)
        self.views.append(view)

    def clear_views(self):
        """
        Clear list of views
        """
        self.views = []

    def show_views(self):
        """
        Shows views
        """
        for view in self.views:
            print(view)

    def add_monitors(self):
        """
        Adds monitors as specified by list of diagnostics points

        The method sorts the points with the component sequence so the
        resulting plots are in the correct order.
        """
        comp_names = [x.name for x in self.instr.make_component_subset()]
        self.monitor_names = []
        self.ordered_point_list = []
        point_names = [x.component for x in self.points]

        for name in comp_names:
            before_point = None
            after_point = None
            if name in point_names:
                for point in self.points:
                    if point.component == name:
                        if point.before:
                            before_point = point
                        else:
                            after_point = point

            if before_point is not None:
                self.ordered_point_list.append(before_point)
            if after_point is not None:
                self.ordered_point_list.append(after_point)

        for point in self.ordered_point_list:
            comp = self.instr.get_component(point.component)
            filename = self.add_monitor(point, comp)
            point.set_filename(filename)

    def add_monitor(self, point, comp_instance):
        """
        Adds monitor_nD event monitor to diagnostics instrument for given point

        Needs point and component instance that match

        Parameters:

        point : DiagnosticsPoint
            Point for which monitor should be added

        comp_instance : Component
            Component instance relative to which the monitor should be placed
        """

        if not point.component == comp_instance.name:
            raise RuntimeError("Given point and component instance should match!")

        user_vars = [None, None, None]
        flags_str = ""
        for index, flag in enumerate(self.flags):
            flags_str += " user" + str(index+1)
            if self.instr.mccode_version == 2:
                # Monitor_nD in McStas 2.X needs to be a raw variable
                user_vars[index] = flag
            else:
                # Monitor_nD in McStas 3.X (and onwards) needs to be a string
                user_vars[index] = '"' + flag + '"'

        if isinstance(point.rays, (float, int)):
            ray_value = point.rays
            if self.instr.mccode_version == 3:
                ray_value += 1 # In McStas 3.0 Monitor_nD grabs one event too few
            ray_str = str(int(ray_value)) # Monitor_nD seems to record one ray less than requested
        else:
            ray_str = point.rays # Case of rays set to all

        if point.before:
            name = "Diag_before_" + point.component
            options = f'"square boarders n x y z vx vy vz t{flags_str}, list {ray_str}"'
            mon = self.instr.add_component(name, "Monitor_nD", before=point.component)
        else:
            name = "Diag_after_" + point.component
            options = f'"previous n x y z vx vy vz t{flags_str}, list {ray_str}"'
            mon = self.instr.add_component(name, "Monitor_nD", after=point.component)

        mon.set_parameters(restore_neutron=1,
                           xwidth=100, yheight=100,
                           options=options,
                           user1=user_vars[0], user2=user_vars[1], user3=user_vars[2],
                           filename='"' + name + ".diag" + '"')


        mon.set_AT(comp_instance.AT_data, RELATIVE=comp_instance.AT_reference)
        if comp_instance.ROTATED_specified:
            mon.set_ROTATED(comp_instance.ROTATED_data, RELATIVE=comp_instance.ROTATED_reference)

        return name

    def run(self):
        """
        Runs diagnostics with all included points

        Saves data in data attribute and the resulting instrument can be found
        in the instr attribute
        """
        self.reset_instr()
        self.remove_previous_use()

        self.add_monitors()

        self.correct_target_index()

        self.data = self.instr.backengine()

        self.read_data()

    def read_data(self):
        """
        Reads the generated data and organizes it as event plotter instances
        """

        self.event_plotters = []

        for point in self.ordered_point_list:
            try:
                event_data = name_search(point.filename, self.data)
            except NameError:
                # Data was not generated, no neutrons reached it or simulation failed
                continue

            point.set_recorded_rays(event_data.metadata.total_N)
            plotter = event_plotter.EventPlotter(point.filename, event_data,
                                                 flag_info=self.flags)

            self.event_plotters.append(plotter)

    def plot(self):
        """
        Plots the generated data for all points with all views
        """

        if len(self.event_plotters) == 0:
            if len(self.views) == 0:
                print("No data to plot! Add views and run.")
            else:
                print("No data to plot! Use the run method to generate data.")

        overview = PlotOverview(self.event_plotters, self.views)
        overview.plot_all()