# -*- coding: utf-8 -*-
"""
Code rotates two motors (Thorlabs TDC001 via USB) and reads in the data from the
LeCroy Oscilloscope (Via Ethernet)
See the Main Code section for more detail.
"""

# Import APTMotor class from PyAPT
from PyAPT import APTMotor
import time
import csv
import win32com.client  # import win32 library


"""
Documentation Links:
https://github.com/mcleung/PyAPT
https://www.thorlabs.com/Software/Motion%20Control/APT_Communications_Protocol_Rev_19.pdf
http://teledynelecroy.com/doc/using-python-with-activedso-for-remote-communication
http://cdn.teledynelecroy.com/files/manuals/wr2_rcm_revb.pdf
http://cdn.teledynelecroy.com/files/manuals/activedso-developers-guide.pdf
"""

"""
Define Functions
"""

def OffsetBattering (chan, maxd, mind, offset):
    """         
    Due to the DC offset, the trace might not be in range
    p.s. I meant to call this bartering, but missed.
    """
    
    # 1. Read in data.
    scope.WriteString("ARM",1) # Arm the scope
    scope.WriteString("FRTR", 1) # Trigger
    askstr = "C%i:INSP? 'SIMPLE'" %(chan) # Format ask string
    scope.WriteString(askstr, 1) # Ask for data
    getstr = "C%i" %(chan) # Format get string
    scopedata = scope.GetScaledWaveform(getstr, 24727*4, 0) # Obtain data
    
    
    # 2. Set the counter
    cnt = 0 # counter to avoid infinite loops
    
    # 3. Start the While loop
    while max(scopedata) > maxd or min(scopedata) < mind: # While data not in range:
        cnt = cnt+1
        if max(scopedata) > maxd:
            offset = offset - 0.010
            maxd = maxd + 0.010
            mind = mind + 0.010 # Move the whole range in case it changes direction later
        elif min(scopedata) < mind:
            offset = offset + 0.010
            mind = mind - 0.010
            maxd = maxd - 0.010
        else:
            print "In loop but not > max or < min."    
            
        # 4. Change offset
        ofstr = "C1:OFST %fV" %(offset) # Format offset
        scope.WriteString(ofstr,1) #Tell scope to offset
        # 5. Take data again
        scope.WriteString("ARM",1) # Arm
        scope.WriteString("FRTR", 1) # Trigger
        askstr = "C%i:INSP? 'SIMPLE'" %(chan) # Format ask string
        scope.WriteString(askstr, 1) # Ask for data
        getstr = "C%i" %(chan) # Format Get String
        scopedata = scope.GetScaledWaveform(getstr, 24727*4, 0) # Obtain data
    
        if cnt == 50: # To prevent infinite loop
            print "ARGH INFINITE OFFSET LOOPS!!"
            break
    return scopedata, maxd, mind, offset
    
"""
Main code
"""

