#!/usr/bin/env python
# *-----------------------------------------------------------------------*
# |                                                                       |
# |  Copyright (c) 2022 by Paul Scherrer Institute (http://www.psi.ch)    |
# |                                                                       |
# |              Author Thierry Zamofing (thierry.zamofing@psi.ch)        |
# *-----------------------------------------------------------------------*
'''
Module to handle FixTarget acquisition

This contains a Widget to handle FixTargetFrames and fiducials, calculate final target positions etc.
'''

#https://stackoverflow.com/questions/27909658/json-encoder-and-decoder-for-complex-numpy-arrays

import logging
_log=logging.getLogger(__name__)

import os, json, base64, yaml
import numpy as np
import pyqtUsrObj as UsrGO
import pyqtgraph as pg
from PyQt5.QtCore import Qt, QFileInfo, pyqtSignal
from PyQt5.QtWidgets import (
  QFileDialog, QWidget, QGridLayout, QItemDelegate, QVBoxLayout, QPushButton, QApplication,
  QMainWindow, QLineEdit, QMessageBox, QAction, QComboBox)


class MyJsonEncoder(json.JSONEncoder):
  """ Special json encoder for numpy types """
  def default(self, obj):
    if isinstance(obj, np.integer):
      return int(obj)
    elif isinstance(obj, np.floating):
      return float(obj)
    elif isinstance(obj, np.ndarray):
        data_b64=base64.b64encode(obj.data)
        return dict(__ndarray__=data_b64,dtype=str(obj.dtype),shape=obj.shape)
        #return obj.tolist()
    elif isinstance(obj, set):
      return list(obj)
    elif type(obj) not in (dict,list,str,int):
      try:
        return obj.obj2json(self)
      except AttributeError:
        _log.error('dont know how to json')
        return repr(obj)
    return json.JSONEncoder.default(self, obj)

  def iterencode(self, o, _one_shot=False):
    list_lvl = 0
    l=super().iterencode(o, _one_shot=_one_shot)
    #l=tuple(l);print(''.join(l)) # helpful to debug
    for s in l:
      if s.startswith('['):
        list_lvl += 1
      if list_lvl > 0:
        s = s[0]+s[1:].replace('\n', '').strip()
        s = s.replace('\n', '').rstrip()
        #self.item_separator):
        #self.key_separator
      if s.endswith(']'):
        list_lvl -= 1
      yield s

def MyJsonDecoder(dct):
    if isinstance(dct, dict):
      if '__class__' in dct:
        cls=dct.pop('__class__')
        cls=UsrGO.__dict__[cls]
        obj=cls.__new__(cls)
        obj.__init__(**dct)
        #try:
        #  obj.json2obj(dct)
        #except AttributeError:
        #  obj.__init__(**dct)
        return obj
      elif '__ndarray__' in dct:
        data = base64.b64decode(dct['__ndarray__'])
        return np.frombuffer(data[:-1], dct['dtype']).reshape(dct['shape'])
    return dct


class MarkerDelegate(QItemDelegate):
  def createEditor(self, parent, option, index):
    comboBox=QLineEdit(parent)
    comboBox.editingFinished.connect(self.emitCommitData)
    return comboBox

  def setEditorData(self, editor, index):
    editor.setText("{}".format(index.model().data(index)))

  def setModelData(self, editor, model, index):
    print("model {}".format(model))
    print(
      "delegate setData: {}x{} => {}".format(
        index.row(), index.column(), editor.text()
      )
    )
    model.setData(index, editor.text())
    model.setData(index, float(editor.text()), Qt.UserRole)

  def emitCommitData(self):
    print("imagedelegate emitting")
    self.commitData.emit(self.sender())


