# This file contains all classes and functions directly related to the operation of the helmholtz test stand.
# The two main classes are Axis and ArduinoCtrl, see their definitions for details.

# import packages:
import numpy as np
import serial
import traceback
from tkinter import messagebox

# import other project files
from User_Interface import ui_print
from pyps2000b import PS2000B
from Arduino import Arduino
import config_handling as config
import globals as g


class Axis:
    # Main class representing an axis (x,y,z) of the test stand
    # contains static and dynamic status information about this axis and methods to control it

    def __init__(self, index, device, PSU_channel, arduino_pin):
        # static information
        self.index = index  # index of this axis, 0->X, 1->Y, 2->Z
        self.device = device  # power supply object for this axis (PS2000B class)
        self.channel = PSU_channel  # power supply unit channel (0 or 1)
        self.ardPin = arduino_pin  # output pin on the arduino for switching polarity on this axis

        self.name = g.AXIS_NAMES[index]  # get name of this axis from list in globals.py (e.g. "X-Axis"
        self.port = g.PORTS[index]  # get serial port of this axis PSU

        # read static information from the configuration object (which has read it from the config file or settings):
        self.resistance = float(config.read_from_config(self.name, "resistance", config.CONFIG_OBJECT))
        self.max_amps = float(config.read_from_config(self.name, "max_amps", config.CONFIG_OBJECT))
        self.max_volts = float(config.read_from_config(self.name, "max_volts", config.CONFIG_OBJECT))

        self.coil_constant = float(config.read_from_config(self.name, "coil_const", config.CONFIG_OBJECT))
        self.ambient_field = float(config.read_from_config(self.name, "ambient_field", config.CONFIG_OBJECT))

        max_field = self.max_amps * self.coil_constant  # calculate max field reachable in this axis
        self.max_field = np.array([-max_field, max_field])  # make array with min/max reachable field (w/o compensation)
        # calculate max and min field that can be reached after compensating for the ambient field
        self.max_comp_field = np.array([self.ambient_field - max_field, self.ambient_field + max_field])  # [min, max]

        # initialize dynamic information, this is updated by self.update_status_info() later
        self.connected = "Not Connected"
        self.output_active = "Unknown"  # power output on the PSU enabled?
        self.remote_ctrl_active = "Unknown"  # remote control on the PSU enabled?
        self.voltage_setpoint = 0  # target voltage on PSU [V]
        self.voltage = 0  # actual voltage on PSU [V]
        self.current_setpoint = 0  # target current on PSU [A]
        self.current = 0  # actual current on PSU [A]

        self.polarity_switched = "Unknown"  # polarity switched on the Arduino?

        self.target_field_comp = 0  # field to be created by coil pair (this is sent to the coils) [T]
        self.target_field = 0  # field that should occur in measurement area (ambient still needs to be compensated) [T]
        self.target_current = 0  # signed current that should pass through coil pair [A]

        if self.device is not None:
            self.update_status_info()

    def update_status_info(self):  # Read out the values of the dynamic parameters stored in this object and update them
        try:  # try to read out the data, this will fail on connection error to PSU
            self.device.update_device_information(self.channel)  # update the information in the device object
            device_status = self.device.get_device_status_information(self.channel)  # get object with new status info

            if device_status.output_active:  # is the power output active?
                self.output_active = "Active"
            else:
                self.output_active = "Inactive"

            # is remote control active, allowing the device to be controlled by this program?
            if device_status.remote_control_active:
                self.remote_ctrl_active = "Active"
            else:
                self.remote_ctrl_active = "Inactive"

            # get currents and voltages:
            self.voltage = self.device.get_voltage(self.channel)
            self.voltage_setpoint = self.device.get_voltage_setpoint(self.channel)
            self.current = self.device.get_current(self.channel)
            self.current_setpoint = self.device.get_current_setpoint(self.channel)

        except (serial.serialutil.SerialException, IndexError):  # Connection error, usually the PSU is unplugged
            if self.connected == "Connected":  # only show error messages if the device was connected before this error
                # Show error as print-out in console and as pop-up:
                ui_print("Connection Error with %s PSU on %s" % (self.name, self.port))
                messagebox.showerror("PSU Error", "Connection Error with %s PSU on %s" % (self.name, self.port))
            # set status attributes to connection error status:
            self.connected = "Connection Error"
            self.output_active = "Unknown"
            self.remote_ctrl_active = "Unknown"
        else:  # no communications error
            self.connected = "Connected"  # PSU is connected

    def print_status(self):  # print out the current status of the device (not used at the moment)
        ui_print("%s, %0.2f V, %0.2f A"
                 % (self.device.get_device_status_information(self.channel),
                    self.device.get_voltage(self.channel), self.device.get_current(self.channel)))

    def power_down(self):  # temporary powerdown, set outputs to 0 but keep connections enabled
        try:
            # set class object attributes to 0 to reflect shutdown in status displays, log files etc.
            self.target_current = 0
            self.target_field = 0
            self.target_field_comp = 0

            if self.device is not None:  # there is a PSU connected for this axis
                self.device.set_voltage(0, self.channel)  # set voltage on PSU channel to 0
                self.device.set_current(0, self.channel)  # set current on PSU channel to 0
                self.device.disable_output(self.channel)  # disable power output on PSU channel
            g.ARDUINO.digitalWrite(self.ardPin, "LOW")  # set arduino pin for polarity switch relay to unpowered state
        except Exception as e:  # some error was encountered
            # show error message:
            ui_print("Error while powering down %s: %s" % (self.name, e))
            messagebox.showerror("PSU Error!", "Error while powering down %s: \n%s" % (self.name, e))

    def set_signed_current(self, value):
        # sets current with correct polarity on this axis, this is the primary way to control the test stand

        # ui_print("Attempting to set current", value, "A")
        self.target_current = value  # show target value in object attribute for status display, logging etc.

        if abs(value) > self.max_amps:  # prevent excessive currents
            self.power_down()  # set output to 0 and deactivate
            raise ValueError("Invalid current value on %s. Tried %0.2fA, max. %0.2fA allowed"
                             % (self.name, value, self.max_amps))

        elif value >= 0:  # switch the e-box relay to change polarity as needed
            g.ARDUINO.digitalWrite(self.ardPin, "LOW")  # command the output pin on the arduino in the electronics box
        elif value < 0:
            g.ARDUINO.digitalWrite(self.ardPin, "HIGH")  # command the output pin on the arduino in the electronics box

        # determine voltage limit to be set on PSU, must be high enough to not limit the current:
        # min. 8V, max. max_volts, in-between as needed with current value (+margin to not limit current)
        maxVoltage = min(max(1.3 * value * self.resistance, 8), self.max_volts)  # limit voltage
        if self.connected == "Connected":  # only try to command the PSU if its actually connected
            self.device.set_current(abs(value), self.channel)  # set desired current
            self.device.set_voltage(maxVoltage, self.channel)  # set voltage limit
            self.device.enable_output(self.channel)  # activate the power output
        else:  # the PSU is not connected
            ui_print(self.name, "not connected, can't set current.")

    def set_field_simple(self, value):  # forms magnetic field as specified by value, w/o cancelling ambient field
        self.target_field = value  # update object attribute for display
        self.target_field_comp = value  # same as above, bc no compensation
        current = value / self.coil_constant  # calculate needed current
        self.set_signed_current(current)  # command the test stand

    def set_field(self, value):  # forms magnetic field as specified by value, corrected for ambient field
        self.target_field = value  # update object attribute for display
        field = value - self.ambient_field  # calculate needed field after compensation
        self.target_field_comp = field  # update object attribute for display
        current = field / self.coil_constant  # calculate needed current
        self.set_signed_current(current)  # command the test stand


