from pyps2000b import PS2000B
from Arduino import Arduino
import settings as g
import pandas
import time
import numpy as np
import serial


class Axis:
    def __init__(self, index, device, PSU_channel, arduino_pin):
        self.index = index
        self.device = device  # power supply object (PS2000B class)
        self.channel = PSU_channel  # power supply unit channel (1 or 2)
        self.ardPin = arduino_pin  # output pin on the arduino for switching polarity on this axis

        self.resistance = 0  # [Ohm]
        # maximum allowable values to pass through this circuit
        self.max_watts = 0  # [W]
        self.max_amps = 0  # [A]
        self.max_volts = 0  # [V]

        self.coil_constant = 0  # coil constant of this axis [T/A]
        self.ambient_field = 0  # ambient field in this axis [T]
        # ToDo: get this info from settings file

        self.output_active = 0  # power output on the PSU enabled?
        self.remote_ctrl_active = 0  # 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 = 0  # polarity switched on the Arduino?

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

        self.name = g.AXIS_NAMES[index]

    def update_status_info(self):  # Read out the values of the parameters stored in this class and update them
        self.device.update_device_information(self.channel)
        device_status = self.device.get_device_status_information(self.channel)
        self.output_active = device_status.output_active
        self.remote_ctrl_active = device_status.remote_control_active

        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)

        self.polarity_switched = g.ARDUINO.digitalRead(self.ardPin)  # ToDo: Test if this actually works

    def print_status(self):  # axis = axis control variable, stored in settings.py
        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
        self.device.voltage1 = 0
        self.device.current1 = 0
        g.ARDUINO.digitalWrite(self.ardPin, "LOW")

    def set_signed_current(self, value):  # sets current with correct polarity on this axis
        device = self.device
        channel = self.channel
        ardPin = self.ardPin
        if abs(value) > self.max_amps:  # prevent excessive currents
            shut_down_all()  # safe all devices
            raise ValueError("Invalid current value. Tried %0.2fA, max. %0.2fA allowed" % (value, self.max_amps))
        elif value >= 0:  # switch polarity as needed
            g.ARDUINO.digitalWrite(ardPin, "LOW")
        elif value < 0:
            g.ARDUINO.digitalWrite(ardPin, "HIGH")
        else:
            raise Exception("This should be impossible.")
        device.set_current(abs(value), channel)
        maxVoltage = min(max(1.1 * self.max_amps * self.resistance, 8), self.max_volts)  # limit voltage
        device.set_voltage(maxVoltage)

    def set_field_simple(self, value):  # forms magnetic field as specified by value, w/o cancelling ambient field
        current = value / self.coil_constant
        self.set_signed_current(current)

    def set_field(self, value):  # forms magnetic field as specified by value, corrected for ambient field
        field = value - self.ambient_field
        current = field / self.coil_constant
        self.set_signed_current(current)


def setup_axes():  # creates device objects for all PSUs and sets their values

    print("Connecting to XY Device on %s..." % g.XY_PORT)
    try:
        g.XY_DEVICE = PS2000B.PS2000B(g.XY_PORT)  # setup PSU
        print("Connection established.")
        g.X_AXIS = Axis(0, g.XY_DEVICE, 0, g.RELAY_PINS[0])  # create axis objects
        g.Y_AXIS = Axis(1, g.XY_DEVICE, 1, g.RELAY_PINS[1])
    except serial.serialutil.SerialException:
        g.X_AXIS = Axis(0, None, 0, g.RELAY_PINS[0])  # create axis objects
        g.Y_AXIS = Axis(1, None, 1, g.RELAY_PINS[1])
        print("XY Device not connected or incorrect port set.")

    if g.Z_PORT == "NC":  # check if device is connected
        g.Z_AXIS = Axis(2, None, 0, g.RELAY_PINS[2])
        print("Z Device not connected or port not set.")

    print("Connecting to Z Device on %s..." % g.XY_PORT)
    try:
        g.Z_DEVICE = PS2000B.PS2000B(g.Z_PORT)
        print("Connection established.")
        g.Z_AXIS = Axis(2, g.Z_DEVICE, 0, g.RELAY_PINS[2])
    except serial.serialutil.SerialException:
        g.Z_AXIS = Axis(2, None, 0, g.RELAY_PINS[2])
        print("Z Device not connected or incorrect port set.")

    g.AXES.append(g.X_AXIS)
    g.AXES.append(g.Y_AXIS)
    g.AXES.append(g.Z_AXIS)

    i = 0
    for axis in g.AXES:
        axis.resistance = g.RESISTANCES[i]
        axis.max_watts = g.MAX_WATTS[i]
        axis.max_amps = np.sqrt(axis.max_watts / axis.resistance)
        axis.max_volts = g.MAX_VOLTS[i]

        axis.coil_constant = g.COIL_CONST[i]
        axis.ambient_field = g.AMBIENT_FIELD[i]
        i = i+1


