#!/usr/bin/env python
# *-----------------------------------------------------------------------*
# |                                                                       |
# |  Copyright (c) 2017 by Paul Scherrer Institute (http://www.psi.ch)    |
# |                                                                       |
# |              Author Thierry Zamofing (thierry.zamofing@psi.ch)        |
# *-----------------------------------------------------------------------*
'''
tools to setup and execute a helical scan of a cristal
#THIS IS JUST TESTING CODE TO SOLVE FINDING THE ROTATION CENTER
'''

import os, sys, json
import numpy as np
import matplotlib as mpl
import matplotlib.pyplot as plt
import mpl_toolkits.mplot3d as plt3d
import matplotlib.animation as anim

from utilities import *

#ax.set_xlabel('Z');ax.set_ylabel('X');ax.set_zlabel('Y')
#plot coordinates: X Y Z
#data              Z X Y

class Trf:
  #https://stackoverflow.com/questions/6802577/python-rotation-of-3d-vector
  @staticmethod
  def rotZ(rad):
    """ Return matrix for rotating about the z-axis by 'radians' radians """
    c = np.cos(rad)
    s = np.sin(rad)
    m=np.identity(4)
    m[0:2,0:2]=((c, -s),(s, c))
    return m

  @staticmethod
  def rotY(rad):
    """ Return matrix for rotating about the z-axis by 'radians' radians """
    c = np.cos(rad)
    s = np.sin(rad)
    m=np.identity(4)
    m[np.ix_((0,2),(0,2))]=((c, -s),(s, c))
    return m

  @staticmethod
  def trans(v):
    m=np.identity(4)
    m[0:3, 3] = v
    return m

class ExtAxes():#mpl.axes._subplots.Axes3DSubplot):
  def __init__(self,ax):
    self.ax=ax
    ax.set_xlabel('Z');ax.set_ylabel('X');ax.set_zlabel('Y')
    ax.view_init(elev=14., azim=10)

  def setCenter(self,v,l):
    ax=self.ax
    #v=center vector, l= length of each axis
    l2=l/2.
    ax.set_xlim(v[2]-l2, v[2]+l2);
    ax.set_ylim(v[0]-l2, v[0]+l2);
    ax.set_zlim(v[1]-l2, v[1]+l2)

  def pltOrig(self,m):
    ax=self.ax
    # m is a 4x4 matrix. the transformed matrix
    r=m[:3,0] #1st
    g=m[:3,1] #2nd
    b=m[:3,2] #3rd
    o=m[:3,3] #origin
    hr=ax.plot((o[2],o[2]+r[2]), (o[0],o[0]+r[0]), (o[1],o[1]+r[1]), 'r')
    hg=ax.plot((o[2],o[2]+g[2]), (o[0],o[0]+g[0]), (o[1],o[1]+g[1]), 'g')
    hb=ax.plot((o[2],o[2]+b[2]), (o[0],o[0]+b[0]), (o[1],o[1]+b[1]), 'b')
    return hr+hg+hb  # this is a list

def my_anim_func(idx,horig):
  print('anim')
  a=idx*.01*2*np.pi
  m=Trf.rotY(a)
  r = m[:3, 0]  # 1st
  g = m[:3, 1]  # 2nd
  b = m[:3, 2]  # 3rd
  o = m[:3, 3]  # origin
  hr,hg,hb=horig
  hr.set_data((o[2], o[2] + r[2]), (o[0], o[0] + r[0]))
  hr.set_3d_properties((o[1], o[1] + r[1]))
  hg.set_data((o[2], o[2] + g[2]), (o[0], o[0] + g[0]))
  hg.set_3d_properties((o[1], o[1] + g[1]))
  hb.set_data((o[2], o[2] + b[2]), (o[0], o[0] + b[0]))
  hb.set_3d_properties((o[1], o[1] + b[1]))

