"""
Direct laser writing module. Import this into your script and then use the 
movement and shutter control functions to perform DLW.
"""
import nidaqmx
import pipython
import sys
import time

from pipython import datarectools, pitools

class Writer:
    """
    Main class for direct laser writing. Sets up communication with the stage
    and shutter and allows control of each.
    """

    def __init__(self, stage_model='E-727', shutter_channel="cDAQ1Mod3/port0/line2")

        # Assign instance variables
        self.stage_model = stage_model
        self.shutter_channel = shutter_channel

        # Initialise stage and shutter
        self.stage = Stage(self.stage_model)
        self.shutter = Shutter(self.shutter_channel)


class Stage:
    """
    Represents a PI stage using the pipython package
    """

    def __init__(self, model):
        """
        Connect to the PI stage over USB. Tested on a PI E-727 piezo stage 
        controller connected over USB. The pipython package has minimal 
        documentation so it is useful to look inside the code. Unfortunately, 
        PI only provides a wheel, no source code, but once installed you can 
        import pipython and run `print pipython.__path__` to get the path 
        where it is installed, and then you can find all the code there.
        """
        # Set up stage
        self._pidevice = pipython.GCSDevice(self.stage_model)

        # Attempt to connect to stage 
        # If this fails, don't worry, we'll try again when some functionality 
        # is called for
        try:
            self._setup()
            self.initialised = True
        except:
            self.initialised = False


    def _setup(self):
        """
        Attempt to connect to and initialise the stage
        """
        # Connect to stage over USB
        usb_devices = self._pidevice.EnumerateUSB()
        if len(usb_devices) < 1:
            raise Exception("No PI stages found")

        # Connect to the first USB device found
        pidevice.ConnectUSB(usb_devices[0].split(' ')[-1])
        pitools.startup(self._pidevice, servostates=True)

        # Enable velocity control for all axes
        self._pidevice.VCO(self._pidevice.axes)

        # Wait for the stage to stabilise
        pitools.waitontarget(self._pidevice, self._pidevice.axes)

    def set_pos(self, coords, wait=True):
        """
        Move to coordinates specified in coords. Coords should be a dictionary 
        with axes indices as keys and positions as values. If wait is specified, 
        wait until the target is reached.

        Example call:
            stage.moveto({1: 35, 2: 9, 3: 20.9})
        will move the stage to (x, y, z) position (35, 9, 20.9).
        """
        if not self.initialised:
            self._setup()

        # Move to coordinates
        self._pidevice.MOV(coords)

        if wait:
            # Wait for stage to reach coordinates
            pitools.waitontarget(self._pidevice, list(coords.keys()))

    def get_pos(self):
        """
        Returns the current position of all axes
        """
        if not self.initialised:
            self._setup()
        return self._pidevice.qPOS(self._pidevice.axes)

    def set_vel(self, vel):
        """
        Set the velocity for some or all axes. When moveto is called, the stage 
        will move at the specified velocities. vel is a dictionary or float.
        """
        if not self.initialised:
            self._setup()
        self._pidevice.VEL(vel)

    def get_vel(self):
        """
        Returns the currently set velocities for all axes
        """
        if not self.initialised:
            self._setup()
        return self._pidevice.qVEL(self._pidevice.axes)

    def set_acc(self, acc):
        """
        Set the acceleration for some or all axes. When moveto is called, the 
        stage will accelerate at the specified rates. acc is a dictionary or
        float.
        """
        if not self.initialised:
            self._setup()
        self._pidevice.ACC(acc)

    def get_acc(self):
        """
        Returns the currently set accelerations for all axes
        """
        if not self.initialised:
            self._setup()
        return self._pidevice.qACC(self._pidevice.axes)


class Shutter:
    """
    Represents a shutter using an NI DAQ with the nidaqmx package
    """

    def __init__(self, channel):
        """
        Set up shutter instance
        """
        self.channel = channel

        try:
            self._setup()
            self.initialised = True
        except:
            self.initialised = False
            

    def _setup(self):
        """
        Connect to DAQ
        """
        self._task = nidaqmx.Task()
        self._task.do_channels.add_do_chan(self.channel,
                line_grouping=nidaqmx.constants.LineGrouping.CHAN_PER_LINE)
        self._task.start()

    def open(self):
        """
        Open the shutter
        """
        if not self.initialised:
            self._setup()
        self._task.write(True)

    def close(self):
        """
        Close the shutter
        """
        if not self.initialised:
            self._setup()
        self._task.write(False)

    def get_status(self):
        """
        Get the current status of the shutter. True/1 = open, False/0 = closed
        """
        if not self.initialised:
            self._setup()
        return bool(self._task.read())