#!/usr/bin/env python

import epics
from time import sleep

class transform:
  def __init__(self, (m1, m2, m3), (t1, t2, t3), (tm1, tm2, tm3)):
    """ 
    real motors (to be transformed): m1, m2, m3
    transform motors: t1, t2, t3
    soft motors (after transformation): tm1, tm2, tm3
    """

    print "Begin Initialization"

    ### Flags for signaling changes in PVs to the main loop

    self.updateReadbackFlag = False
    self.dmovChangeFlag = False
    self.moveSoftMotorFlag = False
    self.stopMotorsFlag = False
    self.syncSoftTargetFlag = False
    self.reinitFlag = False
    # Is only used if SYNC should happen conditionally
    self.softMotorInitiatedMove = False
    self.abortFlag = False
    self.initCompleteFlag = False
    self.syncInProgress = False
    
    ### PVs to set
    
    # real motor VAL (for moving motors)
    self.m_val = []
    for m in (m1, m2, m3):
      tempStr = "%s.VAL" % m
      self.m_val.append( epics.PV(tempStr) )
    
    # real motor STOP (for stopping motors)
    self.m_stop = []
    for m in (m1, m2, m3):
      tempStr = "%s.STOP" % m
      self.m_stop.append( epics.PV(tempStr) )

    # transform motor STOP (for stopping motors)
    self.t_stop = []
    for t in (t1, t2, t3):
      tempStr = "%s.STOP" % t
      self.t_stop.append( epics.PV(tempStr) )

    # soft motor readback PVs (for outputing transformed position readback)
    self.tm_rdbk = []
    for tm in (tm1, tm2, tm3):
      tempStr = "%s:rdbk.VAL" % tm
      self.tm_rdbk.append( epics.PV(tempStr) )
 
    # soft motor dmov PVs (for providing calculated moving indicator)   
    self.tm_dmov = []
    for tm in (tm1, tm2, tm3):
      tempStr = "%s:dmov.VAL" % tm
      self.tm_dmov.append( epics.PV(tempStr) )
    
    # soft motor sync PVs (for synchronizing target positions)   
    self.tm_sync = []
    for tm in (tm1, tm2, tm3):
      tempStr = "%s.SYNC" % tm
      self.tm_sync.append( epics.PV(tempStr) )

    # soft motor proc PVs (for synchronizing target positions)   
    self.tm_proc = []
    for tm in (tm1, tm2, tm3):
      tempStr = "%s.PROC" % tm
      self.tm_proc.append( epics.PV(tempStr) )

    ### PVs to monitor

    # real motor RBV (for calculating transformed position)
    self.m_rbv = []
    for m in (m1, m2, m3):
      tempStr = "%s.RBV" % m
      self.m_rbv.append( epics.PV(tempStr, callback=self.rdbkCallback) )
    
    # real motor DMOV (for calculating transform motor dmov)
    self.m_dmov = []
    for m in (m1, m2, m3):
      tempStr = "%s.DMOV" % m
      self.m_dmov.append( epics.PV(tempStr, callback=self.dmovCallback) )

    # transform motor RBV (for calculating transformed position)
    self.t_rbv = []
    for t in (t1, t2, t3):
      tempStr = "%s.RBV" % t
      self.t_rbv.append( epics.PV(tempStr, callback=self.rdbkCallback) )
    
    # transform motor DMOV (for calculating transform motor dmov)
    self.t_dmov = []
    for t in (t1, t2, t3):
      tempStr = "%s.DMOV" % t
      self.t_dmov.append( epics.PV(tempStr, callback=self.dmovCallback) )

    # soft motor VAL (monitor for move request detection)
    self.tm_val = []
    for tm in (tm1, tm2, tm3):
      tempStr = "%s.VAL" % tm
      self.tm_val.append( epics.PV(tempStr, callback=self.driveCallback, connection_callback=self.connectionCallback) )
    
    # soft motor STOP (monitor for stop request detection)
    self.tm_stop = []
    for tm in (tm1, tm2, tm3):
      tempStr = "%s.STOP" % tm
      self.tm_stop.append( epics.PV(tempStr, callback=self.stopCallback, connection_callback=self.connectionCallback) )
    
    # soft motor RBV (used for manually syncing target positions)
    self.tm_rbv = []
    for tm in (tm1, tm2, tm3):
      tempStr = "%s.RBV" % tm
      self.tm_rbv.append( epics.PV(tempStr) )

    # Manually update readbacks before a callback when they're out of sync causes motors to move
    #!print "begin manual update"
    self.updateReadbacks()
    #!print "end manual update"

    print "End Initialization"
    self.initCompleteFlag = True


  ### Callback functions (to be called when PV values change--only set flags that are watched in main loop)

  def driveCallback(self, **kw):
    #!print "driveCallback:\n", "  ", kw['pvname'], kw['value']

    #!print "_ updateReadbackFlag =", self.updateReadbackFlag
    #!print "_ dmovChangeFlag =", self.dmovChangeFlag
    #!print "_ moveSoftMotorFlag =", self.moveSoftMotorFlag
    #!print "_ stopMotorsFlag =", self.stopMotorsFlag
    #!print "_ syncSoftTargetFlag =", self.syncSoftTargetFlag
    #!print "_ reinitFlag =", self.reinitFlag
    #!print "_ softMotorInitiatedMove =", self.softMotorInitiatedMove
    #!print "_ abortFlag =", self.abortFlag
    #!print "_ initCompeteFlag =", self.initCompleteFlag
    #!print "_ syncInProgress =", self.syncInProgress

    if self.initCompleteFlag == True and self.syncInProgress == False:
      self.moveSoftMotorFlag = True
      self.softMotorInitiatedMove = True
      #!print "  move real motors because NOT syncing"
    else:
      #!print "  not moving real motors because a sync is in progress or initializing"
      pass
      
  def rdbkCallback(self, **kw):
    #!print "rdbkCallback:\n", "  ", kw['pvname'], kw['value']
    self.updateReadbackFlag = True

  def stopCallback(self, **kw):
    #!print "stopCallback:\n", "  ", kw['pvname'], kw['value']
    if kw['value'] == 1:
    	self.stopMotorsFlag = True

  def dmovCallback(self, **kw):
    print "dmovCallback:\n", "  ", kw['pvname'], kw['value']
    self.dmovChangeFlag = True

  def connectionCallback(self, **kw):
    #!print "connection change:\n", "  ", kw['pvname'], kw['conn']
    if kw['conn'] == True:
      self.reinitFlag = True
    else:
      print "\t", self.reinitFlag, kw

      # Abort instead of reconnecting for now
      #!self.abortFlag = True

  ### Functions that actually do stuff (and reset callback flags)
  
  def updateReadbacks(self):
    # Get real and transform motor readback values
    m_positions = [x.get() for x in self.m_rbv]
    t_positions = [x.get() for x in self.t_rbv]
    
    # Do the forward transform
    tm_positions = self.fwdTrans(m_positions, t_positions)
    
    #!print "updating readbacks:", tm_positions
    
    # Write the soft (transformed) readback positions
    for i in range(len(tm_positions)):
      self.tm_rdbk[i].put(tm_positions[i])
    
    self.updateReadbackFlag = False
  
  # This is only useful if LOCK = YES  
  def syncSoftTargets(self):
    # Synchronize soft targets if a real motor initated the move 
    #!if self.softMotorInitiatedMove == False:
    # Always sync soft-motor target position (is this bad for spec?)
    if True:
      #
      self.syncInProgress = True

      # The following sleep is required to prevent problems when making small moves
      ###sleep(0.05)
    
      for i in range(len(self.tm_rbv)):
        #!print i
        self.tm_sync[i].put(1)
	# A delay is necessary before setting syncInProgress to False
	###sleep(0.05)
      
      #
      sleep(0.1)
      self.syncInProgress = False
      
    # For soft-motor-initiate moves don't do anything
    else:
      # reset soft-motor-initiated-move indicator
      self.softMotorInitiatedMove = False
    
    self.syncSoftTargetFlag = False

  def updateDmov(self):
    # The following sleep is required to prevent problems when making small moves
    sleep(0.05)
    # Get real and transform motor dmov values
    m_dmovs = [x.get() for x in self.m_dmov]
    t_dmovs = [x.get() for x in self.t_dmov]
    
    print m_dmovs, t_dmovs
    
    # Determine if any motors are moving
    if ( ( 0 in m_dmovs ) or ( 0 in t_dmovs ) ):
      dmov = 0
    else:
      dmov = 1

    #!print dmov
      
    # Write the soft (transformed) dmovs
    for i in range(len(self.tm_dmov)):
      print "setting dmov for tm%i to %i" % (i, dmov)
      self.tm_dmov[i].put(dmov)
    
    ### Sanity check
    #!print "soft motor initiated move?", self.softMotorInitiatedMove
    
    # sync the target soft motor positions when done moving
    if dmov == 1:
      # Update readback positions only happens first if readback occurs before sync in main loop
      self.updateReadbackFlag = True
      self.syncSoftTargetFlag = True
    
    self.dmovChangeFlag = False
    
  def moveMotors(self):
    # Get transform motor readbacks and desired soft motor target positions
    tm_positions = [x.get() for x in self.tm_val]
    t_positions = [x.get() for x in self.t_rbv]

    # Do the inverse transform
    m_positions = self.invTrans(tm_positions, t_positions)
    
    #!print "soft motor positions: ", tm_positions
    # NOTE: may need to round real motor positions
    #!print "moving real motors to: ", m_positions
    
    # Write the real motor target positions
    for i in range(len(self.m_val)):
      #!print "Actually changing the real motor target positions"
      self.m_val[i].put(m_positions[i])
      # also manually the dmov since this needs to be in sync very quickly
      self.tm_dmov[i].put(0)
      
    self.moveSoftMotorFlag = False
    
  def stopMotors(self):
    # Send a 1 to the stop field of the real and transform motors
    for i in range(3):
      self.m_stop[i].put(1)
      self.t_stop[i].put(1)
    
    self.stopMotorsFlag = False
    
  def resyncAfterReconnect(self):
    print "resyncing motors"
    sleep(0.5)
    self.updateReadbacks()
    self.reinitFlag = False


  ### The transform functions (overload them for custom transforms)
  
  #
  # Note: transform functions must gracefully handle transform over 
  #       the position ranges available to the real motors. 
  #
  
  # Transform real coordinates into soft coordinates
  def fwdTrans(self, m_positions, t_positions):
    retval = [m_positions[x] + t_positions[x] for x in range(3)]
    return retval[:]
  
  # Transform soft coordinates into real coordinates
  def invTrans(self, tm_positions, t_positions):
    retval = [tm_positions[x] - t_positions[x] for x in range(3)]
    return retval[:]

  ### The main loop (essentially a state program)

  def main(self):
    #!print "ENTERING MAIN LOOP"
    # Reset reinitFlag so initial connect doens't cause reinit to happen
    self.reinitFlag = False
    
    while True:
      # Idle loop
      epics.poll(evt=1.e-5, iot=0.1)
      
      if self.abortFlag == True:
        break
      # The order of these if statements is important. The ones that appear first take precendence.
      if self.reinitFlag == True:
        #!print "idle: resyncing after reconnect"
	#!print "  self.reinitFlag = ", self.reinitFlag
	self.resyncAfterReconnect()
      if self.stopMotorsFlag == True:
        #!print "idle: stop motors"
        self.stopMotors()
      if self.moveSoftMotorFlag == True:
        #!print "idle: move motors"
        self.moveMotors()
      if self.dmovChangeFlag == True:
	#!print "idle: update dmov"
        self.updateDmov()
      if self.updateReadbackFlag == True:
        #!print "idle: update readback"
        self.updateReadbacks()
      if self.syncSoftTargetFlag == True:
        #!print "idle: sync soft targets"
	self.syncSoftTargets()

if __name__ == "__main__":
  prefix = 'kmp2:'
  realMotors = ["%s%s" % (prefix, x) for x in ('m1', 'm2', 'm3') ]
  transformMotors = ["%s%s" % (prefix, x) for x in ('m4', 'm5', 'm6') ]
  softMotors = ["%s%s" % (prefix, x) for x in ('lab:m1', 'lab:m2', 'lab:m3') ]
  print realMotors, transformMotors, softMotors
  MyTrans = transform(realMotors, transformMotors, softMotors)
  MyTrans.main()