#!/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 matplotlib.widgets import Slider

from utilities import *

d2r=2*np.pi/360


#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(x,y,z):
    m=np.identity(4)
    m[0:3, 3] =(x,y,z)
    return m

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 sequencer(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)


  @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

    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


  def test_coord_trf(self):
      self.calcParam()
      param = self.param
      y, w, dx, dz = (4.3, .1, 0.2, 0.3)
      print 'input  : fy:%.3g dx:%.3g dz:%.3g w:%.3g' % (y, dx, dz, w / d2r)
      (fy, w, cx, cz) = self.inv_transform(y, w, dx, dz)
      print 'inv_trf: fy:%.3g cx:%.3g cz:%.3g w:%.3g' % (fy, cx, cz, w / d2r)
      (y, w, dx, dz) = self.fwd_transform(fy, w, cx, cz)
      print 'fwd_trf: fy:%.3g dx:%.3g dz:%.3g w:%.3g' % (y, dx, dz, w / d2r)

      # plt.ion()
      # fig = plt.figure()

      # a=anim.FuncAnimation(fig,self.my_anim_func3,100,fargs=(horig,pt),interval=20,blit=False)
      # plt.show()

  #  def my_anim_func3(self,idx):
  #    self.hCrist,pt=self.pltCrist(cx=0,ty=0,cz=0,w=10*idx*d2r,h=self.hCrist)


  def interactive_fy_cx_cz_w(self):
    fig = plt.figure()
    self.manip=False#True#False
    self.ax=ax=plt3d.Axes3D(fig,[0.02, 0.15, 0.96, 0.83])
    ax.set_xlabel('Z');ax.set_ylabel('X');ax.set_zlabel('Y')
    ax.view_init(elev=14., azim=10)
    self.calcParam()
    param=self.param
    # param[i]=(z_i, y_i, x_i, r_i,phi_i)
    ctr=param[:,0:3].mean(0)[::-1]
    self.axSetCenter(ctr,10)

    axFy = plt.axes([0.1, 0.01, 0.8, 0.02])
    axCx = plt.axes([0.1, 0.04, 0.8, 0.02])
    axCz = plt.axes([0.1, 0.07, 0.8, 0.02])
    axW  = plt.axes([0.1, 0.10, 0.8, 0.02])
    if self.manip:
      lz=ax.get_xlim()
      lx=ax.get_ylim()
      ly=ax.get_zlim()
    else:
      lx = ly=lz=[-5,5]
    self.sldFy = sFy = Slider(axFy, 'fy', ly[0], ly[1], valinit=(ly[0]+ly[1])/2.)
    self.sldCx = sCx = Slider(axCx, 'cx', lx[0], lx[1], valinit=(lx[0]+lx[1])/2.)
    self.sldCz = sCz = Slider(axCz, 'cz', lz[0], lz[1], valinit=(lz[0]+lz[1])/2.)
    self.sldW =sW  = Slider(axW, 'ang', -180., 180.0, valinit=0)
    sFy.on_changed(self.update_fy_cx_cz_w)
    sCx.on_changed(self.update_fy_cx_cz_w)
    sCz.on_changed(self.update_fy_cx_cz_w)
    sW.on_changed(self.update_fy_cx_cz_w)

    hCrist,pt=self.pltCrist()

    self.hCrist=hCrist;self.fig=fig
    # 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.sin(phi_i) # 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.cos(phi_i)  # z= z_i+r_i*sin(phi_i*w)
    print p
    ofs=(p[1]+p[0])/2. # = center of the cristal

    m=Trf.trans(*ofs); self.hOrig=self.pltOrig(m)
    plt.show()

  def update_fy_cx_cz_w(self,val):
    fy = self.sldFy.val
    cx = self.sldCx.val
    cz = self.sldCz.val
    w = self.sldW.val
    if self.manip:
      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)  # 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)  # z= z_i+r_i*sin(phi_i*w)
      print p
      ofs = (p[1] + p[0]) / 2.  # = center of the cristal

      m = Trf.trans(cx,fy,cz)
      m= m.dot(Trf.rotY(w*d2r))
      self.hOrig = self.pltOrig(m,self.hOrig)
    else:
      self.hCrist,pt=self.pltCrist(fy,cx,cz,w*d2r,self.hCrist)
    #l.set_ydata(amp * np.sin(2 * np.pi * freq * t))
    self.fig.canvas.draw_idle()

  def interactive_y_dx_dz_w(self):
    fig = plt.figure()
    self.ax=ax=plt3d.Axes3D(fig,[0.02, 0.15, 0.96, 0.83])
    ax.set_xlabel('Z');ax.set_ylabel('X');ax.set_zlabel('Y')
    ax.view_init(elev=14., azim=10)
    self.calcParam()
    param=self.param

    # param[i]=(z_i, y_i, x_i, r_i,phi_i)
    ctr=(0,0,0)
    self.axSetCenter(ctr,10)

    axY = plt.axes([0.1, 0.01, 0.8, 0.02])
    axDx = plt.axes([0.1, 0.04, 0.8, 0.02])
    axDz = plt.axes([0.1, 0.07, 0.8, 0.02])
    axW  = plt.axes([0.1, 0.10, 0.8, 0.02])

    lx=[-1,1];ly=[0,1];lz=[-1,1]
    ly = param[:,1]
    self.sldY = sY = Slider(axY, 'y', ly[0], ly[1], valinit=(ly[0]+ly[1])/2.)
    self.sldDx = sDx = Slider(axDx, 'dx', lx[0], lx[1], valinit=(lx[0]+lx[1])/2.)
    self.sldDz = sDz = Slider(axDz, 'dz', lz[0], lz[1], valinit=(lz[0]+lz[1])/2.)
    self.sldW =sW  = Slider(axW, 'ang', -180., 180.0, valinit=0)
    sY.on_changed(self.update_y_dx_dz_w)
    sDx.on_changed(self.update_y_dx_dz_w)
    sDz.on_changed(self.update_y_dx_dz_w)
    sW.on_changed(self.update_y_dx_dz_w)


    # 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.sin(phi_i) # 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.cos(phi_i)  # z= z_i+r_i*sin(phi_i*w)
    ofs=(p[1]+p[0])/2. # = center of the cristal
    print 'p, ofs',p,ofs

    m=Trf.trans(0,0,0); self.hOrig=self.pltOrig(m)
    hCrist,pt=self.pltCrist(cx=-ofs[0],fy=-ofs[1],cz=-ofs[2])

    self.hCrist=hCrist;self.fig=fig
    plt.show()

  def update_y_dx_dz_w(self,val):
    y = self.sldY.val
    dx = self.sldDx.val
    dz = self.sldDz.val
    w = self.sldW.val
    w=w*d2r
    (fy, w, cx, cz)=self.inv_transform(y,w,dx,dz)

    self.hCrist,pt=self.pltCrist(-fy,-cx,-cz,w,self.hCrist)
    self.fig.canvas.draw_idle()

  def fwd_transform(self,fy,w,cx,cz):
    #cx,cy: coarse stage
    #input:  fy,w,cx,cz
    #output: y,w,dx,dz
    # param[i]=(z_i, y_i, x_i, r_i,phi_i)

    param=self.param
    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.sin(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.cos(phi_i+w)  # z= z_i+r_i*sin(phi_i*w)
    v=p[1]-p[0]
    #for y = 0..1:
    #v=v*y
    #for y = y_0..y_1:
    y_0=param[0,1];y_1=param[1,1];v=v*(fy-y_0)/(y_1-y_0)

    v=v+p[0]
    dx=cx-v[0]
    dz=cz-v[2]
    y=v[1]
    res=(y,w,dx,dz)
    return res

  def inv_transform(self,y,w,dx=0,dz=0):
    #input:  y,w,dx,dz
    #output: y,w,cx,cz
    #dx,dy: deviation from cristal center line
    param=self.param
    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.sin(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.cos(phi_i+w)  # z= z_i+r_i*sin(phi_i*w)
    v=p[1]-p[0]
    #for y = 0..1:
    #v=v*y
    #for y = y_0..y_1:
    y_0=param[0,1];y_1=param[1,1];v=v*(y-y_0)/(y_1-y_0)

    v=v+p[0]
    cx=dx+v[0]
    cz=dz+v[2]
    fy=v[1]
    res=(fy,w,cx,cz)
    return res


  def calcParam(self):
    n = 3.;
    per = 1.;
    w = 2 * np.pi * per / n * np.arange(n)
    p = ((2.3, .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
    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:]
    print param


  def pltOrig(self,m,h=None):
    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
    if h is None:
      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
    else:
      hr, hg, hb = h
      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]))
      return h

  def pltCrist(self,fy=0,cx=0,cz=0,w=0,h=None):
    #h are the handles
    if h is None:
      h=[] #handels
      ax=self.ax
      param=self.param
      pt = np.ndarray((4, 3))
      for i in range(2):
        (z, y, x, r, phi) = param[i]
        x+=cx;y+=fy;z+=cz;phi+=w
        pt[i] = (x, y, z)
        pt[i + 2] = (x + r * np.sin(phi), y, z + r * np.cos(phi))
        obj = mpl.patches.Circle((z, x), r, facecolor=mpl.colors.colorConverter.to_rgba('r', alpha=0.3))
        h1=ax.add_patch(obj)
        h2=plt3d.art3d.pathpatch_2d_to_3d(obj, z=y, zdir="z")
        h.append(obj)
      #hs=ax.scatter(pt[:, 2], pt[:, 0], pt[:, 1])
      hs=ax.plot(pt[:, 2], pt[:, 0], pt[:, 1],'.')
      hp=ax.plot(pt[2:, 2], pt[2:, 0], pt[2:, 1])
      h+=(hs[0],hp[0])
    else:
      ax=self.ax
      param=self.param
      pt = np.ndarray((4, 3))
      for i in range(2):
        (z, y, x, r, phi) = param[i]
        x+=cx;y+=fy;z+=cz;phi+=w
        pt[i] = (x, y, z)
        pt[i + 2] = (x + r * np.sin(phi), y, z + r * np.cos(phi))
        h[i].remove()
        obj = mpl.patches.Circle((z, x), r, 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")
        h[i]=obj
      h[2].set_data(pt[:, 2], pt[:, 0])#, pt[:, 1]))
      h[2].set_3d_properties(pt[:, 1])

      h[3].set_data(pt[2:, 2], pt[2:, 0])#, pt[:, 1]))
      h[3].set_3d_properties(pt[2:, 1])
      #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]))

      #h[3].set_3d_properties(zs=pt[:, 1]))
      #h[4].set_data((pt[2:, 2], pt[2:, 0]))#, pt[2:, 1]))
      #hp=ax.plot(pt[2:, 2], pt[2:, 0], pt[2:, 1], label='zs=0, zdir=z')
      #h+=(hs,hp[0])
    return (h,pt)

  @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)


  def axSetCenter(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 my_anim_func(idx,hs,horig):
  print('anim')
  a=idx*.01*2*np.pi
  m=Trf.rotY(a)
  hs.pltOrig(m,horig)

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

    hs=HelicalScan(args)
    #hs.sequencer()
    hs.test_find_rot_ctr()
    hs.test_find_rot_ctr(n=5. ,per=1.,bias=2.31,ampl=4.12,phi=24.6)
    hs.test_coord_trf()
    hs.interactive_fy_cx_cz_w()
    hs.interactive_y_dx_dz_w()


  #------------------ Main Code ----------------------------------
  #ssh_test()
  ret=parse_args()
  exit(ret)