if __name__ == "__main__":
    """
    Code rotates two motors (Thorlabs TDC001 via USB) and reads in the data from the
    LeCroy Oscilloscope (Via Ethernet). It outputs a data file called SxSy.csv 
    where x is the position (in degrees) of spinner one (Eve) and y is spinner 2
    (Alice/Bob). 
    The file contains data for each detector ABE, then this repeated so there are 
    three readings for each detector at each spinner position. 
    The spinner does NOT move and come back, but stays stationary.
    The order of data in the file is A1,B1,E1,A2,B2,E2,A3,B3,E3. So 9 rows in all.
    """
    try: # This try/catch means that it will automatically disconnect from the instruments if it crashes.
        
        """
        initialise scope
        """
        scope = win32com.client.Dispatch("LeCroy.ActiveDSOCtrl.1") #scope talking stuff
        scope.MakeConnection("IP:169.254.161.99") #scope IP address
        scope.WriteString("C1:VDIV .002", 1) # VDIV = volts per division
        scope.WriteString("C2:VDIV .002", 1)
        scope.WriteString("C3:VDIV .002", 1)
        scope.WriteString("C1:CPL D50", 1) #Set to DC offset 50 Ohm
        scope.WriteString("C2:CPL D50", 1)
        scope.WriteString("C3:CPL D50", 1)
        scope.WriteString("TDIV 500ns",1) # TDIV = Time Division
        #scope.WriteString("") #Memory size? Sample points? What is this called?
        #I want to set that here, but don't know how
    
        """
        Set Important Variables
        """
        MA = 0 # Where to start M1
        MB = 30 # Where to start M2 - ALSO  CHANGE AT HOME M2
        RA = 50 # How far M1 travels
        RB = 50 # How far M2 travels
        
        # And less important ones.
        maxd_c1 = 0.008 # Max 
        mind_c1 = -0.008 
        offset_c1 = 0 # We don't want to be wasting time by resetting them each time
        maxd_c2 = 0.008 # So we initialise them up here.
        mind_c2 = -0.008
        offset_c2 = 0
        maxd_c3 = 0.008
        mind_c3 = -0.008
        offset_c3 = 0
        
        """
        initialise + home Motor 1
        """
        print  "Motor 1:"
        # Create object corresponding to the motor.
        Motor1 = APTMotor(83857513, HWTYPE=31)
        # The number should correspond to the serial number.
        # Use help APTMotor to obtain full list of hardware (HW) supported.
        # TDC001 (the spinner motor we have) is number 31
        """
        APT crap, see https://github.com/mcleung/PyAPT and
        https://www.thorlabs.com/Software/Motion%20Control/APT_Communications_Protocol_Rev_19.pdf
        """
    
        # Obtain current position of motor1
        print(Motor1.getPos())
    
        # Home Motor1
        Motor1.mbAbs(MA)
    
        # Wait for it to home
        time.sleep(5)
        print(Motor1.getPos())
        
        
        """ 
        initialise Motor 2
        """
        print "Motor 2:"
        Motor2 = APTMotor(83854321, HWTYPE=31)
        print(Motor2.getPos())
    
    
        """
        jog motor 1
        """
        for i in range(RA):
            MA = MA+1;
            Motor1.mbRel(1)
            time.sleep(.5)
            print "Motor 1:"
            print(Motor1.getPos())
    
            """
            Home Motor 2
            """
            MB = 30  # *************  Re initialise at the beginning of each loop!!!!! ************************
            print "Motor 2:"
            Motor2.mbAbs(MB)
            time.sleep(5)
            print(Motor2.getPos())
        
            """
            jog motor 2
            """
            for i in range(RB):
                MB = MB+1;
                Motor2.mbRel(1)
                time.sleep(.5)
                print(Motor2.getPos())
    
                """
                read and append data
                """
                # Adjust offsets and read first set of data
                scopedata, maxd_c1, mind_c1, offset_c1 = OffsetBattering (1, maxd_c1, mind_c1, offset_c1)
                scopedata2, maxd_c2, mind_c2, offset_c2 = OffsetBattering (2, maxd_c2, mind_c2, offset_c2)
                scopedata3, maxd_c3, mind_c3, offset_c3 = OffsetBattering (3, maxd_c3, mind_c3, offset_c3)
                
                # Read second set of data
                scope.WriteString("ARM",1) # Arm the scope
                scope.WriteString("FRTR", 1) # Trigger
                scope.WriteString("C1:INSP? 'SIMPLE'", 1) # Ask for data
                scopedata_2 = scope.GetScaledWaveform("C1", 24727*4, 0) # Get Data
                scope.WriteString("C2:INSP? 'SIMPLE'", 1) # Ask for data
                scopedata2_2 = scope.GetScaledWaveform("C2", 24727*4, 0) # Get Data
                scope.WriteString("C3:INSP? 'SIMPLE'", 1) # Ask for data
                scopedata3_2 = scope.GetScaledWaveform("C3", 24727*4, 0) # Get Data

                # Read third set of data
                scope.WriteString("ARM",1) # Arm the scope
                scope.WriteString("FRTR", 1) # Trigger
                scope.WriteString("C1:INSP? 'SIMPLE'", 1) # Ask for data
                scopedata_3 = scope.GetScaledWaveform("C1", 24727*4, 0) # Get Data
                scope.WriteString("C2:INSP? 'SIMPLE'", 1) # Ask for data
                scopedata2_3 = scope.GetScaledWaveform("C2", 24727*4, 0) # Get Data
                scope.WriteString("C3:INSP? 'SIMPLE'", 1) # Ask for data
                scopedata3_3 = scope.GetScaledWaveform("C3", 24727*4, 0) # Get Data
    
                filename = "S%iS%i.csv"%(MA,MB)
                with open (filename, 'wb') as datafile:
                    wr = csv.writer(datafile)
                    wr1 = wr.writerow(scopedata)
                    wr2 = wr.writerow(scopedata2)
                    wr3 = wr.writerow(scopedata3)
                    wr4 = wr.writerow(scopedata_2)
                    wr5 = wr.writerow(scopedata2_2)
                    wr6 = wr.writerow(scopedata3_2)
                    wr7 = wr.writerow(scopedata_3)
                    wr8 = wr.writerow(scopedata2_3)
                    wr9 = wr.writerow(scopedata3_3)
    
        """
        disconnect
        """
        Motor1.cleanUpAPT()
        Motor2.cleanUpAPT()
        scope.Disconnect()
        
    except:
        Motor1.cleanUpAPT()
        Motor2.cleanUpAPT()
        scope.Disconnect()
        raise