# sweep1d.py

import time

from PyQt5.QtCore import QObject, pyqtSlot

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

try:
    from qcodes.instrument_drivers.american_magnetics import AMIModel430 as AMI430
except ImportError:  # pragma: no cover - fallback for older QCoDeS versions
    from qcodes.instrument_drivers.american_magnetics.AMI430 import AMI430
from functools import partial

from qcodes_contrib_drivers.drivers.OxfordInstruments.IPS120 import (
    OxfordInstruments_IPS120,
)

from ..Drivers.M4G import M4G


class Sweep1D(BaseSweep, QObject):
    """Class extending BaseSweep to sweep one independent parameter.

    The independent parameter defaults to time until a new parameter is
    assigned to 'set_param'. Sweep1D will sweep this parameter to its stop
    value (and back) while measuring any changes in the followed parameters.
    The data will be live-plotted and saved to a database if desired.

    Attributes:
    ---------
    set_param:
        QCoDeS Parameter to be swept.
    begin:
        The starting value of the sweeping parameter.
    end:
        The ending value of the sweeping parameter.
    step:
        The step size between measured datapoints.
    bidirectional:
        Determines whether or not to run the sweep in both directions.
    runner:
        Assigns the Runner Thread.
    plotter:
        Assigns the Plotter Thread.
    datasaver:
        Initiated by Runner Thread to enable saving and export of data.
    inter_delay:
        Time (in seconds) to wait between data points.
    save_data:
        Flag used to determine if the data should be saved or not.
    plot_data:
        Flag to determine whether or not to live-plot data.
    complete_func:
        Sets a function to be executed upon completion of the sweep.
    x_axis_time:
        Defaults to 0 to allow the 'set_param' to be plotted on the x-axis.
    parent:
        Sets a parent Sweep2D object.
    continual:
        Causes sweep to continuously run back and forth between start and stop points.
    plot_bin:
        Sets the number of data points taken between updates of the plot. Defaults to 1.
    back_multiplier:
        Factor to scale the step size after flipping directions.
    setpoint:
        The first parameter value where a measurement is obtained after starting.
    direction:
        Either 0 or 1 to indicate the direction of the sweep.
    ramp_sweep:
        Defines the sweep used to set the parameter to its starting value.
    instrument:
        The instrument used to control/monitor the 'set_param'.

    Methods:
    ---------
    start(persist_data=None, ramp_to_start=True, ramp_multiplier=1)
        Starts the sweep; runs from the BaseSweep start() function.
    pause()
        Pauses any currently active sweeps.
    kill()
        Ends the threads spawned by the sweep and closes any active plots.
    step_param()
        Iterates the parameter.
    step_AMI430()
        Used to control sweeps of AMI430 Instrument.
    step_IPS120()
        Used to control sweeps of IPS120 Instrument.
    flip_direction()
        Flips the direction of the sweep, to do bidirectional sweeping.
    ramp_to(value, start_on_finish=False, persist=None, multiplier=1)
        Ramps the set_param to a given value, at a rate dictated by the multiplier.
    ramp_to_zero()
        Ramps the set_param to 0, at the same rate as already specified.
    done_ramping(self, value, start_on_finish=False, pd=None)
        Alerts the sweep that the ramp is finished.
    get_param_setpoint()
        Obtains the current value of the setpoint.
    reset(new_params=None)
        Resets Sweep1D, optional argument to change sweep details.
    """

    def __init__(
        self,
        set_param,
        start,
        stop,
        step,
        bidirectional=False,
        continual=False,
        x_axis_time=0,
        err=1e-2,
        *args,
        **kwargs,
    ):
        """Initializes the sweep.

        There are 6 additional parameters not included in the BaseSweep initialization.

        Parameters
        ---------
        set param:
            The independent parameter to be swept.
        start:
            The value that the sweep will begin from.
        stop:
            The value that the sweep will end at.
        step:
            The step spacing for each measurement.
        x_axis_time:
            Determines if followed parameters are plotted against time(1) or 'set_param'(0).
        err:
            Tolerance for considering rounding errors when determining when the sweep has finished.
        """
        # Initialize the BaseSweep
        QObject.__init__(self)
        BaseSweep.__init__(
            self, set_param=set_param, x_axis_time=x_axis_time, *args, **kwargs
        )

        # Validate sweep range against parameter's validator bounds
        self._validate_param_sweep_range(set_param, start, stop)

        self.begin = start
        self.end = stop
        self.step = step

        # Store the start for snapping. Updated on direction flips to ensure
        # the step grid always includes the current sweep's starting point.
        self._snap_origin = start

        # Make sure the step is in the right direction
        if (self.end - self.begin) > 0:
            self.step = abs(self.step)
        else:
            self.step = (-1) * abs(self.step)

        self.setpoint = self.begin - self.step
        self.setpoint = self._snap_to_step(self.setpoint, self._snap_origin, self.step)
        self.bidirectional = bidirectional
        self.continuous = continual
        self.direction = 0
        self.ramp_sweep = None
        self.instrument = self.set_param.instrument
        self.err = err

        self.emit_step_info(
            self.set_param.label,
            self.begin,
            self.end,
            self.step,
            getattr(self.set_param, "unit", None),
        )

        self.magnet_initialized = False
        self._ami_completion_pending = False
        self._m4g_completion_pending = False
        if not self.continuous:
            self.progressState.progress = 0.0
            self.update_progress()

    def __str__(self):
        return f"1D Sweep of {self.set_param.label} from {self.begin} to {self.end}, with step size {self.step}."

    def __repr__(self):
        return f"Sweep1D({self.set_param.label}, {self.begin}, {self.end}, {self.step})"

    def start(self, persist_data=None, ramp_to_start=True, ramp_multiplier=1):
        """Starts the sweep; runs from the BaseSweep start() function.

        Parameters
        ---------
        persist_data:
            Used to set Sweep2D parameter.
        ramp_to_start:
            Gently sweeps parameter to desired starting value.
        ramp_multiplier:
            Factor to control ramping speed compared to sweep speed.
        """
        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

        self._ami_completion_pending = False
        self._m4g_completion_pending = False

        if (
            ramp_to_start is True
            and not isinstance(self.instrument, OxfordInstruments_IPS120)
            and not isinstance(self.instrument, AMI430)
            and not isinstance(self.instrument, M4G)
        ):
            self.print_main.emit(
                f"Ramping to our starting setpoint value of {self.begin} ({self.set_param.unit})"
            )
            self.ramp_to(
                self.begin,
                start_on_finish=True,
                persist=persist_data,
                multiplier=ramp_multiplier,
            )
        else:
            self.print_main.emit(
                f"Sweeping {self.set_param.label} to {self.end} ({self.set_param.unit})"
            )
            BaseSweep.start(self, persist_data=persist_data)

    def pause(self):
        """Pauses any currently active sweeps."""
        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)
            self.done_ramping(self.ramp_sweep.setpoint)
            # self.setpoint=self.ramp_sweep.setpoint
            # self.ramp_sweep.plotter.clear()
            # self.ramp_sweep = None
            # self.progressState.state = SweepState.READY
            # self.print_main.emit(f"Stopped the ramp, the current setpoint  is {self.setpoint} {self.set_param.unit}")

        BaseSweep.pause(self)

        if isinstance(self.instrument, AMI430):
            self.instrument.pause()
            self.magnet_initialized = False
        elif isinstance(self.instrument, OxfordInstruments_IPS120):
            self.try_set(self.instrument.activity, 0)
            self.magnet_initialized = False
        elif isinstance(self.instrument, M4G):
            self.instrument.write("SWEEP PAUSE")

    def stop(self):
        """Stop/pause the sweep. Alias for pause() for backward compatibility.

        Handles stopping of ramp sweeps if in ramping state and instrument-specific pause operations.
        """
        self.pause()

    def kill(self):
        """Ends the threads spawned by the sweep and closes any active plots."""
        # Always tear down active ramp sweep, regardless of state
        # (error may have occurred during ramping, leaving ramp_sweep active)
        # 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
            self.ramp_sweep = None

        BaseSweep.kill(self)

    def step_param(self):
        """Iterates the parameter and checks for our stop condition.

        Returns:
        ---------
        Data pair in form of (<set param>, <setpoint>), or None if we have reached the end of our sweep.
        """
        # The AMI Magnet sweeps very slowly, so we implement sweeps of it  differently
        # If we are sweeping the magnet, let's deal with it here
        if isinstance(self.instrument, AMI430) and "field" in self.set_param.full_name:
            return self.step_AMI430()
        elif isinstance(self.instrument, M4G) and "field" in self.set_param.full_name:
            return self.step_M4G()

        # Check if we should continue (not at end yet)
        at_end = not (
            abs(self.setpoint - self.end) - abs(self.step / 2)
            > abs(self.step) * self.err
        )

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

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

        # End of sweep handling
        if self.continuous:
            self.flip_direction()
            return self.step_param()
        elif self.bidirectional and self.direction == 0:
            self.flip_direction()
            return self.step_param()
        # If neither of the above are triggered, it means we are at the end of the sweep
        else:
            self.print_main.emit(
                f"Finished the sweep! {self.set_param.label} = {safe_get(self.set_param)} "
                f"({self.set_param.unit})"
            )
            self.flip_direction()
            return None

    def step_AMI430(self):
        """Used to control sweeps of AMI430 Instrument.

        The endpoint is determined prior to the sweep, and the current
        field is measured during ramping.

        Returns:
        ---------
        The parameter-value pair that was measured.
        """
        if self._ami_completion_pending:
            self._ami_completion_pending = False
            return None

        # Check if we have set the magnetic field yet
        if self.magnet_initialized is False:
            self.instrument.set_field(self.end, block=False)
            self.magnet_initialized = True
            time.sleep(self.inter_delay)

            while safe_get(self.instrument.ramping_state) != "ramping":
                time.sleep(self.inter_delay)

        # Grab our data
        dt = safe_get(self.set_param)
        data_pair = (self.set_param, dt)
        self.setpoint = dt

        # Check our stop conditions- being at the end point
        if safe_get(self.instrument.ramping_state) == "holding":
            self.magnet_initialized = False
            self.print_main.emit(
                f"Done with the sweep, {self.set_param.label}={self.set_param.get()} "
                f"({self.set_param.unit})"
            )
            self.flip_direction()
            self._ami_completion_pending = True

        # Return our data pair, just like any other sweep
        return [data_pair]

    def step_M4G(self, _retry_count=0):
        """Function to deal with sweeps of Cryomagnetics M4G Magnet Supply. Instead of setting intermediate points, we set the endpoint at the
        beginning, then ask it for the current field while it is ramping.
        """
        MAX_M4G_RETRIES = 100

        if self._m4g_completion_pending:
            self._m4g_completion_pending = False
            return None

        # Check if we have set the magnetic field yet
        if self.magnet_initialized is False:
            self.print_main.emit("Checking the remote connection to the magnet supply.")
            self.instrument.is_remote()
            time.sleep(1)

            # Start the field sweep
            self.instrument.field(self.end)
            self.magnet_initialized = True
        else:
            self.instrument.field(self.end)

        # Grab our data with proper error handling
        try:
            dt = safe_get(self.set_param)
            data_pair = (self.set_param, dt)
            self.setpoint = dt
        except Exception as e:
            if _retry_count >= MAX_M4G_RETRIES:
                error_msg = (
                    f"M4G field read failed after {MAX_M4G_RETRIES} retries: {e}"
                )
                self.print_main.emit(error_msg)
                self.mark_error(error_msg)
                return None
            self.print_main.emit(f"M4G read error (retry {_retry_count + 1}/{MAX_M4G_RETRIES}): {e}")
            time.sleep(self.inter_delay)
            return self.step_M4G(_retry_count=_retry_count + 1)
        # Check our stop conditions- being at the end point
        if (
            abs(self.instrument.field() - self.end) <= 0.001
        ):  # Check to see if field is within 1 mT of final value
            time.sleep(0.2)
            if (
                abs(self.instrument.field() - self.end) <= 0.001
            ):  # Wait for a little bit, then check again, if true then we are done with the sweep
                self.print_main.emit(
                    f"Done with the sweep, {self.set_param.label} = {self.set_param.get()} (T)"
                )

                if self.continuous:
                    self.flip_direction()
                    return self.step_M4G()
                elif self.bidirectional and self.direction == 0:
                    self.flip_direction()
                    return self.step_M4G()
                else:
                    self.magnet_initialized = False
                    self._m4g_completion_pending = True

        return [data_pair]

    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
            self.step /= self.back_multiplier
        else:
            self.direction = 1
            self.step *= self.back_multiplier

        self.send_updates(no_sp=True)

        temp = self.begin
        self.begin = self.end
        self.end = temp
        self.step = -1 * self.step

        # Update snap origin to new begin on every flip. This ensures the step
        # grid always includes the current sweep's starting point.
        self._snap_origin = self.begin

        self.setpoint -= self.step
        self.setpoint = self._snap_to_step(self.setpoint, self._snap_origin, self.step)

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

    def ramp_to(self, value, start_on_finish=False, persist=None, multiplier=1):
        """Ramps the set_param to a given value, at a rate dictated by the multiplier.

        Parameters
        ---------
        value:
            The desired starting value for the sweep.
        start_on_finish:
            Flag to determine whether to begin the sweep as soon as ramp is finished.
        multiplier:
            Factor to alter the step size, used to ramp quicker than the sweep speed.
        """
        # 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

        # Check if we are already at the value
        curr_value = safe_get(self.set_param)
        if abs(value - curr_value) - abs(self.step / 2) < abs(self.step) * self.err:
            # self.print_main.emit(f"Already within {self.step} of the desired ramp value. Current value: {curr_value},
            # ramp setpoint: {value}.\nSetting our setpoint directly to the ramp value.")
            self.done_ramping(value, start_on_finish, persist)
            return

        # Create a new sweep to ramp our outer parameter to zero
        self.ramp_sweep = Sweep1D(
            self.set_param,
            curr_value,
            value,
            multiplier * self.step,
            inter_delay=self.inter_delay,
            complete_func=partial(self.done_ramping, value, start_on_finish, persist),
            save_data=False,
            plot_data=self.plot_data,
        )
        self.ramp_sweep.follow_param(self._params)

        self.progressState.state = SweepState.RAMPING
        self.ramp_sweep.start(ramp_to_start=False)

        self.print_main.emit(
            f"Ramping {self.set_param.label} to {value} ({self.set_param.unit}) . . . "
        )

    def ramp_to_zero(self):
        """Ramps the set_param to 0, at the same rate as already specified."""
        self.end = 0
        if self.setpoint - self.end > 0:
            self.step = (-1) * abs(self.step)
        else:
            self.step = abs(self.step)

        self.print_main.emit(f"Ramping {self.set_param.label} to 0 . . . ")
        self.start()

    @pyqtSlot()
    def done_ramping(self, value, start_on_finish=False, pd=None):
        """Alerts the sweep that the ramp is finished.

        Parameters
        ---------
        value:
            The starting value for the sweep, also the current parameter value.
        start_on_finish:
            Sweep will be called to start immediately after ramping when set to True.
        pd:
            Sets persistent data if running Sweep2D.
        """
        # 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
        # Grab the beginning
        # value = self.ramp_sweep.begin

        # Check if we are at the value we expect, otherwise something went wrong with the ramp
        actual_value = safe_get(self.set_param)
        position_error = abs(actual_value - value) - abs(self.step / 2)
        tolerance = abs(self.step) * self.err
        if position_error > tolerance:
            error_msg = (
                f"Ramping failed (possible that the direction was changed while ramping). "
                f"Expected {self.set_param.label} final value: {value}. Actual value: "
                f"{actual_value}. Error: {position_error:.6g}, Tolerance: {tolerance:.6g}. "
                f"If tolerance is too tight, consider increasing the 'err' parameter (current: {self.err})."
            )
            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(
            f"Done ramping {self.set_param.label} to {value} ({self.set_param.unit})"
        )
        if not self.try_set(self.set_param, value):
            if self.ramp_sweep is not None:
                self.ramp_sweep.kill()
                self.ramp_sweep = None
            return
        self.setpoint = value - self.step
        self.setpoint = self._snap_to_step(self.setpoint, self._snap_origin, self.step)
        # if self.ramp_sweep is not None and self.ramp_sweep.plotter is not None:
        #    self.ramp_sweep.plotter.clear()

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

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

    def get_param_setpoint(self):
        """Obtains the current value of the setpoint."""
        return f"{self.set_param.label} = {safe_get(self.set_param)} ({self.set_param.unit})"

    def reset(self, new_params=None):
        """Resets Sweep1D, optional argument to change sweep details.

        Parameters
        ---------
        new_params:
            A list of 4 values to determine how we sweep.
            [ start value, stop value, step, frequency ]
        """
        # Set our new values if desired
        if new_params is not None:
            self.begin = new_params[0]
            self.end = new_params[1]
            self.step = new_params[2]
            self.inter_delay = 1 / new_params[3]
            # Update snap origin for the new grid
            self._snap_origin = new_params[0]

        # Reset our setpoint
        self.setpoint = self.begin - self.step
        self.setpoint = self._snap_to_step(self.setpoint, self._snap_origin, self.step)

        # Reset our plots
        self.plotter = None
        self.runner = None

    def estimate_time(self, verbose=True):
        """Estimate remaining time from the current sweep state."""
        if self.progressState.state == SweepState.DONE:
            remaining = 0
        else:
            if isinstance(self.instrument, AMI430):
                rate = safe_get(self.instrument.ramp_rate)
                if not rate:
                    return 0.0
                units = safe_get(self.instrument.ramp_rate_units)
                current = self.setpoint
                target = self.end
                distance = abs(target - current)
                if units == 0:
                    remaining = distance / rate
                else:
                    remaining = distance * 60 / rate

                if self.bidirectional and self.direction == 0:
                    distance_back = abs(self.end - self.begin)
                    if units == 0:
                        remaining += distance_back / rate
                    else:
                        remaining += distance_back * 60 / rate
                return remaining

            if isinstance(self.instrument, M4G):
                try:
                    rate = abs(float(self.instrument.range0_rate()))
                except Exception:
                    return 0.0
                if rate <= 0:
                    return 0.0
                current_value = self.setpoint
                distance = abs(self.end - current_value)
                remaining = distance / rate
                if self.bidirectional and self.direction == 0:
                    remaining += abs(self.end - self.begin) / rate
                return remaining

            current_value = self.setpoint

            step_size = abs(self.step)
            if step_size == 0:
                return 0.0

            distance = abs(self.end - current_value)
            remaining = (distance / step_size) * self.inter_delay

            if self.bidirectional and self.direction == 0:
                effective_back_multiplier = (
                    self.back_multiplier if self.back_multiplier not in (None, 0) else 1.0
                )
                if effective_back_multiplier > 0:
                    distance_back = abs(self.end - self.begin)
                    remaining += (
                        distance_back / (step_size * 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["set_param"] = {
            "param": self.set_param.name,
            "instr_module": self.set_param.instrument.__class__.__module__,
            "instr_class": self.set_param.instrument.__class__.__name__,
            "instr_name": self.set_param.instrument.name,
            "start": self.begin,
            "stop": self.end,
            "step": self.step,
        }
        json_dict["attributes"]["bidirectional"] = self.bidirectional
        json_dict["attributes"]["continual"] = self.continuous
        json_dict["attributes"]["x_axis_time"] = self.x_axis
        return json_dict

    @classmethod
    def from_json(cls, json_dict, station):
        attrs = dict(json_dict.get("attributes", {}))
        bidirectional = attrs.pop("bidirectional", False)
        continual = attrs.pop("continual", False)
        x_axis_time = attrs.pop("x_axis_time", 0)

        sp = json_dict["set_param"]
        p = BaseSweep._load_parameter_by_type(
            sp["param"],
            sp["instr_name"],
            sp["instr_module"],
            sp["instr_class"],
            station,
        )
        return cls(
            p,
            sp["start"],
            sp["stop"],
            sp["step"],
            bidirectional=bidirectional,
            continual=continual,
            x_axis_time=x_axis_time,
            **attrs,
        )

    def _params_to_exclude_from_follow(self) -> set:
        """Exclude the set_param from follow_params export for 1D sweeps."""
        try:
            return {f"{self.set_param.instrument.name}.{self.set_param.name}"}
        except Exception:
            return set()

    def __del__(self):
        """Destructor. Should delete all child threads and close all figures when the sweep object is deleted."""
        self.kill()
