#!/usr/bin/env python
# *-----------------------------------------------------------------------*
# |                                                                       |
# |  Copyright (c) 2022 by Paul Scherrer Institute (http://www.psi.ch)    |
# |                                                                       |
# |              Author Thierry Zamofing (thierry.zamofing@psi.ch)        |
# *-----------------------------------------------------------------------*
import os,epics
import time

import numpy as np

class SimCam:
  def __init__(self):
    pass

  def generate(self,sz=(1500,1000),t=100,mode=0):
    w,h=sz
    self._imgSeq=imgSeq=np.ndarray(shape=(t,h,w),dtype=np.uint16)
    x = np.linspace(-5, 5, w)
    y = np.linspace(-5, 5, h)
    # full coordinate arrays
    xx, yy = np.meshgrid(x, y)

    for i in range(t):
      #imgSeq[i,:,:] = 100*np.sqrt(np.sin(xx+.1*i)**2 + np.sin(yy+.01*i)**2)#+xx*t+yy*t)
      #imgSeq[i,:,:] = 100*np.sqrt(np.sin(xx+.1*i)**2 + np.sin((1+.1*np.sin(.2*i))*yy+.001*i**2)**2)#+xx*t+yy*t)
      #imgSeq[i,:,:] = 100*np.sqrt(np.sin(xx+2*np.sin(i/t*2*np.pi))**2 + np.sin(yy)**2)
      px=2*np.sin(i/t*2*np.pi)
      fx=1
      py=2*np.sin(i/t*2*np.pi)
      fy=1+.3*np.sin(i/t*2*np.pi*2)
      imgSeq[i,:,:] = 100*np.sqrt(np.sin(xx*fx+px)**2 + np.sin(yy*fy+py)**2)
      #np.random.bytes(100)
    wr=w//4
    hr=h//4
    imgSeq[:,0:hr,0:wr]+=np.random.randint(0,100,(t,hr,wr),dtype=np.uint16)



  def pushEPICS(self,sz=None,t=100,prefix='SwissMxSim:'):
    if prefix[-1]!=':': prefix+=':'
    os.environ['EPICS_CA_ADDR_LIST'] = 'localhost'
    pv = dict()
    pv['pic'] = epics.PV(prefix + "FPICTURE")
    for k, v in (('w', 'WIDTH'), ('h', 'HEIGHT'),):
      pv[k] = epics.PV(prefix + v)
    if sz:
      w,h=sz
      pv['w'].put(w)
      pv['h'].put(h)
    else:
      w, h = int(pv['w'].value), int(pv['h'].value)

    print('generate data...')
    self.generate((w,h),t)
    print('done. Now cycle and publish the data to EPICS')
    imgSeq=self._imgSeq
    while(True):
      for i in range(t):
        a = imgSeq[i,:,:].ravel()
        pv['pic'].put(a)
        print(pv['pic'].get())
        #time.sleep(.01)

if __name__ == "__main__":

  def QtUI(simCam):
    global img, i, fps, updateTime
    from pyqtgraph.Qt import QtCore, QtGui
    import numpy as np
    import pyqtgraph as pg
    import pyqtgraph.ptime as ptime
    import sys

    app = QtGui.QApplication([])
    ## Create window with GraphicsView widget
    win = pg.GraphicsLayoutWidget()
    win.show()  ## show widget alone in its own window
    win.setWindowTitle('pyqtgraph example: ImageItem')
    view = win.addViewBox()

    ## lock the aspect ratio so pixels are always square
    view.setAspectLocked(True)

    ## Create image item
    img = pg.ImageItem(border='g')
    view.addItem(img)

    ## Set initial view bounds
    view.setRange(QtCore.QRectF(0, 0, 600, 600))

    i=0; fps=0; updateTime = ptime.time()
    def updateData():
      global img, i, updateTime, fps
      img.setImage(simCam._imgSeq[i])
      i = (i + 1) % simCam._imgSeq.shape[0]
      QtCore.QTimer.singleShot(1, updateData)
      now = ptime.time()
      fps2 = 1.0 / (now - updateTime)
      updateTime = now
      fps = fps * 0.9 + fps2 * 0.1
      print("%d %0.1f fps" % (i,fps))
    updateData()

    ## Start Qt event loop unless running in interactive mode.
    if (sys.flags.interactive != 1) or not hasattr(QtCore, 'PYQT_VERSION'):
      QtGui.QApplication.instance().exec_()

  def QtUI2(simCam):
    global img, i, fps, updateTime
    import sys
    from pyqtgraph.Qt import QtCore, QtGui
    import pyqtgraph as pg
    import pyqtgraph.ptime as ptime
    print(pg.__version__)

    # Interpret image data as row-major instead of col-major
    pg.setConfigOptions(imageAxisOrder='row-major')

    app = QtGui.QApplication([])

    ## Create window with ImageView widget
    win = QtGui.QMainWindow()
    win.resize(800,800)
    imv = pg.ImageView()
    win.setCentralWidget(imv)
    win.show()
    win.setWindowTitle('pyqtgraph example: ImageView')

    ## Display the data and assign each frame a time value from 1.0 to 3.0
    i=0; fps=0; updateTime = ptime.time()
    imv.setImage(simCam._imgSeq[i])
    def updateData():
      global img, i, updateTime, fps
      imv.setImage(simCam._imgSeq[i], autoRange=False, autoLevels=False)
      i = (i + 1) % simCam._imgSeq.shape[0]
      QtCore.QTimer.singleShot(1, updateData)
      now = ptime.time()
      fps2 = 1.0 / (now - updateTime)
      updateTime = now
      fps = fps * 0.9 + fps2 * 0.1
      print("%d %0.1f fps" % (i,fps))
    updateData()

    ## Set a custom color map
    colors = [ (0, 0, 0), (45, 5, 61), (84, 42, 55), (150, 87, 60), (208, 171, 141), (255, 255, 255) ]
    cmap = pg.ColorMap(pos=np.linspace(0.0, 1.0, 6), color=colors)
    imv.setColorMap(cmap)

    ## Start Qt event loop unless running in interactive mode.
    if (sys.flags.interactive != 1) or not hasattr(QtCore, 'PYQT_VERSION'):
        QtGui.QApplication.instance().exec_()



  import argparse
  parser = argparse.ArgumentParser()
  parser.add_argument("--ui", "-u", help="qt local speed test", type=int, default=0)
  parser.add_argument("--prefix","-p",help="EPICS camera-PV prefix",type=str,default="SwissMxSim",)
  args = parser.parse_args()

  if args.ui:
    simCam=SimCam()
    print('generate data...')
    simCam.generate()
    print('done')
    if args.ui==1:
      QtUI(simCam)
    else:
      QtUI2(simCam)
  else:
    obj = SimCam()
    obj.pushEPICS((1500,1000),100,args.prefix)
# -*- coding: utf-8 -*-
"""
Demonstrates very basic use of ImageItem to display image data inside a ViewBox.
"""

## Add path to library (just for examples; you do not need this)