class WndFixTarget(QWidget):
  prefixSelected = pyqtSignal(str)
  dataFileLoaded = pyqtSignal(str)
  prelocatedDataUpdated = pyqtSignal()
  markersDeleted = pyqtSignal()
  markerAdded = pyqtSignal(bool, list)  # emits is_fiducial(boolean), [x, y, cx, cy]
  selectedRowChanged = pyqtSignal(int)
  moveFastStageRequest = pyqtSignal(float, float)

  def __init__(self, parent=None,data=None):
    super(WndFixTarget, self).__init__(parent)
    self._xtals_transformed = True

    self._current_row = None

    layout = QVBoxLayout()

    self.setLayout(layout)
    frame = QWidget()
    bl = QGridLayout()
    frame.setLayout(bl)

    btnLd = QPushButton("Load Datafile")
    btnLd.clicked.connect(lambda: self.load_file(None))
    bl.addWidget(btnLd, 0, 0)

    btnSv = QPushButton("Save Datafile")
    btnSv.clicked.connect(lambda: self.save_file())
    bl.addWidget(btnSv, 0, 1)

    btnDump = QPushButton("dump")
    btnDump.clicked.connect(self.dump_data)
    bl.addWidget(btnDump, 0, 2)

    self._cbType=cb=QComboBox()

    bl.addWidget(cb, 1,0,1,1)
    self._txtParam=param=QLineEdit()
    bl.addWidget(param, 1, 1,1,1)

    self._btnAdd=btnAdd = QPushButton("add obj")
    #btnAdd.clicked.connect(lambda x: _log.warning("TODO: IMPLEMENT") )
    bl.addWidget(btnAdd, 1, 2,1,1)

    self._btnDelAll=btnDelAll = QPushButton("del all")
    #btnDelAll.clicked.connect(lambda x: _log.warning("TODO: IMPLEMENT") )
    bl.addWidget(btnDelAll, 2, 2,1,1)
    self._btnFit=btnFit = QPushButton("fit with fiducials")
    #btnDelAll.clicked.connect(lambda x: _log.warning("TODO: IMPLEMENT") )
    bl.addWidget(btnFit, 2, 1,1,1)

    layout.addWidget(frame)

    self._data=data
    self._tree=tree=pg.DataTreeWidget(data=data)
    layout.addWidget(tree, stretch=2)
    #Context menu: by default the function
    # def contextMenuEvent(self, event) is called... but this requests to overload the function
    # and to build an own context menu
    # much simpler is to use Qt.ActionsContextMenu
    # and to handle the context in the parent class

    tree.setContextMenuPolicy(Qt.ActionsContextMenu)

    act = QAction("center in view", self)
    act.triggered.connect(self.tree_ctx_center)
    tree.addAction(act)

    act = QAction("delete", self)
    act.triggered.connect(self.tree_ctx_delete)
    tree.addAction(act)

    act = QAction("update param", self)
    act.triggered.connect(self.tree_ctx_update)
    tree.addAction(act)

  def get_param(self):
    param=self._txtParam.text().replace('(','[').replace(')',']').strip()
    if param=='' or param[0]!='{':
      param='{'+param+'}'
    #mft._cbType.addItems(["Fiducial", "FixTarget(12.5x12.5)", "FixTarget(23.0x23.0)", "FixTarget(<param>)", "Grid(<param>)", "SwissMX-path"])
    #bm_pos_eu=self._goBeamMarker._pos_eu
    #bm_size_eu=self._goBeamMarker._size_eu
    try:
      #parse the parameters: as yaml string.
      # allows : without space, allows () as []
      # no {} to define a dictionary
      # e.g. 'a:ggf,b:5,c:[5,6.1],d(8,9,3)'

      param=param.replace(':', ': ')  # allow gen:4 without space
      param=yaml.safe_load(param)  # "ofs":[10, 5],"width":200,"fidScl":0.5,"fiducial":[[18,7],[25,16],[70, 20]]
    except BaseException as e:
      _log.error(f'{e}:{param}')
      param=dict()
    return param

  def tree_get_path(self):
    path=[]
    it=self._tree.currentItem()
    while it:
      d=it.data(0,0)
      try:
        d=int(d)
      except ValueError:
        pass
      if d!= '':
        path.append(d)
      it=it.parent()
    return path

  def tree_ctx_delete(self):
    app=QApplication.instance()
    path=self.tree_get_path()
    if len(path)==1:
      try:
        wnd=app._mainWnd
      except AttributeError:
        _log.info('_mainWnd not handeled')
      else:
        vb=wnd.vb
        grp=wnd._goTracked
        go=grp.childItems()[path[0]]
        vb.removeItem(go)
        data=grp.childItems()
        if not len(data):
          grp.setFlag(grp.ItemHasNoContents)
        self._tree.setData(data)

  def tree_ctx_center(self):
    app=QApplication.instance()
    path=self.tree_get_path()
    if len(path)==1:
      try:
        wnd=app._mainWnd
      except AttributeError:
        _log.info('_mainWnd not handeled')
      else:
        vb=wnd.vb
        grp=wnd._goTracked
        go=grp.childItems()[path[0]]
        #vb.autoRange(items=(go,))
        r1=vb.viewRect()
        r2=vb.itemBoundingRect(go)
        #r=vb.viewRect()
        vb.setRange(rect=r2)
        r1.translate(r2.center()-r1.center())
        vb.setRange(r1)

  def tree_ctx_update(self):
    app=QApplication.instance()
    path=self.tree_get_path()
    if len(path)==1:
      try:
        wnd=app._mainWnd
      except AttributeError:
        _log.info('_mainWnd not handeled')
      else:
        grp=wnd._goTracked
        go=grp.childItems()[path[0]]
        go._param=self.get_param()
        data=grp.childItems()
        self._tree.setData(data)

  def load_file(self, filename=None):
    app = QApplication.instance()
    if filename is None:
      filename, _ = QFileDialog.getOpenFileName(self,"Load data file",None,
        "json files (*.json);;all files (*)",)
        #"json files (*.json);;yaml files (*.yaml);;pickle files (*.pkl);;text files (*.txt);;all files (*)",)
    if not filename:  # cancelled dialog
      return

    ext=filename.rsplit('.',1)[1].lower()

    if ext=='json':
      with open(filename, 'r') as f:
        data=json.load(f,object_hook=MyJsonDecoder)
    else:
      raise(IOError('unsupported file type'))


    self._tree.setData(data)
    try:
      wnd=app._mainWnd
    except AttributeError:
      _log.info('_mainWnd not handeled')
      pass
    else:
      grp=wnd._goTracked
      for go in data:
        grp.addItem(go)
        if type(go)==UsrGO.Fiducial:
          go.sigRegionChangeFinished.connect(wnd.cb_fiducial_update_z)
      data=grp.childItems()
      self._tree.setData(data)
      #wnd._goTracked['objLst']=self._data
    return


    try:
      cfg.setValue("folders/last_prelocation_folder", QFileInfo(filename).canonicalPath())
      cfg.setValue("folders/last_prelocation_sheet", QFileInfo(filename).absolutePath())
    except TypeError as e:
      pass

    self.clearMarkers()

    _log.info(f"loading prelocated coords from {filename}")
    df = pd.read_csv(filename, comment="#", header=None, delim_whitespace=False, delimiter="[\t,;]")

    try:
      prefix = QFileInfo(filename).baseName()
    except:
      prefix = filename.name
      filename = prefix

    logger.info(f"prefix => {prefix}")

    gonio_coords_available = True
    for data in df.as_matrix(): # FIXME FutureWarning: Method .as_matrix will be removed in a future version. Use .values instead.
      row = self.markersTable.rowCount()
      self.markersTable.setRowCount(row + 1)
      original = np.copy(data)
      data = list(data)

      if len(data) in [2, 3]:
        gonio_coords_available = False
        data.extend([0, 0])

      if len(data) == 4:
        data.insert(0, False)  # fiducial flag
      elif len(data) == 5:  # fiducial already there, convert to bool
        data[0] = bool(data[0])
      else:
        QMessageBox.warning(
          self,
          "Wrong number of points in data file",
          "I was expecting either 2, 3, 4 or 5 data points per line."
          "\n\nFailed around a line with: {}".format(list(original)),
        )

      self.addSingleMarker(row, data)

    self._xtals_transformed = False
    self.prefixSelected.emit(prefix)

    # only emit this signal if goniometer coordinates already read from file
    if gonio_coords_available:
      logger.debug(f"dataFileLoaded.emit => {filename}")
      self.dataFileLoaded.emit(filename)

  def save_file(self, filename=None):
    app = QApplication.instance()
    # filename = folders.get_file("prelocated-save.dat")
    #data_folder = settings.value("folders/last_prelocation_folder")
    data_folder=''
    if filename is None:
      filename, _ = QFileDialog.getSaveFileName(self,"Save data file",data_folder,
        "json files (*.json);;all files (*)",)
        #"json files (*.json);;yaml files (*.yaml);;pickle files (*.pkl);;text files (*.txt);;all files (*)",)
    if not filename:
      return

    #settings.setValue("folders/last_prelocation_folder", QFileInfo(filename).canonicalPath())
    #_log.info("Saving data in {}".format(filename))
    #data = self.get_data(as_numpy=True)
    #df = pd.DataFrame(data)
    #df.to_csv(filename, float_format="%.6f")
    #import numpy as np
    base,ext=os.path.splitext(filename)
    if not ext.lower():
      ext='json'
      filename=base+'.'+ext
    try:
      wnd=app._mainWnd
    except AttributeError:
      _log.info('_mainWnd not handeled')
      data=self._data
    else:
      grp=wnd._goTracked
      data=grp.childItems()

    if ext.lower()=='json':
      with open(filename, 'w') as f:
        json.dump(data, f,cls=MyJsonEncoder, indent=2)#separators=(',', ':')
    else:
      raise(IOError('unsupported file type'))
    #elif  ext=='yaml':
    #  with open(filename, 'w') as f:
    #    yaml.dump(self._data, f)
    #elif   ext=='pkl':
    #  with open(filename, 'wb') as f:
    #    pickle.dump(self._data, f)
    #print(self._data)


  def dump_data(self):
    print(self._data)
    #for ref, x, y, cx, cy, ox, oy in self.get_data():
    #  print(f"fiducial:{ref} [{x}, {y}, {cx}, {cy}]")

  # **********3 OBSOLETE

  def delete_selected(self):
    row = self._current_row
    try:
      row += 0
    except:
      _log.warning("select a row first")
      QMessageBox.warning(self,"Select a marker first","You must select a marker first before updating its goniometer position",)
      return
    self.markersTable.removeRow(row)
    self.prelocatedDataUpdated.emit()

  def selection_changed(self):
    row = self.markersTable.currentRow()
    if row < 0:
      return
    self._current_row = row
    self.selectedRowChanged.emit(row)
    _log.debug("selection changed: current row {}".format(row))