def activate_all():  # enables remote control and output on all PSUs and channels
    g.XY_DEVICE.enable_all()
    g.Z_DEVICE.enable_all()


def deactivate_all():  # disables remote control and output on all PSUs and channels
    g.XY_DEVICE.disable_all()
    g.Z_DEVICE.disable_all()


def setup_arduino():
    try:
        g.ARDUINO = Arduino()  # search for connected arduino and set handle
    except Exception:
        print("There seems to be no Arduino connected.")
    else:
        for pin in g.RELAY_PINS:
            g.ARDUINO.pinMode(pin, "Output")
            g.ARDUINO.digitalWrite(pin, "LOW")
        print("Arduino ready.")


def safe_arduino():  # sets output pins to low and closes serial connection
    for pin in g.RELAY_PINS:
        g.ARDUINO.digitalWrite(pin, "LOW")


def print_status_3():
    print("X-Axis:")
    g.X_AXIS.print_status()
    print("Y-Axis:")
    g.Y_AXIS.print_status()
    print("Z-Axis:")
    g.Z_AXIS.print_status()


def set_to_zero(device):  # sets voltages and currents to 0
    device.voltage1 = 0
    device.current1 = 0
    device.voltage2 = 0
    device.current2 = 0


def power_down_all():  # temporary, set all outputs to 0 but keep connections enabled
    set_to_zero(g.XY_DEVICE)
    set_to_zero(g.Z_DEVICE)
    safe_arduino()


def shut_down_all():  # shutdown at program end or on error, set outputs to 0 and disable connections
    set_to_zero(g.XY_DEVICE)
    set_to_zero(g.Z_DEVICE)
    deactivate_all()
    safe_arduino()
    g.ARDUINO.close()


def set_field_simple(vector):  # forms magnetic field as specified by vector, w/o cancelling ambient field
    for i in [0, 1, 2]:
        g.AXES[i].set_field_simple(vector[i])


def set_field(vector):  # forms magnetic field as specified by vector, corrected for ambient field
    for i in [0, 1, 2]:
        g.AXES[i].set_field(vector[i])


def set_current_vec(i_vec):  # sets needed currents on each axis for given vector
    for i in [0, 1, 2]:
        g.AXES[i].set_signed_current(i_vec[i])


def execute_csv(filepath, printing=0):  # runs through csv file containing times and desired field vectors
    # csv format: time (s); xField (T); yField (T); zField (T)
    # decimal commas
    print("Reading File:", filepath)
    file = pandas.read_csv(filepath, sep=';', decimal=',', header=0)  # read csv file
    array = file.to_numpy()  # convert csv to array
    t_zero = time.time()
    t_ref = t_zero
    i = 0
    print("Starting Execution...")
    while i < len(array):
        t = time.time() - t_zero
        if t >= array[i, 0]:
            field_vec = array[i, 1:4]
            print("t = %0.2f s, target field vector = " % (array[i, 0]), field_vec)
            set_field(field_vec)
            i = i + 1
        if t - t_ref >= 1 and printing == 1:  # print status every second
            print_status_3()
            t_ref = t
    print("File executed, powering down channels.")
    power_down_all()  # set currents and voltages to 0, set arduino pins to low