#!/usr/bin/env python
# *-----------------------------------------------------------------------*
# |                                                                       |
# |  Copyright (c) 2022 by Paul Scherrer Institute (http://www.psi.ch)    |
# |              Based on Zac great first implementation                  |
# |              Author Thierry Zamofing (thierry.zamofing@psi.ch)        |
# *-----------------------------------------------------------------------*
'''
coordinate systems, optical center, xray axis, pixel sizes etc.

ppm= pixel per mm
update_pix2pos

modes:
  0x01: update_pix2pos
  0x02: update_optical_center
'''
import logging
import numpy as np

_log=logging.getLogger(__name__)


class geometry:

  def __init__(self):
    pass

  def find_optical_center(self, p):
    # p is an array of
    # at zoom out: (p1x,p1y),(p2x,p2y),(p3x,p3y),...
    # at zoom in : (p1x,p1y),(p2x,p2y),(p3x,p3y),...
    #
    # the pixel positions are given in chip pixel coordinates (0/0= top/right)
    # and ignore roi and binning
    # the returned (cx,cy) is the pixel koordinate that does not change with zoom
    # this coordinate represents also the origin of other coordinates
    pass

  def update_optical_center(self, meas, debug=False):
    if debug:
      import sys
      import matplotlib.pyplot as plt
      if sys.gettrace(): # in debuggingmode
        plt.interactive(True) #uncomment to debug
        print('DEBUG MODE')
      fig, ax=plt.subplots()
      ax.invert_yaxis()
      ax.grid()
      ax.axis('equal')
      for k, v in opt_ctr_meas.items():
        v=np.array(v)
        # ax.plot(v[:, 0], v[:, 1])
        ax.fill(v[:, 0], v[:, 1], fill=False)

    numPoints=len(next(iter(meas.values())))
    numZoom=len(meas)
    M=np.ndarray((numPoints, numZoom, 2), np.float)
    K=np.array(tuple(meas.keys()), np.float)
    for i, (k, v) in enumerate(meas.items()):
      M[:, i, :]=np.array(v)

    if debug:
      for i in range(numPoints):
        ax.plot(M[i, :, 0], M[i, :, 1])

    # parametrizex lines:
    # (USE ONLY 2 zoom levels (first and last!!!)
    # (x,y)=(p0x+a*q, p0y+b*q, q=0 -> (p0x,p0y), q=1 (p1x,p1y)
    # -> p1x=p0x+a*1 -> a=p1x-p0x
    # -> p1y=p0y+b*1 -> b=p1y-p0y
    AB=M[:,-1,:]-M[:,0,:] # parameters a1..an b1..bn

    # least square fitting for x,y,q
    # x=p0x+a0*q -> x-a0*q=p0x
    # y=p0y+b0*q
    # x=pnx+an*q
    # y=pny+bn*q
    # [1  0 -a0]   [x]   [p0x]
    # [.. .. ..] * [y] = [..]
    # [1  0 -an]   [q]   [pnx]
    # [0  1 -b0]         [p0y]
    # [.. .. ..]         [..]
    # [0  1 -bn]         [pny]
    #    A    * aa  =  y
    A=np.ndarray((numPoints*2, 3), np.float)
    A[:numPoints,:2]=(1,0)
    A[numPoints:,:2]=(0,1)
    A[:,2]=-AB.T.ravel()
    A=np.asmatrix(A)
    y=np.asmatrix(M[:, 0, :].T.ravel()).T
    aa=(A.T*A).I*A.T*y

    if debug:
      # plot least square line fitting
      ax.autoscale(False)
      P=np.ndarray((2, 2), np.float)

      for i in range(numPoints):
        q=aa[2]
        P[0, :]=M[i, 0, :]
        P[1, :]=M[i, 0, :]+q*AB[i,:]
        ax.plot(P[:, 0], P[:, 1])
      ax.plot(aa[0], aa[1], marker='*',color='r')
      plt.show()
    self._opt_ctr=np.asarray(aa[:2]).ravel()
    _log.debug('optical center:{}'.format(tuple(self._opt_ctr)))


# # line fitting y= ax+b
# # calculate line fitting (least square)
# # y= a*x+b

# # [x1  1]   [a]   [y1]
# # [..  1] * [b] = [..]
# # [xn  1]         [yn]
# #    A    * aa  =  y

# # calculate least square line fitting for each point
# A=np.asmatrix(np.ndarray((numZoom, 2), np.float))
# A[:, 1]=1
# AA=np.ndarray((numPoints, 2), np.float)  # fitting results a1..an b1..bn
# for i in range(numPoints):
#   A[:, 0]=M[i, :, 0].reshape(-1, 1)
#   y=np.asmatrix(M[i, :, 1]).T
#   aa=(A.T*A).I*A.T*y
#   AA[i]=aa.reshape(2)

