import os
from multiprocessing import Pool, Process, Queue, Array
import psutil
import matplotlib.pyplot as plt
from matplotlib.figure import Figure
from matplotlib.backends.backend_agg import FigureCanvasAgg
from matplotlib.colors import Normalize
import numpy as np
import ctypes
import time
import timeit
import logging
from argparse import Namespace
import traceback
import asyncio
import fabio
import ipywidgets as ipw
from pynx.utils.array import rebin, pad2
from pynx.utils.math import primes, test_smaller_primes
from pynx.holotomo.utils import load_data, get_params_id16b
from pynx.holotomo.runner import prep_holotomodata

try:
    # Get the real number of processor cores available
    # os.sched_getaffinity is only available on some *nix platforms
    nproc = len(os.sched_getaffinity(0)) * psutil.cpu_count(logical=False) // psutil.cpu_count(logical=True)
except AttributeError:
    nproc = os.cpu_count()

# This is used to pass shared arrays and their shape between parallel process
shared_arrays = {}


# TODO: need to put this somewhere else
class NoEventContextManager:
    """Context manager to disable events in widgets"""

    def __init__(self, obj):
        """

        :param obj: the widget for which the events will be disabled
        """
        self.obj = obj

    def __enter__(self):
        self.obj._enable_events = False

    def __exit__(self, exc_type, exc_value, exc_tb):
        self.obj._enable_events = True


def data_watcher(qin, qout, prefix_dir, prefix_frame, nb_frame, next_scan, logger):
    """
    :param prefix_dir: e.g. '/path/to/data_{scan:02d}_1' where scan is the scan number (could be absent from pattern)
    :param prefix_frame: e.g. 'data_{scan:02d}_{i:04d}_1.edf', where scan is the scan number, and i the frame index
    :param nb_frame: number of expected frames per scan
    :param next_scan: next scan to watch for
    """
    ct = 0
    while True:
        try:
            next_dir = prefix_dir.format(scan=next_scan)
            next_last_file = os.path.join(next_dir, prefix_frame.format(scan=next_scan, i=nb_frame - 1))
            # logger.info(f"{next_last_file} ? {os.path.exists(next_dir)}")
            if os.path.exists(next_dir):
                # Watch for the last file
                if os.path.exists(next_last_file):
                    logger.info(f"Found new scan #{next_scan}")
                    qout.put(next_scan)
                    next_scan += 1
                    while qin.empty():
                        time.sleep(0.5)
            time.sleep(0.5)
            ct += 1
            # print(ct, os.path.exists(next_dir), os.path.exists(next_last_file), next_last_file)
            stop = False
            while not qin.empty():
                if qin.get() == 'STOP':
                    logger.info("Stopping data watcher")
                    stop = True
            if stop:
                break
        except Exception as ex:
            logger.error(traceback.format_exc())
            logger.error("Error in data watcher")
            raise

    logger.debug("Stopped data watcher")


def load1(kw):
    i, scan, par = kw['i'], kw['scan'], kw['par']
    ref, dark = kw['ref'], kw['dark']
    iobs_array, iobs_shape = shared_arrays['iobs_array']
    iobs = np.frombuffer(iobs_array.get_obj(), dtype=ctypes.c_float).reshape(iobs_shape)

    pth = os.path.join(par.prefix_dir, par.prefix_frame).format(scan=scan, i=i * par.binning)
    tmp = load_data([pth], dark, [1], binning=par.binning, ref=ref, distorted=0,
                    normalise=par.normalise, timing=False)
    if par.padx == 0 and par.pady == 0:
        iobs[i] = tmp
    else:
        iobs[i] = np.expand_dims(pad2(tmp[0], (par.pady, par.padx), mode='reflect_linear'), 0)


def load_scan(scan, logger, par, ref, dark):
    t0 = timeit.default_timer()
    iobs_array, iobs_shape = shared_arrays['iobs_array']
    nb_frame, junk, ny, nx = iobs_shape
    binning = par.binning
    with iobs_array.get_lock():
        with Pool(nproc) as pool:
            res = pool.map(load1, [{'i': i, 'scan': scan, 'par': par,
                                    'ref': ref, 'dark': dark} for i in range(nb_frame)])  # , chunksize=1
        # Double flat-field
        iobs = np.frombuffer(iobs_array.get_obj(), dtype=ctypes.c_float).reshape(iobs_shape)
        if par.doubleff:
            tmp = iobs.mean(axis=0)
            iobs /= tmp / tmp.mean()
        dt = timeit.default_timer() - t0
        logger.info(
            f"Loaded scan #{scan} in dt={dt:.2f}s {ny * nx * nb_frame * binning ** 2 * 16 / 1024 ** 3 / dt:.2f} GB/s"
            f"[first frame:{par.prefix_frame.format(scan=scan, i=0)}], <iobs>={iobs.mean():.2f}")


def data_loader(qin, qout, qin_watcher, qout_watcher, logger, par, ref, dark):
    logger.debug("Started data loader")
    while True:
        try:
            if not qout_watcher.empty():
                scan = qout_watcher.get()
                load_nok = True
                nb_try = 0
                while load_nok and nb_try < 20:
                    try:
                        load_scan(scan, logger, par, ref, dark)
                        load_nok = False
                    except FileNotFoundError:
                        nb_try += 1
                        logger.error(f"Error loading scan {scan} - missing files ? "
                                     f"Wait 1s and try again [try #f{nb_try}/20]")
                        time.sleep(1)
                qin_watcher.put('Find me another scan !')
                if nb_try == 20:
                    logger.error(f"Give up loading scan {scan} - missing files ? ")
                else:
                    qout.put(f"Scan loaded: {scan}")
                    # Wait to continue so algo can grab array
                    while qin.empty():
                        time.sleep(0.1)
            stop = False
            while not qin.empty():
                # Will also read the 'OK' from the ht/algo worker
                if qin.get() == 'STOP':
                    logger.info("Stopping data loader")
                    stop = True
            if stop:
                break
            time.sleep(0.1)
        except Exception as ex:
            logger.error(traceback.format_exc())
            logger.error("Error in data watcher")
            raise
    logger.info("Stopped data loader")


