# simul_sweep.py
import copy
import math
import time
from functools import partial

from PyQt5.QtCore import QObject, pyqtSlot

from ..tools.util import _autorange_srs, safe_get
from .base_sweep import BaseSweep
from .progress import SweepState


class SimulSweep(BaseSweep, QObject):
    """Child of BaseSweep used to simultaneously run sweeps on multiple parameters.

    Parameters must be split into equivalent step intervals. The first key in
    the '_p' dictionary is used as the independent variable passed to BaseSweep
    in the same manner as the Sweep1D 'set_param'.

    Attributes:
    ---------
    _p:
        Dictionary used to pass parameters and their associated information.
    n_steps:
        Integer value of number of steps for each parameter to take, instead of offering step sizes in _p
    bidirectional:
        Set to True to run the sweep in both directions.
    continuous:
        Set to True to keep sweep running indefinitely.
    *args:
        Optional sweep arguments to be passed in BaseSweep.
    **kwargs:
        Optional keyword arguments to be passed in BaseSweep.
    simul_params:
        Used to store the parameters from the input '_p' dictionary.
    set_params_dict:
        Stores input parameter dictionary.
    direction:
        Tells sweep which direction to run; either 0 or 1.
    ramp_sweep:
        Iterates parameters to their starting values before sweep begins.
    runner:
        Assigns the desired Runner Thread.
    plotter:
        Assigns the desired Plotter Thread.
    err:
        Tolerance for considering rounding errors when determining when the sweep has finished.

    Methods:
    ---------
    start(persist_data, ramp_to_start, ramp_multiplier)
        Begins a BaseSweep, creates necessary measurement objects and threads.
    stop()
        Stops the BaseSweep.
    kill()
        Ends all threads and closes any active plots.
    send_updates(no_sp)
        Passed in this class.
    step_param()
        Iterates each parameter based on step size.
    update_values()
        Returns updated parameter-value pairs, default parameter is time.
    ramp_to_zero(params=None)
        Ramps value of all parameters to zero.
    ramp_to(vals_dict, start_on_finish=False, persist=None, multiplier=1)
        Begins ramping parameters to assigned starting values.
    done_ramping(self, vals_dict, start_on_finish=False, pd=None)
        Ensures that each parameter is at its start value and ends ramp.
    flip_direction()
        Changes the direction of the sweep.
    """

    def __init__(
        self,
        _p,
        n_steps=None,
        err=0.1,
        bidirectional=False,
        continual=False,
        *args,
        **kwargs,
    ):
        if len(_p.keys()) < 1 or not all(isinstance(p, dict) for p in _p.values()):
            raise ValueError(
                "Must pass at least one Parameter and the associated values as dictionaries."
            )

        self.simul_params = []
        # Deep copy value dicts to avoid mutating caller's original dictionaries
        self.set_params_dict = {p: copy.deepcopy(v) for p, v in _p.items()}
        self.direction = 0
        self.n_steps = n_steps
        self.ramp_sweep = None
        self.err = err

        # Take the first parameter, and set it as the normal Sweep1D set param
        sp = list(_p.keys())[0]
        # Force time to be on the x axis
        kwargs["x_axis_time"] = 1
        QObject.__init__(self)
        BaseSweep.__init__(self, set_param=sp, *args, **kwargs)

        # Validate sweep ranges for all parameters against their validator bounds
        for p, v in _p.items():
            self._validate_param_sweep_range(
                p, v["start"], v["stop"],
                param_label=getattr(p, "label", p.name)
            )

        self.bidirectional = bidirectional
        self.continuous = continual

        if n_steps is not None:
            for p, v in self.set_params_dict.items():
                v["step"] = (v["stop"] - v["start"]) / n_steps
        else:
            _n_steps = []
            for key, p in _p.items():
                _n_steps.append(int(abs(p["stop"] - p["start"]) / abs(p["step"])))

            if not all(steps == _n_steps[0] for steps in _n_steps):
                raise ValueError(
                    "Parameters have a different number of steps for the sweep. The Parameters must have "
                    "the same number of steps to sweep them simultaneously."
                    f"\nStep numbers: {_n_steps}"
                )
            self.n_steps = _n_steps[0]

        for p, v in self.set_params_dict.items():
            self.simul_params.append(p)

            # Store start as snap origin. Updated on direction flips to ensure
            # the step grid always includes the current sweep's starting point.
            v["_snap_origin"] = v["start"]

            # Make sure the step is in the right direction
            if (v["stop"] - v["start"]) > 0:
                v["step"] = abs(v["step"])
            else:
                v["step"] = (-1) * abs(v["step"])

            v["setpoint"] = safe_get(p) - v["step"]
            v["setpoint"] = self._snap_to_step(v["setpoint"], v["_snap_origin"], v["step"])

        for p, v in self.set_params_dict.items():
            self.emit_step_info(
                p.label,
                v["start"],
                v["stop"],
                v["step"],
                getattr(p, "unit", None),
            )

        self.follow_param([p for p in self.simul_params if p is not self.set_param])
        self.persist_data = []
        if not self.continuous:
            self.progressState.progress = 0.0
            self.update_progress()

    def __str__(self):
        p_desc = ""
        n_params = len(self.simul_params)

        for n, item in enumerate(self.set_params_dict.items()):
            p, v = item
            if n < n_params - 2:
                p_desc += f"{p.label} from {v['start']} to {v['stop']}, with step {v['step']}, "
            elif n == n_params - 2:
                p_desc += f"{p.label} from {v['start']} to {v['stop']}, with step {v['step']}, and "
            elif n == n_params - 1:
                p_desc += f"{p.label} from {v['start']} to {v['stop']}, with step {v['step']}."

        return f"SimulSweep of {p_desc}"

    def __repr__(self):
        return f"SimulSweep({[(p.label, v['start'], v['stop'], v['step']) for p, v in self.set_params_dict.items()]})"

    def start(self, persist_data=None, ramp_to_start=True, ramp_multiplier=1):
        """Starts the sweep. Runs from the BaseSweep start() function."""
        if self.progressState.state == SweepState.RAMPING:
            self.print_main.emit(
                "Still ramping. Wait until ramp is done to start the sweep."
            )
            return
        if self.progressState.state == SweepState.RUNNING:
            self.print_main.emit("Sweep is already running.")
            return

        if ramp_to_start is True:
            self.print_main.emit("Ramping to our starting setpoints.")
            vals_dict = {}
            for p in self.simul_params:
                vals_dict[p] = self.set_params_dict[p]["start"]
            self.ramp_to(
                vals_dict,
                start_on_finish=True,
                persist=persist_data,
                multiplier=ramp_multiplier,
            )
        else:
            self.print_main.emit("Starting our sweep.")
            super().start(persist_data=persist_data, ramp_to_start=False)

    def pause(self):
        if self.progressState.state == SweepState.RAMPING and self.ramp_sweep is not None:
            self.print_main.emit("Stopping the ramp.")
            self.ramp_sweep.pause()
            self.ramp_sweep.kill()

            while self.ramp_sweep.check_running():
                time.sleep(0.2)
            vals_dict = {}
            for p, v in self.ramp_sweep.set_params_dict.items():
                vals_dict[p] = v["setpoint"]
            self.done_ramping(vals_dict)

        super().pause()

    def kill(self):
        # Mark killed immediately to short-circuit any pending callbacks (e.g., done_ramping)
        # ERROR state transitions to KILLED since user explicitly called kill()
        # Use getattr to handle case where __init__ failed before progressState was set
        progress_state = getattr(self, "progressState", None)
        if progress_state is not None and progress_state.state not in (SweepState.DONE, SweepState.KILLED):
            self.progressState.state = SweepState.KILLED

        # Always tear down an active ramp sweep, even if we are no longer in the RAMPING state.
        # This prevents its runner QThread from outliving the parent sweep.
        # Use getattr to handle case where __init__ failed before ramp_sweep was set
        if getattr(self, "ramp_sweep", None) is not None:
            try:
                self.ramp_sweep.pause()
            except Exception:
                pass
            try:
                self.ramp_sweep.kill()
            except Exception:
                pass
            # Wait briefly for the ramp runner to finish to avoid Qt teardown races
            try:
                runner = getattr(self.ramp_sweep, "runner", None)
                if runner is not None and runner.isRunning():
                    runner.wait(500)
            except Exception:
                pass
            self.ramp_sweep = None

        super().kill()

    def send_updates(self, no_sp=True):
        pass

    def step_param(self):
        """Iterates the parameter."""
        rets = []
        ending = False
        self.persist_data = []
        start_dir = self.direction

        for p, v in self.set_params_dict.items():
            # Check if we should continue (not at end yet)
            at_end = not (
                abs(v["setpoint"] - v["stop"]) - abs(v["step"] / 2)
                > abs(v["step"]) * self.err
            )

            if not at_end:
                next_setpoint = v["setpoint"] + v["step"]
                next_setpoint = self._snap_to_step(next_setpoint, v["_snap_origin"], v["step"])

                # If next setpoint exceeds bounds, treat as end of sweep
                sweep_min = min(v["start"], v["stop"])
                sweep_max = max(v["start"], v["stop"])
                if next_setpoint < sweep_min or next_setpoint > sweep_max:
                    at_end = True
                else:
                    v["setpoint"] = next_setpoint
                    if not self.try_set(p, v["setpoint"]):
                        return None
                    rets.append((p, v["setpoint"]))

            # End of sweep handling for this parameter
            if at_end:
                if self.bidirectional and self.direction == 0:
                    self.flip_direction()
                    return self.step_param()
                elif self.continuous and self.direction == start_dir:
                    self.flip_direction()
                    return self.step_param()
                else:
                    ending = True

        if ending is True:
            self.print_main.emit("Done with the sweep!")
            for p, v in self.set_params_dict.items():
                self.print_main.emit(f"{p.label} = {v['setpoint']} ({p.unit})")
            if self.bidirectional or self.continuous:
                self.flip_direction()
            return None

        return rets

    def update_values(self):
        """Iterates our data points, changing our setpoint if we are sweeping, and refreshing
        the values of all our followed parameters. If we are saving data, it happens here,
        and the data is returned.

        Returns:
            data - A list of tuples with the new data. Each tuple is of the format
                   (<QCoDeS Parameter>, measurement value). The tuples are passed in order of
                   time, then set_param (if applicable), then all the followed params.
        """
        t = time.monotonic() - self.t0

        data = [("time", t)]

        sp_data = self.step_param()
        if sp_data is not None:
            data += sp_data
        else:
            self.mark_done()
            return None

        for i, (l, _, gain) in enumerate(self._srs):
            _autorange_srs(l, 3)

        for i, p in enumerate(self._params):
            if p not in self.simul_params:
                v = safe_get(p)
                data.append((p, v))

        if self.save_data and self.progressState.state == SweepState.RUNNING:
            self.runner.datasaver.add_result(*data)

        # self.send_updates()

        return data

    def ramp_to_zero(self, params=None):
        if params is None:
            params = self.simul_params
        vals_dict = {}
        for p in params:
            vals_dict[p] = 0
        self.ramp_to(vals_dict)

    def ramp_to(self, vals_dict, start_on_finish=False, persist=None, multiplier=1):
        # Ensure we aren't currently running
        if self.progressState.state == SweepState.RAMPING:
            self.print_main.emit(
                "Currently ramping. Finish current ramp before starting another."
            )
            return
        if self.progressState.state == SweepState.RUNNING:
            self.print_main.emit("Already running. Stop the sweep before ramping.")
            return

        ramp_params_dict = {}
        n_steps = 1
        for p, v in vals_dict.items():
            if p not in self.set_params_dict.keys():
                self.print_main.emit("Cannot ramp parameter not in our sweep.")
                return

            p_step = self.set_params_dict[p]["step"]
            if abs(v - safe_get(p)) - abs(p_step / 2) < abs(p_step) * self.err:
                continue

            ramp_params_dict[p] = {}
            ramp_params_dict[p]["start"] = safe_get(p)
            ramp_params_dict[p]["stop"] = v

            p_steps = abs(
                (ramp_params_dict[p]["stop"] - ramp_params_dict[p]["start"])
                / p_step
                * multiplier
            )
            if p_steps > n_steps:
                n_steps = math.ceil(p_steps)

        if len(ramp_params_dict.keys()) == 0:
            self.print_main.emit("Already at the values, no ramp needed.")
            self.done_ramping(vals_dict, start_on_finish=start_on_finish, pd=persist)
            return

        # Store ramp steps for validation in done_ramping
        self.ramp_steps = {}
        for p, v in ramp_params_dict.items():
            v["step"] = (v["stop"] - v["start"]) / n_steps
            self.ramp_steps[p] = v["step"]

        self.ramp_sweep = SimulSweep(
            ramp_params_dict,
            n_steps=n_steps,
            inter_delay=self.inter_delay,
            # Only spawn plotter windows if the parent sweep is plotting
            plot_data=self.plot_data,
            save_data=False,
            complete_func=partial(
                self.done_ramping,
                vals_dict,
                start_on_finish=start_on_finish,
                pd=persist,
            ),
            suppress_output=self.suppress_output,
        )
        # Only follow parameters that are not being ramped to avoid circular dependencies
        follow_params = [p for p in self._params if p not in ramp_params_dict.keys()]
        if follow_params:
            self.ramp_sweep.follow_param(follow_params)
        self.progressState.state = SweepState.RAMPING
        self.ramp_sweep.start(ramp_to_start=False)

    @pyqtSlot()
    def done_ramping(self, vals_dict, start_on_finish=False, pd=None):
        # Check if the ramp sweep itself failed with an error
        if self.ramp_sweep is not None and self.ramp_sweep.progressState.state == SweepState.ERROR:
            error_msg = self.ramp_sweep.progressState.error_message or "Ramp sweep failed"
            self.print_main.emit(f"Ramp sweep failed: {error_msg}")
            self.ramp_sweep.kill()
            self.ramp_sweep = None
            self.mark_error(f"Ramp to start failed: {error_msg}")
            return

        if self.progressState.state != SweepState.KILLED:
            self.progressState.state = SweepState.READY

        # Check if we are at the value we expect, otherwise something went wrong with the ramp
        for p, v in vals_dict.items():
            # Use ramp step if available (for parameters that were ramped),
            # otherwise use sweep step (for parameters that were already at target)
            p_step = getattr(self, 'ramp_steps', {}).get(p, self.set_params_dict[p]["step"])
            actual_value = safe_get(p)
            position_error = abs(actual_value - v) - abs(p_step / 2)
            tolerance = abs(p_step) * self.err
            if position_error > tolerance:
                error_msg = (
                    f"Ramping failed (possible that the direction was changed while ramping). "
                    f"Expected {p.label} final value: {v}. Actual value: {actual_value}. "
                    f"Error: {position_error:.6g}, Tolerance: {tolerance:.6g} (err={self.err}). "
                    f"If tolerance is too tight, consider increasing 'err' parameter."
                )
                self.print_main.emit(error_msg)

                if self.ramp_sweep is not None:
                    self.ramp_sweep.kill()
                    self.ramp_sweep = None

                self.mark_error(error_msg)
                return

        self.print_main.emit("Done ramping!")
        for p, v in vals_dict.items():
            if not self.try_set(p, v):
                if self.ramp_sweep is not None:
                    self.ramp_sweep.kill()
                    self.ramp_sweep = None
                return
            self.set_params_dict[p]["setpoint"] = v - self.set_params_dict[p]["step"]
            self.set_params_dict[p]["setpoint"] = self._snap_to_step(
                self.set_params_dict[p]["setpoint"],
                self.set_params_dict[p]["_snap_origin"],
                self.set_params_dict[p]["step"]
            )

        if self.ramp_sweep is not None:
            self.ramp_sweep.kill()
            self.ramp_sweep = None

        # Clear ramp steps after successful ramp
        if hasattr(self, 'ramp_steps'):
            del self.ramp_steps

        if start_on_finish is True:
            self.start(ramp_to_start=False, persist_data=pd)

    def flip_direction(self):
        """Flips the direction of the sweep, to do bidirectional sweeping."""
        # If backwards, go forwards, and vice versa
        if self.direction:
            self.direction = 0
            for p, v in self.set_params_dict.items():
                temp = v["start"]
                v["start"] = v["stop"]
                v["stop"] = temp
                v["step"] = -1 * v["step"] / self.back_multiplier
                v["_snap_origin"] = v["start"]
                v["setpoint"] -= v["step"]
                v["setpoint"] = self._snap_to_step(v["setpoint"], v["_snap_origin"], v["step"])
        else:
            self.direction = 1
            for p, v in self.set_params_dict.items():
                temp = v["start"]
                v["start"] = v["stop"]
                v["stop"] = temp
                v["step"] = -1 * v["step"] * self.back_multiplier
                v["_snap_origin"] = v["start"]
                v["setpoint"] -= v["step"]
                v["setpoint"] = self._snap_to_step(v["setpoint"], v["_snap_origin"], v["step"])

        if self.plot_data is True and self.plotter is not None:
            self.plotter.add_break(self.direction)

    def estimate_time(self, verbose=True):
        """Returns an estimate of the amount of time the sweep will take to complete.

        Parameters
        ----------
        verbose:
            Controls whether the function will print out the estimate in the form hh:mm:ss (default True)

        Returns:
        -------
        Time estimate for the sweep, in seconds
        """

        if self.progressState.state == SweepState.DONE:
            remaining = 0
        else:
            cfg = list(self.set_params_dict.values())[0]
            step_mag = abs(cfg["step"])
            current = cfg["setpoint"]
            target = cfg["stop"]
            distance = abs(target - current)
            steps = max(distance / step_mag, 0)
            remaining = steps * self.inter_delay

            effective_back_multiplier = self.back_multiplier if self.back_multiplier not in (None, 0) else 1.0

            if (
                self.bidirectional
                and effective_back_multiplier > 0
                and (
                    self.direction == 0
                    or (
                        self.t0 == 0
                        and self.progressState.state != SweepState.RUNNING
                    )
                )
            ):
                remaining += self.n_steps / effective_back_multiplier * self.inter_delay

        if verbose:
            hours, minutes, seconds = self._split_hms(remaining)
            self.print_main.emit(
                f"Estimated time remaining for {repr(self)}: {hours}h:{minutes:02d}m:{seconds:02d}s"
            )

        return remaining

    # --- JSON export/import hooks ---
    def _export_json_specific(self, json_dict: dict) -> dict:
        json_dict["attributes"]["bidirectional"] = self.bidirectional
        json_dict["attributes"]["continual"] = self.continuous
        json_dict["set_params"] = {}
        for p, items in self.set_params_dict.items():
            key = f"{p.instrument.name}.{p.name}"
            entry = dict(items)
            entry["instr_module"] = p.instrument.__class__.__module__
            entry["instr_class"] = p.instrument.__class__.__name__
            entry["instr_name"] = p.instrument.name
            json_dict["set_params"][key] = entry
        return json_dict

    @classmethod
    def from_json(cls, json_dict, station):
        attrs = dict(json_dict.get("attributes", {}))
        # Pull SimulSweep-specific attributes but keep BaseSweep kwargs
        bidirectional = attrs.pop("bidirectional", False)
        continual = attrs.pop("continual", False)

        set_params_dict = {}
        for key, items in json_dict.get("set_params", {}).items():
            param_name = key.split(".", 1)[1] if "." in key else key
            p = BaseSweep._load_parameter_by_type(
                param_name,
                items["instr_name"],
                items["instr_module"],
                items["instr_class"],
                station,
            )
            set_params_dict[p] = {
                "start": items["start"],
                "stop": items["stop"],
                "step": items["step"],
            }

        return cls(
            set_params_dict, bidirectional=bidirectional, continual=continual, **attrs
        )

    def _params_to_exclude_from_follow(self) -> set:
        """Exclude all simultaneously-swept params from follow_params."""
        try:
            return {
                f"{p.instrument.name}.{p.name}" for p in self.set_params_dict.keys()
            }
        except Exception:
            return set()