# if debug:
#   # plot least square line fitting
#   ax.autoscale(False)
#   rngx=(M[:, :, 0].min(), M[:, :, 0].max())
#   Y=np.stack((AA[:, 0]*rngx[0]+AA[:, 1], AA[:, 0]*rngx[1]+AA[:, 1]))
#   P=np.ndarray((2, 2), np.float)
#   P[:, 0]=rngx

#   for i in range(numPoints):
#     P[:, 1]=Y[:, i]
#     ax.plot(P[:, 0], P[:, 1])
#   plt.show()

# # TODO: calculate all cross points of least square line fitting
# # a1*x + b1 = y
# # a2*x + b2 = y

# # [a1  -1] * [x]   [-b1]
# # [a2  -1] * [y] = [-b2]
# P=np.ndarray(((numPoints-1)*numPoints//2, 2), np.float)
# k=0
# D=np.ndarray((2, 2), np.float)
# D[:, 1]=-1
# # for i in range(numPoints):
# #  for j in range(i+1,numPoints):
# #    P[k,:]=
# #    k+=1


  def interp_zoom(self, zoom):
    # this calculates the interpolated pix2pos coordinate transformation matrix
    # the returned value is a 2x2 matrix:
    # [a b] [mx] [px]
    # [   ]*[  ]=[  ]
    # [c d] [my] [py]

    K, AA=self._lut_pix2pos
    pix2pos=self._pix2pos=np.asmatrix(np.ndarray((2, 2), np.float))
    pix2pos[0, 0]=np.interp(zoom, K, AA[:, 0, 0])
    pix2pos[0, 1]=np.interp(zoom, K, AA[:, 0, 1])
    pix2pos[1, 0]=np.interp(zoom, K, AA[:, 1, 0])
    pix2pos[1, 1]=np.interp(zoom, K, AA[:, 1, 1])
    _log.debug('interpolation array:{}'.format(tuple(pix2pos.ravel())))

  def update_pix2pos(self, meas):
    # calculates _lut_pix2pos out of measurements
    # meas is a structure as in the sample code
    # [pxx pxy]
    # [pyx pyy]
    #
    # [a b] [mx] [px]
    # [   ]*[  ]=[  ]
    # [c d] [my] [py]
    #
    # [a*mx1  b*my1]  [px1]
    # [  ...   ... ]  [...]
    # [a*mxn  b*myn] =[pxn]
    # [c*mx1  d*my1]  [py1]
    # [  ...   ... ]  [...]
    # [c*mxn  d*myn]  [pyn]
    #
    # [mx1 my1   0   0 ] [a] [px1]
    # [... ..    0   0 ] [b] [...]
    # [mxn myn   0   0 ]*[c]=[pxn]
    # [ 0   0   mx1 my1] [d] [py1]
    # [ 0   0   ... .. ]     [...]
    # [ 0   0   mxn myn]     [pyn]
    #
    #          A        *aa = y

    # https://de.wikipedia.org/wiki/Methode_der_kleinsten_Quadrate

    # r= A*aa-y
    # mit aa=(A.T*A)^-1*A.T*y

    # A=np.zeros((d.shape[0]*2,d.shape[1]))
    # A[:d.shape[1]-1,:2]=A[d.shape[1]-1:,2:]=d[:,:2]
    #
    # y=d[:,2:].T.ravel()
    # y=np.asmatrix(y).T
    #
    # A=np.asmatrix(A)
    # aa=(A.T*A).I*A.T*y
    # a,b,c,d=aa
    # a,b,c,d=np.asarray(aa).ravel()
    #
    # aa=aa.reshape(2,-1)
    _log.debug('raw data\nmeas:{}'.format(meas))
    if len(meas)<2:
      _log.error('need at least 2 zoom levels:\nmeas:{}'.format(meas))
      return
    AA=np.ndarray((len(meas), 2, 2), np.float)
    K=np.array(tuple(meas.keys()), np.float)
    self._lut_pix2pos=(K, AA)
    for i, (k, v) in enumerate(meas.items()):
      m=np.array(v)  # measurements
      if m.shape[0]<3:
        _log.error('zoom:{}: need at least 3 points per zoom levels:\nmeas:{}'.format(k, v))
        return
      d=m[1:]-m[0]  # distances

      A=np.zeros((d.shape[0]*2, d.shape[1]), np.float)
      A[:d.shape[0], :2]=A[d.shape[0]:, 2:]=d[:, :2]

      y=d[:, 2:].T.ravel()
      y=np.asmatrix(y).T
      A=np.asmatrix(A)

      aa=(A.T*A).I*A.T*y
      AA[i]=aa.reshape(2, -1)
      # bx_coefs = np.polyfit(nbx[0], nbx[1], 3)
      # np.poly1d(bx_coefs)
    _log.debug('least square data:\nK:{}\nAA:{}'.format(K, AA))

  def autofocus(self):
    # cam   camera object
    # mot   motor object
    # rng   region (min max relative to current position) to seek
    # n     number of images to take in region
    # roi   region of interrest to calculate sharpness
    # mode  mode to calculate sharpness (sum/max-min/hist?  of edge detection in roi)
    pass

  def pix2pos(self, p, zoom=None):
    # returns the position m(x,y) in meter relative to the optical center at a given zoom level of the pixel p(x,y)
    # if zoom=None, the last zoom value is used
    if zoom is not None:
      self.interp_zoom(zoom)
    return np.asarray(self._pix2pos*np.mat(p).T).ravel()

  def pos2pix(self, p, zoom=None):
    # returns the pixel p(x,y) of the position m(x,y) in meter relative to the optical center at a given zoom level
    # if zoom=None, the last zoom value is used
    if zoom is not None:
      self.interp_zoom(zoom)
    return np.asarray(self._pix2pos.I*np.mat(p).T).ravel()

  def optctr2xray(self):
    # returns the vector m(x,y) of the optical center to the xray
    pass