def run_ht(qin, qout, qin_loader, qout_loader, logger, par, ref):
    # Will run in a separate process
    try:
        logger.info(f"Starting algo process")
        # Import pycuda, init GPU
        # import pycuda.autoinit  # or autoprimaryctx ?
        from pynx.holotomo.operator import ScaleObjProbe, AP, BackPropagatePaganin, \
            BackPropagateCTF, ShowObj, SinoFilter
        from pynx.holotomo import HoloTomo, HoloTomoData
        pu = ScaleObjProbe().processing_unit

        ph_array, ph_shape = shared_arrays['ph_array']
        ny0, nx0 = ph_shape[1:]  # without padding
        iobs_array, iobs_shape = shared_arrays['iobs_array']
        ny, nx = iobs_shape[-2:]
        nb_frame = par.nb_frame

        ph = np.frombuffer(ph_array.get_obj(), dtype=ctypes.c_float).reshape(ph_shape)

        logger.info(f"Started algo process")
        # Wait for initial data
        while qout_loader.empty():
            logger.debug(f"Algo: waiting for next scan...")
            time.sleep(1)
        scan = int(qout_loader.get().split(':')[-1])
        logger.info(f"Algo: got first data, scan #{scan}")

        # Create HolotomoData - use only 1 stack
        with iobs_array.get_lock():
            if par.padx > 0 or par.pady > 0:
                # Ref has not yet been padded
                ref1 = np.expand_dims(pad2(ref[0], (par.pady, par.padx), mode='reflect_linear'), 0)
            else:
                ref1 = ref
            data = prep_holotomodata(iobs_array=iobs_array, iobs_shape=iobs_shape,
                                     ref=ref1,
                                     pixel_size=par.pixel_size * par.binning,
                                     wavelength=par.wavelength,
                                     detector_distance=par.distance,
                                     dx=None, dy=None, proj_idx=None,
                                     padding=(par.pady, par.padx), algorithm="Paganin",
                                     gpu_mem_f=1, pu=pu, logger=logger, stack_size=nb_frame)
        qin_loader.put('OK')  # Allows the next scan to be loaded

        # Create Holotomo object
        probe = np.ones((1, 1, ny, nx), dtype=np.complex64)
        logger.info("Creating Holotomo object")
        p = HoloTomo(data=data, obj=None, probe=probe, pu=pu)
        if not par.normalise:
            logger.info("Adding a direct beam constraint for the probe")
        p.constrain_probe_direct_beam = not par.normalise
        update_probe = not par.normalise

        obj_inertia = 0.1
        probe_inertia = 0.01
        verbose = 10

        # Start optimisation
        logger.info("First optimisation")
        pag = BackPropagatePaganin(delta_beta=par.delta_beta)
        apn = AP(update_object=False, update_probe=True,
                 calc_llk=verbose, delta_beta=par.delta_beta, obj_min=None, obj_max=None,
                 obj_smooth=par.obj_smooth, obj_inertia=obj_inertia, weight_empty=1,
                 show_obj_probe=0, probe_inertia=probe_inertia)

        ap = AP(update_object=True, update_probe=update_probe,
                calc_llk=verbose, delta_beta=par.delta_beta, obj_min=None, obj_max=None,
                obj_smooth=par.obj_smooth, obj_inertia=obj_inertia, weight_empty=1,
                show_obj_probe=0, probe_inertia=probe_inertia)
        ap0 = AP(update_object=True, update_probe=update_probe,
                 calc_llk=verbose, delta_beta=0, obj_min=None, obj_max=None,
                 obj_smooth=par.obj_smooth, obj_inertia=obj_inertia, weight_empty=1,
                 show_obj_probe=0, probe_inertia=probe_inertia)
        algo = ap0 ** 2 * ap ** 8 * apn ** 2 * pag if update_probe else ap0 ** 2 * ap ** 8 * pag
        algo = SinoFilter('ram-lak') * algo

        while True:
            # Algorithms
            p = algo * p

            # Update phase array
            logger.info(f"Getting phase array for scan: {scan}")
            with ph_array.get_lock():
                idx, tmp = p.get_obj_phase_unwrapped(crop_padding=True, dtype=np.float32, sino_filter=None)
                ph.flat = tmp
                qout.put(scan)
            while qout_loader.empty() and qin.empty():
                time.sleep(0.1)
            if not qout_loader.empty():
                scan = int(qout_loader.get().split(':')[-1])
                logger.info(f"Algo: got new scan: {scan}")
                # Update Iobs & do 2 cycles
                with iobs_array.get_lock():
                    iobs = np.frombuffer(iobs_array.get_obj(), dtype=ctypes.c_float).reshape(iobs_shape)
                    data.replace_data(iobs, ref1, dx=None, dy=None, idx=None, init_obj=False, pu=None)
                p._cu_timestamp_counter = 0  # Force uploading new data to GPU
                qin_loader.put('OK')  # Allows the next scan to be loaded

            stop = False
            while not qin.empty():
                if qin.get() == 'STOP':
                    logger.info("Stopping algorithms worker")
                    stop = True
            if stop:
                break
    except Exception as ex:
        logger.error(traceback.format_exc())
        logger.error("Error in pynx.holotomo algorithms")
        raise

    # Loop - update data or run algorithms according to instructions from queue
    logger.debug("Finished algorithms")