class HelicalScan:
    def __init__(self,args):
      if args.cfg:
        fh=open(args.cfg,'r')
        s=fh.read()
        cfg=json.loads(s, object_hook=ConvUtf8)
        s=json.dumps(cfg, indent=2, separators=(',', ': '));print(s)
      else:
        fn='/tmp/shapepath4'
        #fn='/home/zamofing_t/Documents/prj/SwissFEL/epics_ioc_modules/ESB_MX/python/data/'+time.strftime('%y-%m-%d-%H_%M_%S')
        #cfg = {"sequencer": ['gen_grid_points(w=5,h=5,pitch=100,rnd=0.4)', 'sort_points()','gen_prog(file="'+fn+'.prg",host="SAR-CPPM-EXPMX1",mode=1,pt2pt_time=10,cnt=1)', 'plot_gather("'+fn+'.npz")']}
        #cfg = {"sequencer": ['test_find_rot_ctr()']}
        #cfg = {"sequencer": ['test_find_rot_ctr(n=5. ,per=1.,bias=2.31,ampl=4.12,phi=24.6)']}
        cfg = {"sequencer": ['test_coord_trf()']}

      self.cfg=dotdict(cfg)
      self.args=args

    def run(self):
      print('args='+str(self.args))
      print('cfg='+str(self.cfg))
      #try:
      #  self.points=np.array(self.cfg.points)
      #except AttributeError:
      #  pass
      try:
        sequencer= self.cfg.pop('sequencer')
      except KeyError:
        print('no command sequence to execute')
      else:
        dryrun=self.args.dryrun
        for cmd in sequencer:
          print('>'*5+' '+cmd+' '+'<'*5)
          if not dryrun:
            eval('self.' + cmd)

    def test_coord_trf(self):
      d2r=2*np.pi/360
      plt.ion()
      fig = plt.figure()
      #ax = fig.gca(projection='3d')
      ax = fig.add_subplot(1,1,1,projection='3d')
      extAx=ExtAxes(ax)
      extAx.setCenter((0,5,15),10)
      n = 3.; per = 1.; w = 2*np.pi*per/n*np.arange(n)
      p=((2.3,0.71,4.12,10.6*d2r),(6.2,.45,3.2,45.28*d2r))       #(y, bias, ampl, phi)
      self.param=param=np.ndarray((len(p),5))
      z=14.5 # fix z position
      pt=np.ndarray((4,3))
      for i in range(2):
        (y, bias, ampl, phi) =p[i]
        x= ampl * np.cos(w+phi) + bias
        print('yMeas_%d='%i+str(y)+' xMeas_%d='%i+str(x))
        #param[i]=(z_i, y_i, x_i, r_i,phi_i)
        param[i,0] =z
        param[i,1] =y
        param[i,2:]=HelicalScan.meas_rot_ctr(x) #(bias,ampl,phase)
        (bias, ampl, phase)=param[i][2:]
        pt[i]=(bias, y, z)
        pt[i+2]=(bias+ampl*np.cos(phase),y, z+ampl*np.sin(phase))
        obj = mpl.patches.Circle((z,bias), ampl, facecolor=mpl.colors.colorConverter.to_rgba('r', alpha=0.3))
        ax.add_patch(obj)
        plt3d.art3d.pathpatch_2d_to_3d(obj, z=y, zdir="z")
      ax.scatter(pt[:,2], pt[:,0], pt[:,1])
      ax.plot(pt[2:,2], pt[2:,0], pt[2:,1], label='zs=0, zdir=z')
      print param
      #plt.show()
      #m=np.identity(4); horig=extAx.pltOrig(m)
      m=Trf.trans((0,0,z)); horig=extAx.pltOrig(m)
      #self.fwd_transform(y         ,w       ,cx         ,cz)
      #                    y_0        ,0deg       ,x_0    ,z_0)
      m=self.fwd_transform(param[0][1],0,pt[2][0],pt[2][2],extAx)
      m=self.fwd_transform(param[0][1],0,pt[2][0],pt[2][2],extAx)

      m=self.fwd_transform(param[0][1],20*d2r,pt[2][0],pt[2][2])
      extAx.pltOrig(m)

      #my_anim_func(0,horig)
      a=anim.FuncAnimation(fig,my_anim_func,25,fargs=(horig,),interval=50,blit=False)
      plt.show()
      #                  y_0        ,120deg    ,x_0        ,z_0)
      self.fwd_transform(param[0][1],2*np.pi/3.,param[0][2],param[0][0])
      #self.fwd_transform(param[1][1],0.,param[1][2],param[1][3])

    def fwd_transform(self,y,w,cx,cz,extAx):
      #cx,cy: coarse stage
      #input:  y,w,cx,cz
      #output: y,w,dx,dz
      m=Trf.trans((cx,y,cz))
      m=m.dot(Trf.rotY(w))
      extAx.pltOrig(m)
      extAx.setCenter(m[0:3,3],1)
      #TODO: NOT WORKING AT ALL NOW...
      param=self.param
      # param[i]=(z_i, y_i, x_i, r_i,phi_i)
      p=np.ndarray((param.shape[0], 3))
      for i in range(2):
        (z_i, y_i, x_i, r_i, phi_i)=param[i]
        p[i,0]=x_i+r_i*np.cos(phi_i+w) # x= x_i+r_i*cos(phi_i+w)+cx
        p[i,1]=y_i # y= y_i
        p[i,2] =z_i+r_i*np.sin(phi_i+w)  # z= z_i+r_i*sin(phi_i*w)
      print p
      v=p[1]-p[0]
      v=v/np.sqrt(v.dot(v)) # v/|v|
      y_0=param[0][1]
      y_1=param[1][1]
      v=v*(y-y_0)/(y_1-y_0) # v(y)=v*(v-y_0)/(y_1-y_0)
      v=p[0]+v
      cz + cx

      #v=v/abs(v)
      print v


      return m

    def inv_transform(y,phi,dx=0,dz=0):
      #input:  y,w,dx,dz
      #output: y,w,cx,cz
      m=np.identity(4)
      #dx,dy: deviation from cristal center line
      #ps= #x,y,z
      #returns y,phi,cx,cz
      pass


    @staticmethod
    def meas_rot_ctr(y,per=1):
      # find the amplitude bias and phase of an equidistant sampled sinus
      # it needs at least 3 measurements e.g. at 0,120 240 deg or 0 90 180 270 deg
      # per is the number of persiods, default is 1 period =360 deg
      n=len(y)
      f = np.fft.fft(y)
      idx=int(per)
      bias=np.absolute(f[0]/n)
      phase=np.angle(f[idx])
      ampl=np.absolute(f[idx])*2/n
      return (bias,ampl,phase)

    @staticmethod
    def test_find_rot_ctr(n=3.,per=1.,bias=4.1,ampl=2.4,phi=37):
      # find the rotation center, amplitude out of n (niminum 3) measurements
      # n    number of equidistant measurements
      # per  number of periods (full rotation of all measurements nut be a interger value for precise measurements)
      # phi  phase
      # bias bias value
      # ampl amplitude
      d2r=2*np.pi/360

      t = np.arange(n)
      w=2*np.pi*per/n*t
      y=ampl*np.cos(w+phi*d2r)+bias
      plt.figure(1)
      plt.subplot(311)
      plt.plot(t,y,'b.-')

      plt.subplot(312)
      f = np.fft.fft(y)
      plt.step(t, f.real,'b.-', t, f.imag,'r.-', where='mid')

      (bias,ampl,phase)=HelicalScan.meas_rot_ctr(y, per)
      print('bias: '+str(bias))
      print('amplitude: '+str(ampl))
      print('phase: '+str(phase*360./2/np.pi))

      plt.subplot(313)
      t2 = np.linspace(0,2*np.pi,64)
      y2=ampl*np.cos(t2+phase)+bias
      plt.plot(t2,y2,'g-')
      plt.stem(w,y,'b-')


      plt.show()
      pass


