#!/usr/bin/env python
# *-----------------------------------------------------------------------*
# |                                                                       |
# |  Copyright (c) 2024 by Paul Scherrer Institute (http://www.psi.ch)    |
# |                                                                       |
# |              Author Thierry Zamofing (thierry.zamofing@psi.ch)        |
# *-----------------------------------------------------------------------*
'''
hit and return motion simulation
'''

import logging
_log=logging.getLogger(__name__)

import os, sys, time, re
import numpy as np
import matplotlib.pyplot as plt


class HitReturnEval:

  def initPlt(self,grid):
    nx,ny=grid['count']
    fig=plt.figure()
    ax = fig.add_subplot(1,1,1)
    #ax.invert_xaxis()
    ax.invert_yaxis()
    #hl=ax[0].plot(x, y, color=col)
    pts=np.array(((0,0),(nx-1,0),(nx-1,ny-1),(0,ny-1),(0,0)))
    h3=ax.plot(pts[:,0],pts[:,1],'g-')
    h2=ax.plot(pts[:,0],pts[:,1],'y-')
    h1=ax.plot(pts[:,0],pts[:,1],'r.')
    #cid = fig.canvas.mpl_connect('button_press_event', self.onclick)
    #fig.obj=self
    plt.axis('equal')
    plt.ion()
    fig.show()
    self._plt=(fig,ax,h1[0],h2[0])

  def updatePlt(self,pts):
    fig, ax, h1, h2=self._plt
    pts=np.array(pts)
    h1.set_xdata((pts[:,0]))
    h1.set_ydata((pts[:,1]))
    h2.set_xdata((pts[:,0]))
    h2.set_ydata((pts[:,1]))
    plt.draw()
    plt.pause(.01)
    #fig.show()
    #fig.update()

  def setup_motion(self,prgId=2,fnPrg=None,mode=0,**kwargs):
    #vgrid:  grid parameters: {orig:(0,0),pitch(10,10),cnt:(10,10),mode:0}
    #vp0:    x/y koordinates of well to start (default 0,0)
    #vsz:    size of wells to cycle (must be 2n x m)
    g=kwargs['grid']
    ox,oy=g['pos']      #origin position in um (or counts if scaled)
    px,py=g['pitch']    #pitch to next position in um (or 1 if scaled)
    nx,ny=g['count']     #total count of wells
    sx,sy=kwargs['ssz']  #section size (in wells)
    #smv=kwargs['smv']    #time(in num of shots) to move to next section (horiz/vert)
    #swait=kwargs['smv']  #wait after <swait> section for motion trigger. swait=0:never wait

    #fast y (up/down) slow x(left/right)
    #    8
    #    7       8         8
    #    6       7         7
    #  0-5     0-6---5   0-6-------5
    #  | |     | 3-4 |   | 3-4 3-4 |
    #  | |     | | | |   | | | | | |
    #  | |     | | | |   | | | | | |
    #  | |     | | | |   | | | | | |
    #  1-2     1-2 1-2   1-2 1-2 1-2

    r=0             #row
    x0,y0,n=0,-1,0  # counter well in region
    dx,dy=0,1       # motion pitch
    x1,y1=0,0       #counter region start
    mv=list()
    while(True):
      x0+=dx;y0+=dy
      print(f'{x0+x1} {y0+y1} ({dx}|{dy}|{n})')
      mv.append((x0+x1+.1*n,y0+y1+.1*n,n))
      self.updatePlt(mv)
      if y0==sy-1:
        if dy==1:   #(1)
          dx,dy=1,0
        else:
          dx,dy=0,-1#(2)
      elif y0==1 and x0<sx-1:
        if dy==-1:  #(3)
          dx,dy=1,0
        else:
          dx,dy=0,1 #(4)
      elif y0==0:
        if n==1 and x0==1: #(7 or 8)
          if r%2==0 and x1+2*sx<=nx:#(7 even)
            x1+=sx
          elif r%2==1 and x1>0:#(7 odd)
            x1-=sx
          elif y1+2*sy<=ny:#(8)
            y1+=sy;r+=1
          else:
            break #finished whole grid
          print('next regionA')
          x0,y0,n=0,0,0
          dx,dy=0,0
        elif x0==1: #(6)
          dx,dy=-1,0;n+=1
        elif x0==0: #(0)
          dx,dy=0,1
        else: #(5)
          dx,dy=-1,0

  def plot_trajectory(self):
    pts = self.pts  # X,Y array
    rec = self.rec  # yA,xA,yD,xD,trig
    fig = plt.figure('trajectory')
    ax = fig.add_subplot(1, 1, 1)
    ax.invert_xaxis()
    ax.invert_yaxis()
    # hl=ax[0].plot(x, y, color=col)
    hl = ax.plot(pts[:, 0], pts[:, 1], 'r.', label='points')
    hl += ax.plot(pts[:, 0], pts[:, 1], 'y--', label='direct')
    hl += ax.plot(rec[:, 3], rec[:, 2], 'b-', label='DesPos')  # desired path
    hl += ax.plot(rec[:, 1], rec[:, 0], 'g.', label='ActPos')  # actual path
    try:
      pvt = self.pvt
    except AttributeError:
      pass
    else:
      hl = ax.plot(pvt[1], pvt[2], 'c--', label='SimPos')  # simulated path

    fig2 = plt.figure('time line')
    ax2 = fig2.add_subplot(1, 1, 1)
    hl2 = ax2.plot(rec[:, 2], 'r-', label='desPos Mot1')
    hl2 += ax2.plot(rec[:, 3], 'g-', label='desPos Mot2')

    idxTrigger = rec[:, 4]
    idxTrigger = np.where(np.diff(idxTrigger) == 1)[0] + 1
    if idxTrigger.shape[0] > 0:
      hl += ax.plot(rec[idxTrigger, 1], rec[idxTrigger, 0], 'xr', label='trig')  # actual path
      hl2 += ax2.plot(rec[:, 4]*10, 'b-', label='trigger')

    ax.xaxis.set_label_text('x-pos um')
    ax.yaxis.set_label_text('y-pos um')
    ax.axis('equal')
    ax.legend(loc='best')
    # cid = fig.canvas.mpl_connect('button_press_event', self.onclick)
    # fig.obj=self
    ax2.legend(loc='best')
    plt.show(block=False)

