#!/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 HitReturnSim:

  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.')
    ha=ax.arrow(4, 2, 0.5, 0.5, head_width=0.05, head_length=0.1, fc='k', ec='k')

    #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],ha]

  def updatePlt(self,pts,vx,vy):
    fig,ax,h1,h2,ha=self._plt
    x,y=pts[-1][:2]
    pts=np.array(pts)
    ha.remove()
    l=(vx**2+vy**2)**.5
    ha=ax.arrow(x,y,.7*vx,.7*vy, head_width=.2, head_length=0.3*l, fc='k', ec='k',zorder=5)
    self._plt[-1]=ha
    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 sim_motion(self,mode=0xff,**kwargs):
    #mode: 0x01: print hints
    #mode: 0x02: plot simulation
    #kwargs:
    #grid  : grid parameters: {orig:(0,0),pitch(10,10),cnt:(10,10),mode:0}
    #ssz   : section size (in wells)
    #smv   : time(in num of shots) to move to next section (horiz/vert)
    #        default is (ssz[0]-1,ssz[1]+1)
    #sdelay: shots count of delay. Default is ssz[0]*ssz[1]

    #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
    g=kwargs['grid']
    if mode&0x02: self.initPlt(g)
    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)
    try:
      tsx, tsy=kwargs['smv']  # time(in num of shots) to move to next section (horiz/vert)
    except KeyError as e:
      tsx=sx-1
      tsy=sy
    try:
      tsd=kwargs['sdelay']
    except KeyError as e:
      tsd=sx*sy
    tsd+=1-sx*sy  # sdl are shots to wait at position #(6)

    tx,ty=nx//sx,ny//sy #total sections
    vsx=vsy=1           # scaled velocity

    t=1             #motion pvt tome scaling
    k=-1            # key position
    cx,cy=0,0       # section counter
    x0,y0,n=0,0,0   # counter well in region
    dx,dy=0,1       # motion pitch
    vx,vy=0,vsy     # scaled velocity
    pts=list()
    while(True):
      if cy%2==0: # even rows
        x=sx*cx+x0;y=sy*cy+y0
      else:       # odd rows
        x=(tx-cx)*sx-x0-1;y=sy*cy+y0
        vx=-vx
      x0+=dx;y0+=dy
      if mode&0x01: print(f'X{x}:{vx} Y{y}:{vy} ({k}|{dx}|{dy}|{n})')
      pts.append((x+.1*n,y+.1*n,n))
      if mode&0x02: self.updatePlt(pts,vx,vy)
      if t !=1:
        if mode&0x01: print('pvt 10')
      if k>=0: #if keypoint not directly followed by an other keypoint, define future velocity
        if k in (0,4):
          vx,vy=0,vsy
        elif k==2:
          vx,vy=0,-vsy
        elif k==5:
          vx,vy=-vsx,0
        elif k==6:
          n+=1
          if tsd!=1:
            if mode&0x01: print(f'pvt {10*tsd}');t=tsd
            for i in range(1,tsd):
              pts.append((x+.9*i/tsd, y+.9*i/tsd, n))
              if mode&0x02: self.updatePlt(pts, vx, vy)

        elif k==7:
          if mode&0x01: print('next h section')
          cx+=1
          if tsx!=1:
            if mode&0x01: print(f'pvt {10*tsx}');t=tsx
            xx=(sx-1)/tsx
            if cy%2!=0: xx=-xx
            for i in range(1,tsx):
              pts.append((x+i*xx, y+.3, n))
              if mode&0x02: self.updatePlt(pts, vx, vy)
          x0,y0,n,dx,dy=0,0,0,0,0
        elif k==8:
          if mode&0x01: print('next v section')
          cy+=1
          if cy>=ty:
            if mode&0x01: print('finished whole grid')
            break
          if tsy!=1:
            if mode&0x01: print(f'pvt {10*tsy}');t=tsy
            xx=(sx-2)/tsy
            if cy%2!=1: xx=-xx
            for i in range(1,tsy):
              pts.append((x+i*xx, y+i*sy/tsy, n))
              if mode&0x02: self.updatePlt(pts, vx, vy)

          cx,x0,y0,n,dx,dy=0,0,0,0,0,0
        else:
          raise ValueError('should never happened')
        if mode&0x01: print(f'vx|vy after keypoint {k}: ({vx}|{vy})')
        k=-1
      if y0==sy-1:
        if dy==1:   #(1)
          dx,dy=1,0
          vx,vy=vsx/2,vsy/2
        else:
          k,dx,dy=2,0,-1#(2)
          vx,vy=vsx/2,-vsy/2
      elif y0==1 and x0>0 and x0<sx-1:
        if dy==-1:  #(3)
          dx,dy=1,0
          vx,vy=vsx/2,-vsy/2
        else:
          k,dx,dy=4,0,1 #(4)
          vx,vy=vsx/2,vsy/2
      elif y0==0:
        if n==1 and x0==1: #(7 or 8)
          if cx<tx-1:
            k,dx,dy=7,-1,0
            vx,vy=0,0
          else: #(8)
            k,dx,dy=8,sx-1,sy
            vx,vy=0,0
        elif x0==1: #(6)
          k,dx,dy=6,-1,0
          vx,vy=-vsx,0
        elif x0==0: #(0)
          k,dx,dy=0,0,1
          vx,vy=0,vsy
        elif x0==sx-1: #(5)
          k,dx,dy=5,-1,0
          vx,vy=-vsx/2,-vsy/2
    return pts

  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=HitReturnSim()
    grid={'pos':(-1000, -1200), 'pitch':(120, 120), 'count':(20, 25)}
    grid={'pos':(0, 0), 'pitch':(1, 1), 'count':(16, 20)}
    obj.sim_motion(grid=grid,ssz=(4,5),smv=(4-1,5+1),sdelay=4*5)
    obj.sim_motion(grid=grid,ssz=(6,10))
    obj.sim_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()