if __name__=='__main__':
  from optparse import OptionParser, IndentedHelpFormatter
  class MyFormatter(IndentedHelpFormatter):
    'helper class for formating the OptionParser'

    def __init__(self):
      IndentedHelpFormatter.__init__(self)

    def format_epilog(self, epilog):
      if epilog:
        return epilog
      else:
        return ""

  def parse_args():
    'main command line interpreter function'
    #usage: gpasciiCommunicator.py --host=PPMACZT84 myPowerBRICK.cfg
    (h, t)=os.path.split(sys.argv[0]);cmd='\n  '+(t if len(h)>3 else sys.argv[0])+' '
    exampleCmd=('-n',
                '-v15'
               )
    epilog=__doc__+'''
Examples:'''+''.join(map(lambda s:cmd+s, exampleCmd))+'\n '

    fmt=MyFormatter()
    parser=OptionParser(epilog=epilog, formatter=fmt)

    parser.add_option('-v', '--verbose', type="int", dest='verbose', help='verbosity bits (see below)', default=0)
    parser.add_option('-n', '--dryrun', action='store_true', help='dryrun to stdout')
    parser.add_option('--xy', action='store_true', help='sort x,y instead y,x')
    parser.add_option('--cfg', help='config file containing json configuration structure')

    (args, other)=parser.parse_args()
    args.other=other

    sp=HelicalScan(args)
    sp.run()
#------------------ Main Code ----------------------------------
  #ssh_test()
  ret=parse_args()
  exit(ret)