from optparse import OptionParser, IndentedHelpFormatter
if __name__=='__main__':
  logging.basicConfig(level=logging.INFO, format='%(levelname)s:%(module)s:%(lineno)d:%(funcName)s:%(message)s ')
  import argparse

  def parse_args():
    'main command line interpreter function'
    (h, t)=os.path.split(sys.argv[0]);cmd='\n  '+(t if len(h)>3 else sys.argv[0])+' '
    exampleCmd=('-v0xff',
                '--host=SAR-CPPM-EXPMX1 -v0x5d',
                '--host=localhost:10001:10002 -v0x59',
                '--host=SAR-CPPM-EXPMX1 -v0x5d -m5',
                )
    epilog=__doc__+'\nExamples:'+''.join(map(lambda s:cmd+s, exampleCmd))+'\n '
    parser=argparse.ArgumentParser(epilog=epilog, formatter_class=argparse.RawDescriptionHelpFormatter)
    parser.add_argument("-m", "--mode", type=lambda x:int(x, 0), help="mode default=0x%(default)x", default=4)
    parser.add_argument("-s", "--sync", type=lambda x:int(x, 0), help="sync default=0x%(default)x", default=2)
    parser.add_argument('-v', '--verbose', type=lambda x:int(x, 0), dest='verbose', help='verbosity bits (see below) default=0x%(default)x', default=0x00)
    parser.add_argument('--host', help='hostname', default=None)

    args=parser.parse_args()
    _log.info('Arguments:{}'.format(args.__dict__))

    obj=HitReturnEval()
    grid={'pos':(-1000, -1200), 'pitch':(120, 120), 'count':(20, 25)}
    grid={'pos':(0, 0), 'pitch':(1, 1), 'count':(16, 20)}
    obj.initPlt(grid)
    obj.setup_motion(grid=grid,ssz=(4,5))
    obj.initPlt(grid)
    obj.setup_motion(grid=grid,ssz=(6,10))
    obj.initPlt(grid)
    obj.setup_motion(grid=grid,ssz=(2,5))
    print('done')

#------------------ Main Code ----------------------------------

  #dp=DebugPlot('/tmp/shapepath.npz');dp.plot_gather(mode=11);plt.show()
  #exit(0)
  #ssh_test()
  ret=parse_args()