class ArduinoCtrl(Arduino):
    # main class to control the electronics box (which means commanding the arduino inside)
    # inherits from the Arduino library

    def __init__(self):
        self.connected = "Unknown"  # connection status attribute, nominal "Connected"
        self.pins = [0, 0, 0]  # initialize list with pins to switch relay of each axis
        for i in range(3):  # get correct pins from the config
            self.pins[i] = int(config.read_from_config(g.AXIS_NAMES[i], "relay_pin", config.CONFIG_OBJECT))

        ui_print("\nConnecting to Arduino...")
        try:  # try to set up the arduino
            Arduino.__init__(self)  # search for connected arduino and connect by initializing arduino library class
            for pin in self.pins:
                self.pinMode(pin, "Output")
                self.digitalWrite(pin, "LOW")
        except Exception as e:  # some error occurred, usually the arduino is not connected
            ui_print("Connection to Arduino failed.", e)
            self.connected = "Not Connected"
        else:  # connection was successfully established
            self.connected = "Connected"
            ui_print("Arduino ready.")

    def update_status_info(self):  # update the attributes stored in this class object
        if self.connected == "Connected":  # only do this if arduino is connected (initialize new instance to reconnect)
            try:  # try to read the status of the pins from the arduino
                for axis in g.AXES:  # go through all three axes
                    if g.ARDUINO.digitalRead(axis.ardPin):  # pin is HIGH --> relay is switched
                        axis.polarity_switched = "True"  # set attribute in axis object accordingly
                    else:  # pin is LOW --> relay is not switched
                        axis.polarity_switched = "False"  # set attribute in axis object accordingly
            except Exception as e:  # some error occurred while trying to read status, usually arduino is disconnected
                # show warning messages to alert user
                ui_print("Error with Arduino:", e)
                messagebox.showerror("Error with Arduino!", "Connection Error with Arduino: \n%s" % e)
                for axis in g.AXES:  # set polarity switch attributes in axis objects to "Unknown"
                    axis.polarity_switched = "Unknown"
                self.connected = "Connection Error"  # update own connection status
            else:  # no error occurred --> data was read successfully
                self.connected = "Connected"  # update own connection status

    def safe(self):  # sets relay switching pins to low to depower most of the electronics box
        for pin in self.pins:
            self.digitalWrite(pin, "LOW")


