#!/usr/bin/env python
# *-----------------------------------------------------------------------*
# |                                                                       |
# |  Copyright (c) 2024 by Paul Scherrer Institute (http://www.psi.ch)    |
# |                                                                       |
# |              Author xxx (xxx.xxx@psi.ch)                              |
# *-----------------------------------------------------------------------*
'''
slic interface class for CTAAcquisition
based on xxx code snipplets
'''

import logging
import numpy as np
from tqdm import tqdm
from time import time
from time import sleep
import sys, os, threading

sys.path.insert(0, os.path.expanduser('/sf/cristallina/applications/slic/slic-package'))

from slic.core.acquisition import SFAcquisition
from slic.core.acquisition.sfacquisition import BSChannels, transpose_dicts, print_response
from slic.core.task import DAQTask

_log=logging.getLogger(__name__)


class CTAAcquisition(SFAcquisition):

  def __init__(self, cta, *args, n_block_size=None, **kwargs):
    self.cta = cta
    self.n_block_size = n_block_size
    self.wait_time=60.0
    self._abort_evt=threading.Event()
    self._abort_pid=None
    super().__init__(*args, **kwargs)

  def abort(self, stop_pid=None):
    '''stop the CTA sequence and end the data retrieval at stop_pid (None: retrieve nothing more)'''
    self._abort_pid=stop_pid
    self._abort_evt.set()
    self.cta.stop()

  def acquire(self, filename, data_base_dir=None, detectors=None, channels=None, pvs=None, scan_info=None, n_block_size=None, n_pulses=100, n_repeat=1, is_scan_step=False, wait=True, **kwargs):

    if n_repeat != 1:

      raise NotImplementedError("Repetitions are not implemented") #TODO

    if not is_scan_step:

      run_number = self.client.next_run()
      _log.info(f"Advanced run number to {run_number}.")

    else:

      run_number = self.client.run_number
      _log.info(f"Continuing run number {run_number}.")
    #   if not filename or filename == "/dev/null":
    #   print("Skipping retrieval since no filename was given.")
    #   return

    if detectors is None:

      _log.debug("No detectors specified, using default detector list.")
      detectors = self.default_detectors

    if pvs is None:

      _log.debug("No PVs specified, using default PV list.")
      pvs = self.default_pvs

    if channels is None:

      _log.debug("No channels specified, using default channel list.")
      channels = self.default_channels
    n_block_size = n_block_size or self.n_block_size
    #print("block size:", n_block_size)

    bschs = BSChannels(*channels)
    bschs.check()
    client = self.client
    #print(kwargs)
    client.set_config(n_pulses, filename, detectors=detectors, channels=channels, pvs=pvs, scan_info=scan_info, **kwargs)
    #print('set client')
    self._abort_evt.clear()
    self._abort_pid=None
    abort_evt=self._abort_evt

    def _acquire():

      filenames=None
      try:
        if abort_evt.is_set():
          return filenames
        self.cta.stop()
        self.cta.start()
        if abort_evt.is_set(): # abort arrived while starting: make sure the sequence is stopped
          self.cta.stop()
        abort_evt.wait(10) ###? (returns early on abort)
        time_start = time()
        start_pid = self.cta.get_start_pid()
        _log.info(f"CTA start pid: {start_pid}")
        stop_pid = start_pid + n_pulses
        #print("CTA stop pid:", stop_pid)
        pids = np.arange(start_pid, stop_pid)
        pids_blocks = split(pids, n_block_size)
        #print(pids_blocks)
        res = {}
        for pb in tqdm(pids_blocks):
          #print(pb)
          if filename:
            aborted = abort_evt.wait(len(pb) / 100 + 0.2) # wait for the pulses of this block, returns early on abort
            if aborted:
              # the retrieval can lag behind the pulses: continue with the following blocks, each cut at the
              # abort pulse id, until a block is empty
              stop_pid = self._abort_pid
              pb = pb[pb <= stop_pid] if stop_pid is not None else pb[:0]
              _log.warning(f"CTA acquisition aborted at pulse id {stop_pid}: retrieve {len(pb)} pulses of this block")
              if len(pb) == 0:
                break
              sleep(0.2)
            res = self.retrieve(filename, pb, run_number=run_number, scan_info=scan_info, **kwargs) # error here
            res = transpose_dicts(res)
            filenames = res.pop("filenames")
            #print(filenames)

          else:
            res = {}
          #print(res)
        print_response(res)
        # possibly optional, because cta should be done!

        while self.cta.is_running() and not abort_evt.is_set():
          abort_evt.wait(self.wait_time)
          delta_time = time() - time_start
          _log.info(f"Waiting since {delta_time:.0f} seconds for CTA sequence to finish")

        return filenames

      except Exception as e:
        _log.error(f"CTA acquisition: {e}")

    def stopper():
      client.stop()
      self.abort()

    task = DAQTask(_acquire, stopper=stopper, filename=filename, hold=False)
    self.current_task = task
    if wait:
      try:
        task.wait()
      except KeyboardInterrupt:
        _log.warning("Stopped current DAQ task")

    return task

def split(a, block_size):
  if block_size is None:
    return [a]

  length = len(a)
  if isinstance(block_size, (list, tuple, np.ndarray)):
    if sum(block_size) != length:
      raise ValueError(
        f"Sum of block sizes {sum(block_size)} does not match PID length {length}."
      )

    blocks = []
    start = 0
    for size in block_size:
      if size <= 0:
        raise ValueError(f"Block size values must be positive, got {size}")
      blocks.append(a[start:start + size])
      start += size
    return blocks

  indices = np.arange(block_size, length, block_size) # must not start at 0, otherwise the first entry is an empty array

  return np.array_split(a, indices)

if __name__ == "__main__":
  from slic.devices.timing.events import CTASequencer

  cta = CTASequencer("SAT-CCTA-ESC")
  daq = CTAAcquisition(cta, "cristallina", "p19150", default_channels=["SAR-CVME-TIFALL6:EvtSet"], append_user_tag_to_data_dir=True)

  cta.cfg.repetitions = n_pulses # etc. etc.
  daq.acquire("test")