# script to maximize the output power of the piezo
import serial
import time
import os, sys, subprocess
import numpy as np

def slowDownJog(ser):
    ser.write('1SU50\r\n')
    time.sleep(0.1)
    ser.write('1SU-50\r\n')
    time.sleep(0.1)
    ser.write('2SU50\r\n')
    time.sleep(0.1)
    ser.write('2SU-50\r\n')
    time.sleep(0.1)

def goToPosn(axis0, posn, ser):

    print('Going to position = ' + str(posn))
    currPosn = getPosn(axis0, ser)
    if posn > currPosn:
        dirn = 1
    else:
        dirn = -1    
    dirn0 = np.copy(dirn)
        
    # go from start to finish at fast jog                                    
    mediumJog(axis0, dirn, ser)
    while np.abs(currPosn - posn) > 2:
        currPosn = getPosn(axis0, ser)
        print(currPosn)
        time.sleep(0.1)
        if posn > currPosn:
	    dirn = 1
        else:
            dirn = -1

        # if sign switches then terminate loop
        if dirn != dirn0:
            break
            
    stopAxis(axis0, ser)


def goToPosnFast(axis0, posn, ser):

    print('Going to position = ' + str(posn))
    currPosn = getPosn(axis0, ser)
    if posn > currPosn:
        dirn = 1
    else:
        dirn = -1    
    dirn0 = np.copy(dirn)
        
    # go from start to finish at fast jog                                    
    fastJog(axis0, dirn, ser)
    while np.abs(currPosn - posn) > 20:
        currPosn = getPosn(axis0, ser)
        print(currPosn)
        time.sleep(0.1)
        if posn > currPosn:
	    dirn = 1
        else:
            dirn = -1

        # if sign switches then terminate loop
        if dirn != dirn0:
            break
            
    stopAxis(axis0, ser)


def goToMax(axis0, ser, delta):

    dirn0 = 1
    deltaCurr = 500
    # start the position and power logging                                                                   
    try:
        # go from start to finish at fast jog                                                                
        mediumJog(axis0, dirn0, ser)
        navg = 10
        swCount = 0
        while abs(deltaCurr) > delta:
            res = np.array([1,2.])
            for aa in range(2):
                a0 = 0
                for ii in range(navg):
                    a0 = a0 + getPower()
                res[aa] = a0
            deltaCurr = np.float(np.diff(res)[0])/navg
            print('Axis ' + str(axis0) + ' average change in power = ' + str(deltaCurr))

            if deltaCurr < 0:
                dirn0 = dirn0*-1
                mediumJog(axis0, dirn0, ser)
                swCount = swCount + 1
                print('Switching direction')

            if swCount > 3:
                break
        stopAxis(axis0,ser)
        
                
    except:
        stopAxis(1,ser)
	time.sleep(0.02)
        stopAxis(2,ser)
	time.sleep(0.02)

    

def fastJog(axis0, dirn, ser):
    axisSTR = str(axis0)
    if dirn < 0:
        dirSTR = '-'
    else:
        dirSTR = ''

    strOUT = axisSTR + 'MV' + dirSTR + '2' + '\r\n'
    ser.write(strOUT)
    time.sleep(0.01)

def mediumJog(axis0, dirn, ser):
    axisSTR = str(axis0)
    if dirn < 0:
        dirSTR = '-'
    else:
        dirSTR = ''

    strOUT = axisSTR + 'MV' + dirSTR + '1' + '\r\n'
    ser.write(strOUT)
    time.sleep(0.01)


    
def stopAxis(axis0, ser):
    axisSTR = str(axis0)

    strOUT = axisSTR + 'ST' + '\r\n'
    ser.write(strOUT)
    time.sleep(0.01)

def getPosn(axis0, ser):
    axisSTR = str(axis0)

    strOUT = axisSTR + 'TP' + '\r\n'
    ser.write(strOUT)
    time.sleep(0.01)

    a = ser.readline()
    b = a.replace('\n', '')
    b = b.replace(axisSTR+'TP' , '')
    bOUT = int(b)
    return bOUT


def getPower():
    a = subprocess.check_output('caget -t C4:TST-PD_RESPONSE', shell=True)
    b = a.replace('\n', '')
    bOUT = np.float(b)
    return bOUT


# create the main function
def main(argv):

    print('Opening serial')
    ser = serial.Serial('/dev/ttyUSB0', baudrate=921600, timeout=0.5)

    # set up for medium jog
    slowDownJog(ser)

    # start the position and power logging
    b = getPosn(1,ser)
    a = getPower()
    powArr = np.array([[b],[a]])

    try: 
        # go from start to finish at fast jog
        for ii in range(2):
            goToMax(1,ser,0.1E-4)
            goToMax(2,ser,0.1E-4)





    except Exception as (e):
        print(e)
        print('Error')
        stopAxis(1,ser)
        time.sleep(0.02)
        stopAxis(2,ser)
        time.sleep(0.02)

    ser.close()


if __name__ == "__main__":
    try:
        main(sys.argv[1:])
    finally:
        print('Export thread')