def value_in_limits(axis, key, value):  # Check if value is within safe limits (set in globals.py)
    # axis is string with axis name, e.g. "X-Axis"
    # key specifies which value to check, e.g. current
    max_value = g.default_arrays[key][1][g.AXIS_NAMES.index(axis)]  # get max value from dictionary in globals.py
    min_value = g.default_arrays[key][2][g.AXIS_NAMES.index(axis)]  # get min value from dictionary in globals.py

    if float(value) > float(max_value):  # value is too high
        return 'HIGH'
    elif float(value) < float(min_value):  # value is too low
        return 'LOW'
    else:  # value is within limits
        return 'OK'


def setup_all():  # main test stand initialization function
    # creates device objects for all PSUs and Arduino and sets their values
    # initializes an object of class Axis for all three axes (x,y,z)

    # Setup Arduino:
    try:  # broad error handling for unforeseen errors, handling in ArduinoCtrl should catch most errors
        if g.ARDUINO is not None:  # the arduino has been initialized before, so we need to first close its connection
            try:
                g.ARDUINO.close()  # close serial link
            except serial.serialutil.SerialException:
                pass
                # serial.flush() in Arduino.close() fails when reconnecting
                # this ignores it and allows serial.close() to execute (I think)
            except AttributeError:
                pass
                # when no Arduino is connected but g.ARDUINO has been initialized then there is nothing to close
                # this throws an exception, which can be ignored

        g.ARDUINO = ArduinoCtrl()  # initialize the arduino object from the control class, connects and sets up
    except Exception as e:  # some unforeseen error occurred
        # show error messages to alert user
        ui_print("Arduino setup failed:", e)
        ui_print(traceback.print_exc())
        messagebox.showerror("Error!", "Arduino setup failed:\n%s \nCheck traceback in console." % e)

    # Setup PSUs and axis objects:
    g.AXES = []  # initialize global list containing the three axis objects

    # get serial ports for the PSUs from config
    g.XY_PORT = config.read_from_config("PORTS", "xy_port", config.CONFIG_OBJECT)
    g.Z_PORT = config.read_from_config("PORTS", "z_port", config.CONFIG_OBJECT)
    g.PORTS = [g.XY_PORT, g.XY_PORT, g.Z_PORT]  # write list with PSU port for each axis (X/Y share PSU)

    # setup PSU and axis objects for X and Y axes:
    ui_print("\nConnecting to XY Device on %s..." % g.XY_PORT)
    try:  # try to connect to the PSU
        if g.XY_DEVICE is not None:  # if PSU has previously been connected we need to close the serial link first
            ui_print("Closing serial connection on XY device")
            g.XY_DEVICE.serial.close()
            g.XY_DEVICE = None
        g.XY_DEVICE = PS2000B.PS2000B(g.XY_PORT)  # setup PSU
        ui_print("Connection established.")
        g.X_AXIS = Axis(0, g.XY_DEVICE, 0, g.ARDUINO.pins[0])  # create axis objects (index, PSU, channel, relay pin)
        g.Y_AXIS = Axis(1, g.XY_DEVICE, 1, g.ARDUINO.pins[1])
    except serial.serialutil.SerialException:  # communications error, usually PSU is not connected or wrong port set
        g.X_AXIS = Axis(0, None, 0, g.ARDUINO.pins[0])  # create axis objects without the PSU
        g.Y_AXIS = Axis(1, None, 1, g.ARDUINO.pins[1])
        ui_print("XY Device not connected or incorrect port set.")

    # same for the Z axis
    ui_print("Connecting to Z Device on %s..." % g.Z_PORT)
    try:
        if g.Z_DEVICE is not None:
            ui_print("Closing serial connection on Z device")
            g.Z_DEVICE.serial.close()
            g.Z_DEVICE = None
        g.Z_DEVICE = PS2000B.PS2000B(g.Z_PORT)
        ui_print("Connection established.")
        g.Z_AXIS = Axis(2, g.Z_DEVICE, 0, g.ARDUINO.pins[2])
    except serial.serialutil.SerialException:
        g.Z_AXIS = Axis(2, None, 0, g.ARDUINO.pins[2])
        ui_print("Z Device not connected or incorrect port set.")

    # put newly created axis objects into a list for access later
    g.AXES.append(g.X_AXIS)
    g.AXES.append(g.Y_AXIS)
    g.AXES.append(g.Z_AXIS)

    ui_print("")  # print new line