# FBP process
def run_fbp(qin, qout, qin_algo, qout_algo, logger, par):
    try:
        # Will run in separate process
        from nabu.reconstruction.fbp import Backprojector
        from nabu.reconstruction.sinogram import SinoBuilder
        from pynx.processing_unit.cu_resources import cu_resources
        from pynx.processing_unit import default_processing_unit as pu

        ph_array, ph_shape = shared_arrays['ph_array']
        nb_frame, ny, nx = ph_shape
        vol_array, vol_shape = shared_arrays['vol_array']

        pu.select_gpu(language='cuda')
        cu_ctx = cu_resources.get_context(pu.cu_device)
        B = Backprojector((nb_frame, nx), rot_center=par.tomo_rot_center / par.binning,
                          angles=np.linspace(0, 2 * np.pi, nb_frame),
                          filter_name="ram-lak", backend_options={'ctx': cu_ctx})
        ph = np.frombuffer(ph_array.get_obj(), dtype=np.float32).reshape(ph_shape)
        vol = np.frombuffer(vol_array.get_obj(), dtype=np.float32).reshape(vol_shape)
        while True:
            while qout_algo.empty() and qin.empty():
                time.sleep(0.1)
            if not qout_algo.empty():
                scan = qout_algo.get()
                logger.info(f"3D FBP for scan {scan}")
                with ph_array.get_lock():
                    with vol_array.get_lock():
                        t0 = timeit.default_timer()
                        for i in range(ny):
                            if False:
                                vol[i] = B.fbp(ph[:, i, :].astype(np.float32))
                            else:
                                # Sino filter has already been applied
                                vol[i] = B.backproj(ph[:, i, :] * np.float32(np.pi / nb_frame))

                dt = timeit.default_timer() - t0
                logger.info(f"Finished volume #{scan} reconstruction, dt={dt:.2f}s")
                qout.put(scan)
            stop = False
            while not qin.empty():
                if qin.get() == 'STOP':
                    logger.info("Stopping Nabu/3D FBP worker")
                    stop = True
            if stop:
                break
    except Exception as ex:
        logger.error(traceback.format_exc())
        logger.error("Error in FBP volume reconstruction")
        raise


class OutputWidgetHandler(logging.Handler):
    """ Custom logging handler sending logs to an output widget
    Original code: https://ipywidgets.readthedocs.io/en/7.6.5/examples/Output%20Widget.html
    """

    def __init__(self, *args, **kwargs):
        super(OutputWidgetHandler, self).__init__(*args, **kwargs)
        layout = {
            'width': '100%',
            'height': '160px',
            'border': '1px solid black'
        }
        self.out = ipw.Output(layout=layout)

    def emit(self, record):
        """ Overload of logging.Handler method """
        formatted_record = self.format(record)
        new_output = {
            'name': 'stdout',
            'output_type': 'stream',
            'text': formatted_record + '\n'
        }
        self.out.outputs = (new_output,) + self.out.outputs

    def show_logs(self):
        """ Show the logs """
        display(self.out)

    def clear_logs(self):
        """ Clear the current logs """
        self.out.clear_output()

    def crop_logs(self, nb=20):
        if len(self.out.outputs) > nb:
            self.out.outputs = self.out.outputs[:20]