if __name__ == "__main__":
  import sys

  class Monster:
    yaml_tag = u'!Monster'
    def __init__(self, name, hp, ac, attacks):
      self.name = name
      self.hp = hp
      self.ac = ac
      self.attacks = attacks
    def __repr__(self):
      return "%s(name=%r, hp=%r, ac=%r, attacks=%r)" % (
        self.__class__.__name__, self.name, self.hp, self.ac, self.attacks)


  class MainWindow(QMainWindow):
    def __init__(self, parent=None):
      super(MainWindow, self).__init__(parent)
      self.move(100,100)
      self.resize(500,800)
      self.centralWidget=QWidget()
      self.wndFixTrg=WndFixTarget(parent=self)

      self.setCentralWidget(self.centralWidget)
      mainLayout=QVBoxLayout()
      mainLayout.addWidget(self.wndFixTrg)
      self.centralWidget.setLayout(mainLayout)


  app = QApplication(sys.argv)
  mainWin = MainWindow()

  t=3
  if t==0:
    data=[{'name':'John Doe', 'occupation':'gardener'}, {'name':'Lucy Black', 'occupation':'teacher'}]
  elif t==1:
    data ={
     'name': 'John Doe', 'occupation': 'gardener',
     'A':(1,2,3),
     'B':[1,2,3],
     'C':{1,2,3},
     'D':
       {'1':(1,2,3),
        '2':'df',
        '3':'dddd'},
     '4':np.array(((1,2,3),(4,5,6))),
     '5':np.array(range(40*2)).reshape(-1,2)
     }
  elif t==2:
    data = {
      'a list': [1,2,3,4,5,6, {'nested1': 'aaaaa', 'nested2': 'bbbbb'}, "seven"],
      'a dict': {
          'x': 1,
          'y': 2,
          'z': 'three'
      },
      'an array': np.random.randint(10, size=(40,10)),
      #'a traceback': some_func1(),
      #'a function': some_func1,
      #'a class': pg.DataTreeWidget,
    }
  elif t==3:
    data=[
      UsrGO.Grid(pos=(120.0, -100.0), size=(1000.0, 500.0), cnt=(10, 5), fiducialSize=2),
      Monster(name='Cave lizard', hp=[3,6], ac=16, attacks=['BITE','HURT'])
    ]
  mainWin.wndFixTrg._data=data
  mainWin.wndFixTrg._tree.setData(data)
  mainWin.show()
  sys.exit(app.exec_())