def set_to_zero(device):  # sets voltages and currents to 0 on all channels of a specific PSU
    device.voltage1 = 0
    device.current1 = 0
    device.voltage2 = 0
    device.current2 = 0


def power_down_all():  # on all PSUs set all outputs to 0 but keep connections enabled
    for axis in g.AXES:
        axis.power_down()  # set outputs to 0 and pin to low on this axis


def shut_down_all():  # safe shutdown at program end or on error
    # set outputs to 0 and disable connections on all devices

    ui_print("\nAttempting to safely shut down all devices. Check equipment to confirm.")
    # start writing string to later show how shutdown on all devices went in a single info pop-up:
    message = "Tried to shut down all devices. Check equipment to confirm."

    # Shut down PSUs:
    if g.XY_DEVICE is not None:  # the PSU has been setup before
        try:  # try to safe the PSU
            set_to_zero(g.XY_DEVICE)  # set currents and voltages to 0 for both channels
            g.XY_DEVICE.disable_all()  # disable power output on both channels
        except BaseException as e:  # some error occurred, usually device has been disconnected
            ui_print("Error while deactivating XY PSU:", e)  # print the problem in the console
            message += "\nError while deactivating XY PSU: %s" % e  # append status to the message to show later
        else:  # device was successfully deactivated
            ui_print("XY PSU deactivated.")
            message += "\nXY PSU deactivated."  # append message to show later
    else:  # the device was not connected before
        # tell user there was no need/no possibility to deactivate:
        ui_print("XY PSU not connected, can't deactivate.")
        message += "\nXY PSU not connected, can't deactivate."

    # same as above
    if g.Z_DEVICE is not None:
        try:
            set_to_zero(g.Z_DEVICE)
            g.Z_DEVICE.disable_all()
        except BaseException as e:
            ui_print("Error while deactivating Z PSU:", e)
            message += "\nError while deactivating Z PSU: %s" % e
        else:
            ui_print("Z PSU deactivated.")
            message += "\nZ PSU deactivated."
    else:
        ui_print("Z PSU not connected, can't deactivate.")
        message += "\nZ PSU not connected, can't deactivate."

    # Shut down Arduino:
    try:
        g.ARDUINO.safe()  # call safe method in ArduinoCtrl class (all relay pins to LOW)
    except BaseException as e:  # some error occurred
        ui_print("Arduino safing unsuccessful:", e)
        message += "\nArduino safing unsuccessful: %s" % e  # append to the message to show later
    # this throws no exception, even when arduino is not connected
    # ToDo (optional): figure out error handling for this
    try:
        g.ARDUINO.close()  # close the serial link
    except BaseException as e:  # something went wrong there
        if g.ARDUINO.connected == "Connected":  # Arduino was connected, some error occurred
            ui_print("Closing Arduino connection failed:", e)
            message += "\nClosing Arduino connection failed: %s" % e
        else:  # Arduino was not connected, so error is expected
            ui_print("Arduino not connected, can't close connection.")
            message += "\nArduino not connected, can't close connection."
    else:  # no problems, connection was successfully closed
        ui_print("Serial connection to Arduino closed.")
        message += "\nSerial connection to Arduino closed."

    messagebox.showinfo("Program ended", message)  # Show a unified pop-up with how the shutdown on each device went