if __name__=="__main__":
  import argparse

  logging.basicConfig(level=logging.DEBUG, format='%(levelname)s:%(module)s:%(lineno)d:%(funcName)s:%(message)s ')

  #(h, t)=os.path.split(sys.argv[0]);cmd='\n  '+(t if len(h)>20 else sys.argv[0])+' '
  #exampleCmd=('', '-m0xf -v0')
  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 (see bitmasks) default=0x%(default)x", default=0xff)

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

  logging.getLogger('matplotlib').setLevel(logging.INFO)
  import numpy as np
  import matplotlib as mpl
  import matplotlib.pyplot as plt

  np.set_printoptions(suppress=True)
  np.set_printoptions(linewidth=196)

  obj=geometry()

  if args.mode&0x01:
    measure={
      1:[(-3.0, -7.6, 116.99110193046, 632.5463827525),
         (-2.0, -7.6, 934.07501517856, 600.7926167715),
         (-2.0, -7.1, 916.54131238102, 191.0366615002),
         (-3.0, -7.1, 103.74668003329, 226.2150231456)],
      200:[(-3.1, -7.3, 113.66321353086, 824.9041423107),
           (-2.3, -7.3, 1065.97386092697, 792.2851118419),
           (-2.3, -6.7, 1033.68452410347, 74.0336610693),
           (-3.1, -6.7, 84.62681572700, 116.6832971512)],
      400:[(-3.4, -6.7, 155.00053674203, 601.3838942136),
           (-3.0, -6.7, 957.95919656052, 573.0827012272),
           (-3.2, -6.5, 541.08684037200, 187.9171307943),
           (-3.2, -6.8, 564.32152887203, 789.1146957326)],
      600:[(-3.3, -6.8, 328.27244399903, 509.5061192017),
           (-3.1, -6.8, 992.78996735279, 488.0323963092),
           (-3.2, -6.9, 672.03111567934, 832.4122409755),
           (-3.2, -6.7, 645.70960116180, 164.2534779331)],
      800:[(-3.2, -6.7, 639.52253576449, 53.4455632943),
           (-3.2, -6.85, 671.47023245203, 882.6335091391),
           (-3.3, -6.75, 105.12470026379, 361.3051859197),
           (-3.1, -6.75, 1195.96864609255, 313.1068618673)],
      1000:[(-3.25, -6.75, 195.05641095116, 353.3492286375),
            (-3.15, -6.75, 1117.27204644084, 314.9636405871),
            (-3.2, -6.8, 675.10991143017, 790.3040145281),
            (-3.2, -6.72, 638.98580653116, 59.3803912957)]}

    obj.update_pix2pos(measure)
    obj.interp_zoom(1)
    obj.interp_zoom(200)
    obj.interp_zoom(100)


  if args.mode&0x02:
    opt_ctr_meas={  # x,y = 0.02, -4.89
      1000:[(1057.4251530483375, 116.10122290395591),
            (117.84916300310408, 190.27827474963223),
            (184.2181041281829, 963.2812360887852),
            (1092.5616512910262, 899.514998537239)],
      800:[(888.2494207687248, 203.2917926172947),
           (329.96950424600305, 248.83910515411347),
           (372.9141132092893, 708.2162858826),
           (906.4683457834523, 675.6824912134438)],
      600:[(781.5385742538922, 251.44180872764602),
           (447.09116505496564, 264.4553265953085),
           (471.81684900352445, 554.6567750441825),
           (798.4561474818535, 535.1364982426887)],
      400:[(722.9777438494109, 286.5783069703348),
           (525.1722722609408, 295.68776947769857),
           (535.5830865550707, 462.26079818377866),
           (729.4845027832422, 450.5486321028824)],
      200:[(689.1425973934884, 308.70128734536104),
           (565.5141776506945, 307.39993555859473),
           (574.6236401580583, 406.30267135282986),
           (693.0466527537872, 399.79591241899857)],
      1:[(673.5263759522934, 307.39993555859473),
         (591.5412133860195, 308.70128734536104),
         (595.4452687463182, 376.3715802572061),
         (672.2250241655271, 373.7688766836736)]}

    obj.update_optical_center(opt_ctr_meas, True)