class HoloTomoLiveWidget(ipw.VBox):
    """
    This widget allows to follow a live acquisition, with a new single-distance scan
    every few seconds. It uses multiple parallel process to watch for new scans, load
    them, perform a quick iterative phasing, reconstruct the 3D volume with Nabu, and
    display 3 slices of the object. It also allows to navigate through the series of
    scans to explore the corresponding volumes.
    """

    def __init__(self, prefix_dir=None, prefix_frame=None, next_scan=None, tomo_rot_center=None):
        super().__init__()
        # used to disable event handling when changing a slider/button value from the API,
        # using NoEventContextManager
        self._enable_events = True
        # Logger (multiprocessing friendly)
        self.logger = logging.getLogger(__name__)
        self.logger.setLevel(logging.INFO)
        self._logger_queue = Queue(-1)
        self.logger.addHandler(logging.handlers.QueueHandler(self._logger_queue))
        self.log_handler = OutputWidgetHandler()
        self.log_output = self.log_handler.out
        self._logger_queue_listener = logging.handlers.QueueListener(self._logger_queue, self.log_handler)
        self._logger_queue_listener.start()

        params = {
            'prefix_dir': '/data/visitor/ma4928/id16b/20240605/RAW_DATA/SIMAP/Marion/Marion_05_insitu_s{scan:03d}_1_',
            'prefix_frame': 'Marion_05_insitu_s{scan:03d}_1_{i:04d}.edf',
            'binning': 2,
            'next_scan': 1,
            'tomo_rot_center': 1021.4,

            # Reconstruction parameters
            'delta_beta': 180,
            'normalise': True,
            'doubleff': True,
            'obj_smooth': 0,

            'padding': 0,  # Needs to be handled properly (shared array sizes etc..)
            'padx': None,
            'pady': None,

            'dark_url': None,  # os.path.join(prefix_dir.format(scan=next_scan),'dark.edf')
            'ref_url': None,
        }
        self.par = Namespace(**params)
        if prefix_dir is not None:
            self.par.prefix_dir = prefix_dir
        if prefix_frame is not None:
            self.par.prefix_frame = prefix_frame
        if next_scan is not None:
            self.par.next_scan = next_scan
        if tomo_rot_center is not None:
            self.par.tomo_rot_center = tomo_rot_center

        # HBox with view and parameters
        self.box0 = None
        self.box0_params = None
        self.fig = None
        self.suptitle = None

        self._make_box()
        self.children = [self.box0, self.log_output]

        # process
        self.watcher = None
        self.watcher_qin = Queue()
        self.watcher_qout = Queue()
        self.loader = None
        self.loader_qin = Queue()
        self.loader_qout = Queue()
        self.algo = None
        self.algo_qin = Queue()
        self.algo_qout = Queue()
        self.fbp = None
        self.fbp_qin = Queue()
        self.fbp_qout = Queue()
        # Volume to display
        self.current_scan = None
        self.vol = None

        self.refresh_figure = False

    def _make_box(self):
        self.box0 = ipw.HBox()
        # Matplotlib figure
        with plt.ioff():
            plt.style.use('dark_background')
            self.fig = plt.figure(figsize=(16, 12))
            self.ax1 = None
            self.ax2 = None
            self.ax3 = None
            self.ax4 = None
        # Parameters & buttons
        self.box0_params = ipw.VBox()
        self.params_button = ipw.Button(description='Start',
                                        disabled=False, layout=ipw.Layout(width="auto"),
                                        button_style='',  # 'success', 'info', 'warning', 'danger' or ''
                                        tooltip='Click to activate the configuration & load the first scan',
                                        icon='play')
        self.params_button.on_click(self.on_start)

        self.toggle_live = ipw.ToggleButton(value=False, description='Live', disabled=True,
                                            tooltip='Enable live update',
                                            icon='check',
                                            layout=ipw.Layout(width="auto"))
        self.toggle_live.observe(self.on_toggle_live)

        params_dir_prefix_label = ipw.Text(value="Scans location:", disabled=True, layout=ipw.Layout(width="auto"))

        self.params_dir_prefix = ipw.Text(value=self.par.prefix_dir,
                                          placeholder='/path/to/scan_prefix{scan:03d}_1_',
                                          continuous_update=False,
                                          disabled=False, layout=ipw.Layout(width="25em"))
        params_prefix_label = ipw.Text(value="Files prefix:", disabled=True, layout=ipw.Layout(width="auto"))
        self.params_files_prefix = ipw.Text(value=self.par.prefix_frame,
                                            continuous_update=False,
                                            disabled=False, layout=ipw.Layout(width="25em"))
        self.params_delta_beta = ipw.FloatText(value=self.par.delta_beta,
                                               description='Delta/beta:',
                                               continuous_update=False,
                                               layout=ipw.Layout(width="auto"))
        self.params_binning = ipw.BoundedIntText(value=self.par.binning,
                                                 description='Binning:',
                                                 min=1, max=4,
                                                 continuous_update=False,
                                                 layout=ipw.Layout(width="auto"))
        self.params_padding = ipw.BoundedIntText(value=self.par.padding,
                                                 min=0, max=100,
                                                 description='Padding:',
                                                 continuous_update=False,
                                                 layout=ipw.Layout(width="auto"))
        self.toggle_normalise = ipw.ToggleButton(value=self.par.normalise, description='Normalise', disabled=False,
                                                 tooltip='Normalise frames to direct beam',
                                                 icon='check' if self.par.normalise else 'xmark',
                                                 layout=ipw.Layout(width="auto"))
        self.toggle_normalise.observe(self.on_toggle_normalise)
        self.toggle_doubleff = ipw.ToggleButton(value=self.par.doubleff, description='Double flat-field',
                                                disabled=False if self.par.normalise else 'xmark',
                                                tooltip='Normalise frames to average',
                                                icon='check',
                                                layout=ipw.Layout(width="auto"))
        self.toggle_doubleff.observe(self.on_toggle_doubleff)
        self.params_first_scan = ipw.IntText(value=self.par.next_scan,
                                             description='First scan:',
                                             continuous_update=False,
                                             layout=ipw.Layout(width="auto"))
        self.params_cor = ipw.FloatText(value=self.par.tomo_rot_center,
                                        description='Center of Rotation:',
                                        tooltip='Rotation center in pixels (not binned !)',
                                        continuous_update=False,
                                        layout=ipw.Layout(width="auto"))

        self.slider_x = ipw.IntSlider(value=0, min=0, max=0, step=1,
                                      description='X-slice:', orientation='horizontal',
                                      continuous_update=False, disabled=True,
                                      layout=ipw.Layout(width="auto"))
        self.slider_y = ipw.IntSlider(value=0, min=0, max=0, step=1,
                                      description='Y-slice:', orientation='horizontal',
                                      continuous_update=False, disabled=True,
                                      layout=ipw.Layout(width="auto"))
        self.slider_z = ipw.IntSlider(value=0, min=0, max=0, step=1,
                                      description='Z-slice:', orientation='horizontal',
                                      continuous_update=False, disabled=True,
                                      layout=ipw.Layout(width="auto"))
        self.slider_scan = ipw.IntSlider(value=0, min=0, max=0, step=1,
                                         description='Scan:', orientation='horizontal',
                                         continuous_update=False, disabled=True,
                                         layout=ipw.Layout(width="auto"))
        self.slider_scan.observe(self.on_slider_scan)
        self.params_refresh_nbscan = ipw.Button(description='Update nb scans',
                                                disabled=False, layout=ipw.Layout(width="auto"),
                                                button_style='',  # 'success', 'info', 'warning', 'danger' or ''
                                                tooltip="Use this if you don't want to wait for the live update. "
                                                        "This searches for the last frame in each folder. This "
                                                        "is only useful if you started while scans are being "
                                                        "accumulated, and you want to jump directly to a scan not "
                                                        "present initially.",
                                                icon='rotate')
        self.params_refresh_nbscan.on_click(self.update_nb_scans)
        self.params_refresh = ipw.Button(description='Refresh figure',
                                         disabled=False, layout=ipw.Layout(width="auto"),
                                         button_style='',  # 'success', 'info', 'warning', 'danger' or ''
                                         tooltip='Use this if the figure does not update anymore',
                                         icon='rotate')
        self.params_refresh.on_click(self.on_refresh_figure)

        self.box0_params.children = [
            self.params_button,
            self.toggle_live,
            params_dir_prefix_label,
            self.params_dir_prefix,
            params_prefix_label,
            self.params_files_prefix,
            self.params_delta_beta,
            self.params_binning,
            self.params_padding,
            self.toggle_normalise,
            self.toggle_doubleff,
            self.params_first_scan,
            self.params_cor,
            self.params_refresh_nbscan,
            self.params_refresh,
            self.slider_x,
            self.slider_y,
            self.slider_z,
            self.slider_scan,
        ]

        self.box0.children = [self.fig.canvas, self.box0_params]

    def on_start(self, v=None):
        self.logger.error(f"Button clicked: {self.params_button.description}")
        try:
            if False:
                if self.params_button.description == 'Start':
                    self.params_button.description = 'Restart [update parameters]'
                    self.params_button.icon = 'repeat'
                    self.toggle_live.value = True
            else:
                with NoEventContextManager(self):
                    # Disable parameter buttons
                    self.params_button.disabled = True
                    self.params_dir_prefix.disabled = True
                    self.params_files_prefix.disabled = True
                    self.params_delta_beta.disabled = True
                    self.params_binning.disabled = True
                    self.params_padding.disabled = True
                    self.toggle_normalise.disabled = True
                    self.toggle_doubleff.disabled = True
                    self.params_first_scan.disabled = True
                    self.params_cor.disabled = True
                    # Toggle live
                    self.toggle_live.value = True
                    self.toggle_live.button_style = 'success'
                    self.toggle_live.disabled = False
            self._on_start()
        except Exception as ex:
            self.logger.error(traceback.format_exc())
            self.logger.error("Error during start/stop")
            raise

    def _on_start(self, v=None):
        try:
            self.logger.info('Starting...')
            # First read parameters
            self.par.prefix_dir = self.params_dir_prefix.value
            self.par.prefix_frame = self.params_files_prefix.value
            self.par.binning = self.params_binning.value
            self.par.padding = self.params_padding.value
            self.par.next_scan = self.params_first_scan.value
            self.par.normalise = self.toggle_normalise.value
            self.par.doubleff = self.toggle_doubleff.value
            self.par.tomo_rot_center = self.params_cor.value

            # Read ID16B parameters
            self.logger.info('Reading experimental scan parameters...')
            params_id16b = get_params_id16b(self.par.prefix_dir.format(scan=self.par.next_scan), planes=(1,), sx0=None,
                                            sx=None, cx=None, pixel_size_detector=None, logger=self.logger)
            self.par.wavelength = params_id16b['wavelength']
            self.par.pixel_size = params_id16b['pixel_size'][0]
            self.par.distance = params_id16b['distance']
            self.par.nb_frame = params_id16b['tomo_n'] // self.par.binning
            self.logger.info('Reading experimental scan parameters... Done')
            # Load dark & ref - try next scan, 0 and 1
            load_ok = False
            for scan in [self.par.next_scan, 0, 1]:
                dark_url = os.path.join(self.par.prefix_dir.format(scan=self.par.next_scan), 'dark.edf')
                ref_url = os.path.join(self.par.prefix_dir.format(scan=self.par.next_scan), 'refHST0000.edf')
                try:
                    self.logger.info(f"Loading dark from {dark_url}")
                    self.dark = fabio.open(dark_url).data.astype(np.float32)
                    self.logger.info(f"Loading ref from {ref_url}")
                    self.ref = fabio.open(ref_url).data.astype(np.float32) - self.dark
                    if self.par.binning > 1:
                        self.dark = rebin(self.dark, self.par.binning)
                        self.ref = rebin(self.ref, self.par.binning)
                    load_ok = True
                except FileNotFoundError:
                    self.logger.error(f"Error loading either dark or ref from scan {scan}")
                if load_ok:
                    break
            if not load_ok:
                raise RuntimeError(f"Error loading either dark or ref from scans 0,1 "
                                   f"or {self.par.next_scan}. Check files/paths ?")
            # reshape as code is built for multi-distance
            ny0, nx0 = self.dark.shape
            self.ref = np.expand_dims(self.ref, 0)
            self.dark = np.expand_dims(self.dark, 0)

            # clipped pixels outside reconstructed cylinder
            ix, iy = np.meshgrid(np.arange(nx0) - nx0 / 2, np.arange(nx0) - nx0 / 2)
            r2 = ix ** 2 + iy ** 2
            self.idx_in2d = (r2 < (nx0 * nx0 / 4)).astype(np.float32)

            # Update slider limits
            self.slider_x.disabled = False
            self.slider_x.max = nx0 - 1
            self.slider_x.value = nx0 // 2
            self.slider_x.observe(self.on_slider_x)

            self.slider_y.disabled = False
            self.slider_y.max = nx0 - 1
            self.slider_y.value = nx0 // 2
            self.slider_y.observe(self.on_slider_y)

            self.slider_z.disabled = False
            self.slider_z.max = ny0 - 1
            self.slider_z.value = ny0 // 2
            self.slider_z.observe(self.on_slider_z)

            self.update_nb_scans()
            self.slider_scan.disabled = False

        except Exception as ex:
            self.logger.error(traceback.format_exc())
            self.logger.error("Error starting")
            raise

        self._stop_workers()
        self._init_padding()
        self._allocate_arrays()
        self._start_workers()
        self.run_check_new_volume()

    def update_nb_scans(self, v=None):
        # How many scans are already there ?
        scan = 0
        while True:
            if not os.path.exists(os.path.join(self.par.prefix_dir,
                                               self.par.prefix_frame).format(scan=scan, i=self.par.nb_frame - 1)):
                break
            scan += 1
        self.slider_scan.max = scan - 1
        self.logger.info(f"Found {self.slider_scan.max + 1} scans")

    def _start_workers(self):
        self._stop_workers()
        self.logger.info(f"Starting workers for parallel processing...")
        try:
            nb_frame, ny, nx = shared_arrays['ph_array'][1]
            self.watcher = Process(target=data_watcher, args=(
                self.watcher_qin, self.watcher_qout, self.par.prefix_dir, self.par.prefix_frame, nb_frame,
                self.par.next_scan, self.logger))
            self.watcher.start()
            self.loader = Process(target=data_loader, args=(
                self.loader_qin, self.loader_qout, self.watcher_qin, self.watcher_qout, self.logger, self.par, self.ref,
                self.dark))
            self.loader.start()
            self.algo = Process(target=run_ht, args=(
                self.algo_qin, self.algo_qout, self.loader_qin, self.loader_qout, self.logger, self.par, self.ref))
            self.algo.start()
            self.fbp = Process(target=run_fbp,
                               args=(self.fbp_qin, self.fbp_qout, self.algo_qin, self.algo_qout, self.logger, self.par))
            self.fbp.start()
        except Exception as ex:
            self.logger.error(traceback.format_exc())
            self.logger.error("Error starting workers")
            raise
        self.logger.info(f"Starting workers for parallel processing...Done")

    def _allocate_arrays(self):
        try:
            self.logger.info("Allocating shared arrays for parallel processing...")
            shared_arrays.clear()

            iobs_shape = (self.par.nb_frame, 1, self.par.ny, self.par.nx)
            iobs_array = Array(ctypes.c_float, int(np.prod(iobs_shape)))

            ph_shape = (self.par.nb_frame, self.par.ny0, self.par.nx0)
            ph_array = Array(ctypes.c_float, int(np.prod(ph_shape)))  # This could also be float16

            vol_shape = (self.par.ny0, self.par.nx0, self.par.nx0)
            vol_array = Array(ctypes.c_float, int(np.prod(vol_shape)))
            shared_arrays['iobs_array'] = iobs_array, iobs_shape
            shared_arrays['ph_array'] = ph_array, ph_shape
            shared_arrays['vol_array'] = vol_array, vol_shape
            self.logger.info("Allocating shared arrays for parallel processing... Done")

        except Exception as ex:
            self.logger.error(traceback.format_exc())
            self.logger.error("Error allocating shared arrays")
            raise

    def _stop_workers(self):
        try:
            if self.watcher is not None:
                self.logger.info("Stopping data watcher...")
                self.watcher_qin.put('STOP')
                self.watcher.join()
                self.watcher = None
            if self.loader is not None:
                self.logger.info("Stopping loader...")
                self.loader_qin.put('STOP')
                self.loader.join()
                self.loader = None
            if self.algo is not None:
                self.logger.info("Stopping algo process...")
                self.algo_qin.put('STOP')
                self.algo.join()
                self.algo = None
            if self.fbp is not None:
                self.logger.info("Stopping Nabu/FBP process...")
                self.fbp_qin.put('STOP')
                self.fbp.join()
                self.fbp = None
            self.logger.info("All process stopped...")
        except Exception as ex:
            self.logger.error(traceback.format_exc())
            self.logger.error("Error stopping workers")
            raise

    def _init_padding(self):
        try:
            ny, nx = self.dark.shape[-2:]  # Before padding
            if self.par.padding:
                pady = padx = self.par.padding
                primesy, primesx = primes(ny + 2 * pady), primes(nx + 2 * padx)
                if max(primesy) > 13:
                    padup = pady
                    while not test_smaller_primes(ny + 2 * padup, required_dividers=[2]):
                        padup += 1
                    paddown = pady
                    while not test_smaller_primes(ny + 2 * paddown, required_dividers=[2]):
                        paddown -= 1
                    s = "The Y dimension (with padding=%d,%d) is incompatible with a radix FFT:\n" \
                        "  ny=%d primes=%s  (should be <=13)\n" \
                        "  Closest acceptable padding values: %d or %d" % \
                        (pady, padx, ny + 2 * pady, str(primesy), paddown, padup)
                    # self.logger.warning(s)
                    # Use closest padding option, unless we pad to 2x the array size, in which case
                    # we can only go down
                    if (padup - pady < pady - paddown) and (self.par.padding != 2) and (pady - paddown > 0):
                        pady = padup
                    else:
                        pady = paddown

                if max(primesx) > 13:
                    padup = padx
                    while not test_smaller_primes(nx + 2 * padup, required_dividers=[2]):
                        padup += 1
                    paddown = padx
                    while not test_smaller_primes(nx + 2 * paddown, required_dividers=[2]):
                        paddown -= 1
                    s = "The X dimension (with padding=%d) is incompatible with a radix FFT:\n" \
                        "  nx=%d primes=%s  (should be <=13)\n" \
                        "  Closest acceptable padding values: %d or %d" % \
                        (padx, nx + 2 * padx, str(primesx), paddown, padup)
                    # self.logger.warning(s)
                    # Use the closest padding option, unless we pad to 2x the array size, in which case
                    # we can only go down
                    if (padup - padx < padx - paddown) and (self.par.padding != 2) and (padx - paddown > 0):
                        padx = padup
                    else:
                        padx = paddown
                self.par.padx = padx
                self.par.pady = pady
                self.logger.info(f"Padding along x:{padx}   and y:{pady}")
            else:
                self.par.padx = 0
                self.par.pady = 0
            self.par.ny0 = ny
            self.par.nx0 = nx
            self.par.ny = ny + 2 * self.par.pady
            self.par.nx = nx + 2 * self.par.padx
        except Exception as ex:
            self.logger.error(traceback.format_exc())
            self.logger.error("Error stopping workers")
            raise

    async def check_new_volume(self):
        while True:
            try:
                if not self.fbp_qout.empty():
                    while not self.fbp_qout.empty():
                        self.current_scan = self.fbp_qout.get()
                        with NoEventContextManager(self):
                            if self.current_scan > self.slider_scan.max:
                                self.slider_scan.max = self.current_scan
                            self.slider_scan.value = self.current_scan
                    self.logger.info(f"Got new volume: {self.current_scan}")
                    vol_array, vol_shape = shared_arrays['vol_array']
                    with vol_array.get_lock():
                        self.vol = np.frombuffer(vol_array.get_obj(), dtype=np.float32).reshape(vol_shape).copy()
                    # NaN mask to clip reconstruction
                    self.vol *= self.idx_in2d
                    self._save_axial_plots()  # only if does not already exists
                    self.plot()
            except Exception as ex:
                self.logger.error(traceback.format_exc())
                self.logger.error("Error logging")
                raise
            await asyncio.sleep(2)

    def run_check_new_volume(self):
        try:
            self.logger.info("Starting loop to check for new volume")
            self.loop = asyncio.get_event_loop()
            self.loop_task = self.loop.create_task(self.check_new_volume())
        except Exception as ex:
            self.logger.error(traceback.format_exc())
            self.logger.error("Error starting volum check loop")
            raise

    def _save_axial_plots(self):
        filename = os.path.split(self.par.prefix_dir)[-1].format(scan=self.current_scan) + '.png'
        if not os.path.exists(filename):
            # Save axial plots for current volumes
            fig = Figure(figsize=(16, 12))
            ny0, nx0 = self.vol.shape[:2]
            ax1, ax2, ax3, ax4 = fig.subplots(2, 2).flat
            c = self.vol[ny0 // 2].astype(np.float32)
            vmin, vmax = np.nanpercentile(c, (1, 99))
            ax1.imshow(c, vmin=vmin, vmax=vmax, cmap='gray')

            # We use a layout similar to what is used in imagej
            c = self.vol[..., nx0 // 2].astype(np.float32).transpose()
            vmin, vmax = np.nanpercentile(c, (1, 99))
            ax2.imshow(c, vmin=vmin, vmax=vmax, cmap='gray')

            c = self.vol[:, nx0 // 2].astype(np.float32)
            vmin, vmax = np.nanpercentile(c, (1, 99))
            ax3.imshow(c, vmin=vmin, vmax=vmax, cmap='gray')

            ax4.set_visible(False)
            fig.suptitle(f"{os.path.split(self.par.prefix_dir)[-1].format(scan=self.current_scan)}")
            fig.tight_layout()
            canvas = FigureCanvasAgg(fig)
            self.logger.info(f"Saving scan {self.current_scan} axial plots to: {filename}")
            canvas.print_figure(filename, dpi=150)

    def on_refresh_figure(self, v=None):
        self.refresh_figure = True
        self.plot()

    def plot(self):
        self.logger.info(f"Plotting volume #{self.current_scan}")
        try:
            if self.vol is not None:
                ny0, nx0 = self.vol.shape[:2]
                if self.refresh_figure:
                    self.fig.clf()
                if self.ax1 is None or self.refresh_figure:
                    # First plot
                    ax1, ax2, ax3, ax4 = self.fig.subplots(2, 2).flat
                    c = self.vol[self.slider_z.value].astype(np.float32)
                    vmin, vmax = np.nanpercentile(c, (1, 99))
                    self.ax1 = ax1.imshow(c, vmin=vmin, vmax=vmax, cmap='gray')

                    # We use a layout similar to what is used in imagej
                    c = self.vol[..., self.slider_x.value].astype(np.float32).transpose()
                    vmin, vmax = np.nanpercentile(c, (1, 99))
                    self.ax2 = ax2.imshow(c, vmin=vmin, vmax=vmax, cmap='gray')

                    c = self.vol[:, self.slider_y.value].astype(np.float32)
                    vmin, vmax = np.nanpercentile(c, (1, 99))
                    self.ax3 = ax3.imshow(c, vmin=vmin, vmax=vmax, cmap='gray')

                    if True:
                        ax4.set_visible(False)
                    else:
                        iobs_array, iobs_shape = shared_arrays['iobs_array']
                        ph_array, ph_shape = shared_arrays['ph_array']
                        with ph_array.get_lock():
                            ph = np.frombuffer(ph_array.get_obj(), dtype=np.float32).reshape(ph_shape)
                            c = ph[self.current_scan].copy()
                        vmin, vmax = np.nanpercentile(c, (1, 99))
                        self.ax4 = ax4.imshow(c, vmin=vmin, vmax=vmax, cmap='gray')
                    self.suptitle = plt.suptitle(
                        f"{os.path.split(self.par.prefix_dir)[-1].format(scan=self.current_scan)}")
                    plt.tight_layout()
                    self.refresh_figure = False
                else:
                    c = self.vol[self.slider_z.value].astype(np.float32)
                    vmin, vmax = np.nanpercentile(c, (1, 99))
                    self.ax1.set_data(c)
                    self.ax1.set_norm(Normalize(vmin=vmin, vmax=vmax))
                    self.ax1.set_extent((0, nx0, 0, nx0))  # In case we zoomed

                    # We use a layout similar to what is used in imagej
                    c = self.vol[..., self.slider_x.value].astype(np.float32).transpose()
                    vmin, vmax = np.nanpercentile(c, (1, 99))
                    self.ax2.set_data(c)
                    self.ax2.set_norm(Normalize(vmin=vmin, vmax=vmax))
                    self.ax2.set_extent((0, ny0, 0, nx0))

                    c = self.vol[:, self.slider_y.value].astype(np.float32)
                    vmin, vmax = np.nanpercentile(c, (1, 99))
                    self.ax3.set_data(c)
                    self.ax3.set_norm(Normalize(vmin=vmin, vmax=vmax))
                    self.ax3.set_extent((0, nx0, 0, ny0))

                    if self.ax4 is not None:
                        iobs_array, iobs_shape = shared_arrays['iobs_array']
                        ph_array, ph_shape = shared_arrays['ph_array']
                        with ph_array.get_lock():
                            ph = np.frombuffer(ph_array.get_obj(), dtype=np.float32).reshape(ph_shape)
                            c = ph[self.current_scan].copy()
                        vmin, vmax = np.nanpercentile(c, (1, 99))
                        self.ax4.set_data(c)
                        self.ax4.set_norm(Normalize(vmin=vmin, vmax=vmax))
                        self.ax4.set_extent((0, nx0, 0, ny0))
                self.suptitle.set_text(f"{os.path.split(self.par.prefix_dir)[-1].format(scan=self.current_scan)}")

                self.fig.canvas.draw()
                self.fig.canvas.flush_events()

        except Exception as ex:
            self.logger.error(traceback.format_exc())
            self.logger.error("Error plotting volume")
            raise

    def on_slider_x(self, v):
        if v['name'] != 'value' or not self._enable_events:
            return
        # self.logger.info(f"on_slider_x: x={self.slider_x.value} {v}")
        try:
            if self.ax2 is not None and self.vol is not None:
                ny0, nx0 = self.vol.shape[:2]
                c = self.vol[..., self.slider_x.value].astype(np.float32).transpose()
                vmin, vmax = np.nanpercentile(c, (1, 99))
                self.ax2.set_data(c)
                self.ax2.set_norm(Normalize(vmin=vmin, vmax=vmax))
                self.ax2.set_extent((0, ny0, 0, nx0))
                self.fig.canvas.draw()
                self.fig.canvas.flush_events()
        except Exception as ex:
            self.logger.error(traceback.format_exc())
            self.logger.error("Error updating plot")
            raise

    def on_slider_y(self, v):
        if v['name'] != 'value' or not self._enable_events:
            return
        # self.logger.info(f"on_slider_y: y={self.slider_y.value} {v}")
        try:
            if self.ax3 is not None and self.vol is not None:
                ny0, nx0 = self.vol.shape[:2]
                c = self.vol[:, self.slider_y.value].astype(np.float32)
                vmin, vmax = np.nanpercentile(c, (1, 99))
                self.ax3.set_data(c)
                self.ax3.set_norm(Normalize(vmin=vmin, vmax=vmax))
                self.ax3.set_extent((0, nx0, 0, ny0))
                self.fig.canvas.draw()
                self.fig.canvas.flush_events()
        except Exception as ex:
            self.logger.error(traceback.format_exc())
            self.logger.error("Error updating plot")
            raise

    def on_slider_z(self, v):
        if v['name'] != 'value' or not self._enable_events:
            return
        # self.logger.info(f"on_slider_z: z={self.slider_z.value} {v}")
        try:
            if self.ax1 is not None and self.vol is not None:
                ny0, nx0 = self.vol.shape[:2]
                c = self.vol[self.slider_z.value].astype(np.float32)
                vmin, vmax = np.nanpercentile(c, (1, 99))
                self.ax1.set_data(c)
                self.ax1.set_norm(Normalize(vmin=vmin, vmax=vmax))
                self.ax1.set_extent((0, nx0, 0, nx0))  # In case we zoomed
                self.fig.canvas.draw()
                self.fig.canvas.flush_events()
        except Exception as ex:
            self.logger.error(traceback.format_exc())
            self.logger.error("Error updating plot")
            raise

    def on_slider_scan(self, v):
        if v['name'] != 'value' or not self._enable_events:
            return
        try:
            self.logger.info(f"Selected scan {self.slider_scan.value}")
            if self.toggle_live.value:
                self.logger.info(f"Manual scan selection -disabling live update")
                self.toggle_live.value = False
            self.watcher_qout.put(self.slider_scan.value)
        except Exception as ex:
            self.logger.error(traceback.format_exc())
            self.logger.error("Error updating plot")
            raise

    def on_toggle_live(self, v):
        if v['name'] != 'value' or not self._enable_events:
            return
        if not self.toggle_live.value:
            if self.watcher is not None:
                self.logger.info('Stopping Live update (data watcher)')
                self.watcher_qin.put('STOP')
                self.watcher.join()
                self.watcher = None
            self.toggle_live.button_style = 'warning'
            self.toggle_live.icon = 'xmark'
        else:
            self.logger.info('Starting Live update')
            self.watcher = Process(target=data_watcher, args=(
                self.watcher_qin, self.watcher_qout, self.par.prefix_dir, self.par.prefix_frame,
                self.par.nb_frame, self.slider_scan.value + 1, self.logger))
            self.watcher.start()
            self.toggle_live.button_style = 'success'
            self.toggle_live.icon = 'check'

    def on_toggle_doubleff(self, v):
        if v['name'] != 'value' or not self._enable_events:
            return
        if not self.toggle_doubleff.value:
            self.toggle_doubleff.icon = 'xmark'
        else:
            self.toggle_doubleff.icon = 'check'

    def on_toggle_normalise(self, v):
        if v['name'] != 'value' or not self._enable_events:
            return
        if not self.toggle_normalise.value:
            self.toggle_normalise.icon = 'xmark'
        else:
            self.toggle_normalise.icon = 'check'