def set_field_simple(vector):  # forms magnetic field as specified by vector, w/o cancelling ambient field
    for i in [0, 1, 2]:
        try:
            g.AXES[i].set_field_simple(vector[i])  # try to set the field on each axis
        except ValueError as e:  # a limit was violated, usually the needed current was too high
            ui_print(e)  # let the user know


def set_field(vector):  # forms magnetic field as specified by vector, corrected for ambient field
    # same as set_field_simple(), but with compensation
    for i in [0, 1, 2]:
        try:
            g.AXES[i].set_field(vector[i])
        except ValueError as e:
            ui_print(e)


def set_current_vec(vector):  # sets currents on each axis according to given vector
    i = 0
    for axis in g.AXES:
        try:
            axis.target_field = 0  # set target field attribute to 0 to show that current, not field is controlled atm
            axis.target_field_comp = 0  # as above

            axis.set_signed_current(vector[i])  # command test stand to set the current
        except ValueError as e:  # current was too high
            ui_print(e)  # print out the error message
        i += 1


def devices_ok(xy_off=False, z_off=False, arduino_off=False):
    # check if all devices are connected, return True if yes
    # checks for individual devices can be disabled by parameters above (default not disabled)
    try:  # handle errors while checking connections
        if not xy_off:  # if check for this device is not disabled
            if g.XY_DEVICE is not None:  # has the handle for this device been set?
                g.X_AXIS.update_status_info()  # update info --> this actually communicates with the device
                if g.X_AXIS.connected != "Connected":  # if not connected
                    return False  # return and exit function
            else:  # if handle has not been set the device is inactive --> not ok
                return False
        if not z_off:  # same as above
            if g.Z_DEVICE is not None:
                g.Z_AXIS.update_status_info()
                if g.Z_AXIS.connected != "Connected":
                    return False
            else:
                return False

        if not arduino_off:  # check not disabled
            g.ARDUINO.update_status_info()  # update status info --> attempts communication
            if g.ARDUINO.connected != "Connected":
                return False
    except Exception as e:  # if an error is encountered while checking the devices
        messagebox.showerror("Error!", "Error while checking devices: \n%s" % e)  # show error pop-up
        return False  # clearly something is not ok
    else:  # if nothing has triggered so far all devices are ok --> return True
        return True