from os import path
from time import time
from math import ceil
import numpy as np
from silx.io.url import DataUrl
from ...utils import remove_items_from_list
from ...resources.logger import LoggerOrPrint
from ...resources.utils import extract_parameters
from ...io.reader import ChunkReader, HDF5Loader
from ...preproc.ccd import Log, CCDFilter
from ...preproc.flatfield import FlatFieldDataUrls
from ...preproc.distortion import DistortionCorrection
from ...preproc.shift import VerticalShift
from ...preproc.double_flatfield import DoubleFlatField
from ...preproc.phase import PaganinPhaseRetrieval
from ...preproc.ctf import CTFPhaseRetrieval, GeoPars
from ...reconstruction.sinogram import SinoBuilder, SinoNormalization
from ...misc.rotation import Rotation
from ...reconstruction.rings import MunchDeringer
from ...misc.unsharp import UnsharpMask
from ...misc.histogram import PartialHistogram, hist_as_2Darray
from ..utils import use_options, pipeline_step, get_subregion
from ..datadump import DataDumpManager
from ..writer import WriterManager
from ..detector_distortion_provider import DetectorDistortionProvider

# For now we don't have a plain python/numpy backend for reconstruction
try:
    from ...reconstruction.fbp_opencl import Backprojector
except:
    Backprojector = None


class ChunkedPipeline:
    """
    Pipeline for "regular" full-field tomography.
    Data is processed by chunks. A chunk consists in K contiguous lines of all the radios.
    In parallel geometry, a chunk of K radios lines gives K sinograms,
    and equivalently K reconstructed slices.
    """

    backend = "numpy"
    FlatFieldClass = FlatFieldDataUrls
    DoubleFlatFieldClass = DoubleFlatField
    CCDCorrectionClass = CCDFilter
    PaganinPhaseRetrievalClass = PaganinPhaseRetrieval
    CTFPhaseRetrievalClass = CTFPhaseRetrieval
    UnsharpMaskClass = UnsharpMask
    ImageRotationClass = Rotation
    VerticalShiftClass = VerticalShift
    SinoDeringerClass = MunchDeringer
    MLogClass = Log
    SinoBuilderClass = SinoBuilder
    SinoNormalizationClass = SinoNormalization
    FBPClass = Backprojector
    ConebeamClass = None  # unsupported on CPU
    HistogramClass = PartialHistogram

    _default_extra_options = {}

    # These steps are skipped if the reconstruction is done in two stages.
    # The first stage will skip these steps, and the second stage will do these stages after merging sinograms.
    _reconstruction_steps = ["sino_rings_correction", "reconstruction", "save", "histogram"]

    def __init__(
        self, process_config, chunk_shape, margin=None, logger=None, use_grouped_mode=False, extra_options=None
    ):
        """
        Initialize a "Chunked" pipeline.

        Parameters
        ----------
        processing_config: `ProcessConfig`
            Process configuration.
        chunk_shape: tuple
            Shape of the chunk of data to process, in the form (n_angles, n_z, n_x).
            It has to account for possible cropping of the data, eg. [:, start_z:end_z, start_x:end_x]
            where start_xz and/or end_xz can be other than None.
        margin: tuple, optional
            Margin to use, in the form ((up, down), (left, right)).
            It is used for example when performing phase retrieval or a convolution-like operation:
            some extra data is kept to avoid boundaries issues.
            These boundaries are then discarded: the data volume is eventually cropped as
            `data[U:D, L:R]` where `((U, D), (L, R)) = margin`
            If not provided, no margin is applied.
        logger: `nabu.app.logger.Logger`, optional
            Logger class
        extra_options: dict, optional
            Advanced extra options.


        Notes
        ------

        Using `margin` results in a lesser number of reconstructed slices.
        More specifically, if `margin = (V, H)`, then there will be `delta_z - 2*V`
        reconstructed slices (if the sub-region is in the middle of the volume)
        or `delta_z - V` reconstructed slices (if the sub-region is on top or bottom
        of the volume).
        """
        self.logger = LoggerOrPrint(logger)
        self._set_params(process_config, chunk_shape, extra_options, margin, use_grouped_mode)
        self._init_pipeline()

    def _set_params(self, process_config, chunk_shape, extra_options, margin, use_grouped_mode):
        self.process_config = process_config
        self.dataset_info = self.process_config.dataset_info
        self.processing_steps = self.process_config.processing_steps.copy()
        self.processing_options = self.process_config.processing_options
        self._set_chunk_shape(chunk_shape, use_grouped_mode)
        self.set_subregion(None)
        self._set_margin(margin)
        self._set_extra_options(extra_options)
        self._callbacks = {}
        self._steps_name2component = {}
        self._steps_component2name = {}

    def _set_chunk_shape(self, chunk_shape, use_grouped_mode):
        if len(chunk_shape) != 3:
            raise ValueError("Expected chunk_shape to be a tuple of length 3 in the form (n_z, n_y, n_x)")
        self.chunk_shape = tuple(int(c) for c in chunk_shape)  # cast to int, as numpy.int64 can make pycuda crash
        # TODO: sanity check (eg. compare to size of radios in dataset_info) ?
        # (n_a, n_z, n_x)
        self.radios_shape = (
            ceil(self.chunk_shape[0] / self.process_config.subsampling_factor),
            self.chunk_shape[1] // self.process_config.binning[1],
            self.chunk_shape[2] // self.process_config.binning[0],
        )
        self.n_angles = self.radios_shape[0]
        self.n_slices = self.radios_shape[1]
        self._grouped_processing = False
        if use_grouped_mode or self.chunk_shape[0] < len(self.process_config.rotation_angles(subsampling=False)):
            # TODO allow a certain tolerance in this case ?
            # Reconstruction is still possible (albeit less accurate) if delta is small
            self._grouped_processing = True
            self.logger.debug("Only a subset of angles is processed - Reconstruction will be skipped")
            self.processing_steps, _ = remove_items_from_list(self.processing_steps, self._reconstruction_steps)

    def _set_margin(self, margin):
        if margin is None:
            U, D, L, R = None, None, None, None
        else:
            ((U, D), (L, R)) = get_subregion(margin, ndim=2)

        # Replace "None" with zeros
        U, D, L, R = U or 0, D or 0, L or 0, R or 0

        self.margin = ((U, D), (L, R))
        self._margin_up = U
        self._margin_down = D
        self._margin_left = L
        self._margin_right = R
        self.use_margin = (U + D + L + R) > 0
        self.n_recs = self.chunk_shape[1] - sum(self.margin[0])
        self.radios_cropped_shape = (self.radios_shape[0], self.radios_shape[1] - U - D, self.radios_shape[2] - L - R)
        if self.use_margin:
            self.n_slices -= sum(self.margin[0])

    def set_subregion(self, sub_region):
        """
        Set the data volume sub-region to process.
        Note that processing margin, if any, is contained within the sub-region.

        Parameters
        -----------
        sub_region: tuple
            Data volume sub-region, in the form ((start_a, end_a), (start_z, end_z), (start_x, end_x))
            where the data volume has a layout (angles, Z, X)
        """
        n_angles = self.dataset_info.n_angles
        n_x, n_z = self.dataset_info.radio_dims
        c_a, c_z, c_x = self.chunk_shape
        if sub_region is None:
            # By default, take the sub-region around central slice
            sub_region = (
                (0, c_a),
                (n_z // 2 - c_z // 2, n_z // 2 - c_z // 2 + c_z),
                (n_x // 2 - c_x // 2, n_x // 2 - c_x // 2 + c_x),
            )
        else:
            sub_region = get_subregion(sub_region, ndim=3)
            # check sub-region
            for i, start_end in enumerate(sub_region):
                start, end = start_end
                if start is not None and end is not None:
                    if end - start != self.chunk_shape[i]:
                        raise ValueError(
                            "Invalid (start, end)=(%d, %d) for sub-region (dimension %d): chunk shape is %s, but %d-%d=%d != %d"
                            % (start, end, i, str(self.chunk_shape), end, start, end - start, self.chunk_shape[i])
                        )
            #
        self.logger.debug("Set sub-region to %s" % (str(sub_region)))
        self.sub_region = sub_region
        self._sub_region_xz = sub_region[2] + sub_region[1]

    def _set_extra_options(self, extra_options):
        self.extra_options = self._default_extra_options.copy()
        self.extra_options.update(extra_options or {})

    #
    # Callbacks
    #

    def register_callback(self, step_name, callback):
        """
        Register a callback for a pipeline processing step.

        Parameters
        ----------
        step_name: str
            processing step name
        callback: callable
            A function. It will be executed once the processing step `step_name`
            is finished. The function takes only one argument: the class instance.
        """
        if step_name not in self.processing_steps:
            raise ValueError("'%s' is not in processing steps %s" % (step_name, self.processing_steps))
        if step_name in self._callbacks:
            self._callbacks[step_name].append(callback)
        else:
            self._callbacks[step_name] = [callback]

    #
    # Memory management
    #

    def _allocate_array(self, shape, dtype, name=None):
        return np.zeros(shape, dtype=dtype)

    def _allocate_recs(self, ny, nx, n_slices=None):
        n_slices = n_slices or self.n_slices
        self.recs = self._allocate_array((n_slices, ny, nx), "f", name="recs")

    #
    # Runtime attributes
    #

    @property
    def sub_region_xz(self):
        """
        Return the currently processed sub-region in the form
        (start_x, end_x, start_z, end_z)
        """
        return self._sub_region_xz

    @property
    def z_min(self):
        return self._sub_region_xz[2]

    @property
    def sino_shape(self):
        return self.process_config.sino_shape(binning=True, subsampling=True, after_ha=True)

    @property
    def sinos_shape(self):
        return (self.n_slices,) + self.sino_shape

    def get_slice_start_index(self):
        return self.z_min + self._margin_up

    #
    # Pipeline initialization
    #

    def _init_pipeline(self):
        self._allocate_radios()
        self._init_data_dump()
        self._init_reader()
        self._init_flatfield()
        self._init_double_flatfield()
        self._init_ccd_corrections()
        self._init_radios_rotation()
        self._init_phase()
        self._init_unsharp()
        self._init_radios_movements()
        self._init_mlog()
        self._init_sino_normalization()
        self._init_sino_rings_correction()
        self._init_reconstruction()
        self._init_histogram()
        self._init_writer()

    def _allocate_radios(self):
        self.radios = np.zeros(self.radios_shape, dtype=np.float32)
        self.data = self.radios  # alias

    def _init_data_dump(self):
        self._resume_from_step = self.processing_options["read_chunk"].get("step_name", None)
        self.datadump_manager = DataDumpManager(
            self.process_config, self.sub_region, margin=self.margin, logger=self.logger
        )
        # When using "grouped processing", sinogram has to be dumped.
        # If it was not specified by user, force sinogram dump
        # Perhaps these lines should be moved directly to DataDumpManager.
        if self._grouped_processing and not self.process_config.dump_sinogram:
            sino_dump_fname = self.process_config.get_save_steps_file("sinogram")
            self.datadump_manager._configure_dump("sinogram", force_dump_to_fname=sino_dump_fname)
            self.logger.debug("Will dump sinogram to %s" % self.datadump_manager.data_dump["sinogram"].fname)

    @use_options("read_chunk", "chunk_reader")
    def _init_reader(self):
        options = self.processing_options["read_chunk"]
        self._update_reader_configuration()

        process_file = options.get("process_file", None)
        if process_file is None:
            # Standard case - start pipeline from raw data
            if self.process_config.nabu_config["preproc"]["detector_distortion_correction"] is None:
                self.detector_corrector = None
            else:
                self.detector_corrector = DetectorDistortionProvider(
                    detector_full_shape_vh=self.process_config.dataset_info.radio_dims[::-1],
                    correction_type=self.process_config.nabu_config["preproc"]["detector_distortion_correction"],
                    options=self.process_config.nabu_config["preproc"]["detector_distortion_correction_options"],
                )
            # ChunkReader always take a non-subsampled dictionary "files".
            self.chunk_reader = ChunkReader(
                self._read_options["files"],
                sub_region=self.sub_region_xz,
                data_buffer=self.radios,
                pre_allocate=False,
                detector_corrector=self.detector_corrector,
                convert_float=True,
                binning=options["binning"],
                dataset_subsampling=options["dataset_subsampling"],
            )
        else:
            # Resume pipeline from dumped intermediate step
            self.chunk_reader = HDF5Loader(
                process_file,
                options["process_h5_path"],
                sub_region=self.datadump_manager.get_read_dump_subregion(),
                data_buffer=self.radios,
                pre_allocate=False,
            )
            self._resume_from_step = options["step_name"]
            self.logger.debug(
                "Load subregion %s from file %s" % (str(self.chunk_reader.sub_region), self.chunk_reader.fname)
            )

    def _update_reader_configuration(self):
        """
        Modify self.processing_options["read_chunk"] to select a subset of the files, if needed
        (i.e when processing only a subset of the images stack)
        """
        self._read_options = self.processing_options["read_chunk"].copy()
        if self.n_angles == self.process_config.n_angles(subsampling=True):
            # Nothing to do if the full angular range is processed in one shot
            return
        if self._resume_from_step is not None:
            if self._resume_from_step == "sinogram":
                msg = "It makes no sense to use 'grouped processing' when resuming from sinogram"
                self.logger.fatal(msg)
                raise ValueError(msg)
            # Nothing to do if we resume the processing from a given step
            return
        input_data_files = {}
        files_indices = sorted(self._read_options["files"].keys())
        angle_idx_start, angle_idx_end = self.sub_region[0]
        for i in range(angle_idx_start, angle_idx_end):
            idx = files_indices[i]
            input_data_files[idx] = self._read_options["files"][idx]
        self._read_options["files"] = input_data_files

    @use_options("flatfield", "flatfield")
    def _init_flatfield(self):
        self._ff_options = self.processing_options["flatfield"].copy()
        # Use chunk_reader.files instead of process_config.projs_indices(subsampling=True), because
        # chunk_reader might read only a subset of the files (in "grouped mode")
        self._ff_options["projs_indices"] = list(self.chunk_reader.files_subsampled.keys())
        if self._ff_options.get("normalize_srcurrent", False):
            a_start_idx, a_end_idx = self.sub_region[0]
            subs = self.process_config.subsampling_factor
            self._ff_options["radios_srcurrent"] = self._ff_options["radios_srcurrent"][a_start_idx:a_end_idx:subs]

        distortion_correction = None
        if self._ff_options["do_flat_distortion"]:
            self.logger.info("Flats distortion correction will be applied")
            self.FlatFieldClass = FlatFieldDataUrls  # no GPU implementation available, force this backend
            estimation_kwargs = {}
            estimation_kwargs.update(self._ff_options["flat_distortion_params"])
            estimation_kwargs["logger"] = self.logger
            distortion_correction = DistortionCorrection(
                estimation_method="fft-correlation", estimation_kwargs=estimation_kwargs, correction_method="interpn"
            )

        # FlatField parameter "radios_indices" must account for subsampling
        self.flatfield = self.FlatFieldClass(
            self.radios_shape,
            flats=self.dataset_info.flats,
            darks=self.dataset_info.darks,
            radios_indices=self._ff_options["projs_indices"],
            interpolation="linear",
            distortion_correction=distortion_correction,
            sub_region=self.sub_region_xz,
            detector_corrector=self.detector_corrector,
            binning=self._ff_options["binning"],
            radios_srcurrent=self._ff_options["radios_srcurrent"],
            flats_srcurrent=self._ff_options["flats_srcurrent"],
            convert_float=True,
        )

    @use_options("double_flatfield", "double_flatfield")
    def _init_double_flatfield(self):
        options = self.processing_options["double_flatfield"]
        avg_is_on_log = options["sigma"] is not None
        result_url = None
        if options["processes_file"] not in (None, ""):
            result_url = DataUrl(
                file_path=options["processes_file"],
                data_path=(self.dataset_info.hdf5_entry or "entry") + "/double_flatfield/results/data",
            )
            self.logger.info("Loading double flatfield from %s" % result_url.file_path())
        if (self.n_angles < self.process_config.n_angles(subsampling=True)) and result_url is None:
            raise ValueError(
                "Cannot use double-flatfield when processing subset of radios. Please use the 'nabu-double-flatfield' command"
            )
        self.double_flatfield = self.DoubleFlatFieldClass(
            self.radios_shape,
            result_url=result_url,
            sub_region=self.sub_region_xz,
            detector_corrector=self.detector_corrector,
            input_is_mlog=False,
            output_is_mlog=False,
            average_is_on_log=avg_is_on_log,
            sigma_filter=options["sigma"],
        )

    @use_options("ccd_correction", "ccd_correction")
    def _init_ccd_corrections(self):
        options = self.processing_options["ccd_correction"]
        self.ccd_correction = self.CCDCorrectionClass(
            self.radios_shape[1:], median_clip_thresh=options["median_clip_thresh"]
        )

    @use_options("rotate_projections", "projs_rot")
    def _init_radios_rotation(self):
        options = self.processing_options["rotate_projections"]
        center = options["center"]
        if center is None:
            nx, ny = self.radios_shape[1:][::-1]  # after binning
            center = (nx / 2 - 0.5, ny / 2 - 0.5)
        center = (center[0], center[1] - self.z_min)
        self.projs_rot = self.ImageRotationClass(
            self.radios_shape[1:], options["angle"], center=center, mode="edge", reshape=False
        )
        self._tmp_rotated_radio = self._allocate_array(self.radios_shape[1:], "f", name="tmp_rotated_radio")

    @use_options("radios_movements", "radios_movements")
    def _init_radios_movements(self):
        options = self.processing_options["radios_movements"]
        self._vertical_shifts = options["translation_movements"][:, 1]
        self.radios_movements = self.VerticalShiftClass(self.radios.shape, self._vertical_shifts)

    @use_options("phase", "phase_retrieval")
    def _init_phase(self):
        options = self.processing_options["phase"]
        if options["method"] == "CTF":
            translations_vh = getattr(self.dataset_info, "ctf_translations", None)
            geo_pars_params = options["ctf_geo_pars"].copy()
            geo_pars_params["logger"] = self.logger
            geo_pars = GeoPars(**geo_pars_params)
            self.phase_retrieval = self.CTFPhaseRetrievalClass(
                self.radios_shape[1:],
                geo_pars,
                options["delta_beta"],
                lim1=options["ctf_lim1"],
                lim2=options["ctf_lim2"],
                logger=self.logger,
                fftw_num_threads=None,  # TODO tune in advanced params of nabu config file
                use_rfft=True,
                normalize_by_mean=options["ctf_normalize_by_mean"],
                translation_vh=translations_vh,
            )
        else:
            self.phase_retrieval = self.PaganinPhaseRetrievalClass(
                self.radios_shape[1:],
                distance=options["distance_m"],
                energy=options["energy_kev"],
                delta_beta=options["delta_beta"],
                pixel_size=options["pixel_size_m"],
                padding=options["padding_type"],
                # TODO tune in advanced params of nabu config file
                fftw_num_threads=None,
            )
        if self.phase_retrieval.use_fftw:
            self.logger.debug(
                "%s using FFTW with %d threads"
                % (self.phase_retrieval.__class__.__name__, self.phase_retrieval.fftw.num_threads)
            )

    @use_options("unsharp_mask", "unsharp_mask")
    def _init_unsharp(self):
        options = self.processing_options["unsharp_mask"]
        self.unsharp_mask = self.UnsharpMaskClass(
            self.radios_shape[1:],
            options["unsharp_sigma"],
            options["unsharp_coeff"],
            mode="reflect",
            method=options["unsharp_method"],
        )

    @use_options("take_log", "mlog")
    def _init_mlog(self):
        options = self.processing_options["take_log"]
        self.mlog = self.MLogClass(
            self.radios_shape, clip_min=options["log_min_clip"], clip_max=options["log_max_clip"]
        )

    @use_options("sino_normalization", "sino_normalization")
    def _init_sino_normalization(self):
        options = self.processing_options["sino_normalization"]
        self.sino_normalization = self.SinoNormalizationClass(
            kind=options["method"],
            radios_shape=self.radios_cropped_shape,
            normalization_array=options["normalization_array"],
        )

    @use_options("sino_rings_correction", "sino_deringer")
    def _init_sino_rings_correction(self):
        options = self.processing_options["sino_rings_correction"]
        fw_params = extract_parameters(options["user_options"])
        fw_sigma = fw_params.pop("sigma", 1.0)
        self.sino_deringer = self.SinoDeringerClass(fw_sigma, sinos_shape=self.sinos_shape, **fw_params)

    @use_options("reconstruction", "reconstruction")
    def _init_reconstruction(self):
        options = self.processing_options["reconstruction"]
        if options["method"] == "FBP" and self.FBPClass is None:
            raise ValueError("No usable FBP module was found")
        if options["method"] == "cone" and self.ConebeamClass is None:
            raise ValueError("No usable cone-beam module was found")

        if options["method"] == "FBP":
            n_slices = self.n_slices
            radios_shape_for_sino_builder = self.radios_cropped_shape
            self.reconstruction = self.FBPClass(
                self.sinos_shape[1:],
                angles=options["angles"],
                rot_center=options["fbp_rotation_axis_position"],
                filter_name=options["fbp_filter_type"],
                slice_roi=self.process_config.rec_roi,
                padding_mode=options["padding_type"],
                extra_options={
                    "scale_factor": 1.0 / options["pixel_size_cm"],
                    "axis_correction": options["axis_correction"],
                    "centered_axis": options["centered_axis"],
                    "clip_outer_circle": options["clip_outer_circle"],
                    "filter_cutoff": options["fbp_filter_cutoff"],
                },
            )
            if options["fbp_filter_type"] is None:
                self.reconstruction.fbp = self.reconstruction.backproj
        if options["method"] == "cone":
            radios_shape_for_sino_builder = self.radios_shape
            n_slices = self.n_slices + sum(self.margin[0])
            # For numerical stability, normalize all lengths with respect to detector pixel size
            pixel_size_m = self.dataset_info.pixel_size * 1e-6
            source_sample_dist = options["source_sample_dist"] / pixel_size_m
            sample_detector_dist = options["sample_detector_dist"] / pixel_size_m
            self.reconstruction = self.ConebeamClass(  # pylint: disable=E1102
                (self.radios_shape[1],) + self.sino_shape,
                source_sample_dist,
                sample_detector_dist,
                angles=options["angles"],
                # TODO one center for each angle to handle "x translations"
                rot_center=options["rotation_axis_position"],
                pixel_size=1,
            )
        self.sino_builder = self.SinoBuilderClass(
            radios_shape=radios_shape_for_sino_builder,
            rot_center=options["rotation_axis_position"],
            halftomo=options["enable_halftomo"],
            angles=options["angles"],
        )
        self._allocate_recs(*self.process_config.rec_shape, n_slices=n_slices)
        if options["method"] == "cone":
            self.sinos = self._allocate_array(self.sino_builder.output_shape, "f", name="sinos")

    @use_options("histogram", "histogram")
    def _init_histogram(self):
        options = self.processing_options["histogram"]
        self.histogram = self.HistogramClass(method="fixed_bins_number", num_bins=options["histogram_bins"])

    @use_options("save", "writer")
    def _init_writer(self):
        options = self.processing_options["save"]
        metadata = {
            "process_name": "reconstruction",
            "processing_index": 0,
            # TODO this one takes too much time to write, not useful for partial files
            # "processing_options": self.processing_options,
            #
            "nabu_config": self.process_config.nabu_config,
            "entry": getattr(self.dataset_info.dataset_scanner, "entry", "entry"),
        }
        writer_extra_options = {
            "jpeg2000_compression_ratio": options["jpeg2000_compression_ratio"],
            "float_clip_values": options["float_clip_values"],
            "tiff_single_file": options.get("tiff_single_file", False),
            "single_output_file_initialized": getattr(self.process_config, "single_output_file_initialized", False),
        }
        self.writer = WriterManager(
            options["location"],
            options["file_prefix"],
            file_format=options["file_format"],
            overwrite=options["overwrite"],
            start_index=self.get_slice_start_index(),
            logger=self.logger,
            metadata=metadata,
            histogram=("histogram" in self.processing_steps),
            extra_options=writer_extra_options,
        )

    #
    # Pipeline execution
    #

    @pipeline_step("chunk_reader", "Reading data")
    def _read_data(self):
        self.logger.debug("Region = %s" % str(self.sub_region))
        t0 = time()
        self.chunk_reader.load_data()
        el = time() - t0
        shp = self.chunk_reader.data.shape
        self.logger.info("Read subvolume %s in %.2f s" % (str(shp), el))

    @pipeline_step("flatfield", "Applying flat-field")
    def _flatfield(self):
        self.flatfield.normalize_radios(self.radios)

    @pipeline_step("double_flatfield", "Applying double flat-field")
    def _double_flatfield(self, radios=None):
        if radios is None:
            radios = self.radios
        self.double_flatfield.apply_double_flatfield(radios)

    @pipeline_step("ccd_correction", "Applying CCD corrections")
    def _ccd_corrections(self, radios=None):
        if radios is None:
            radios = self.radios
        _tmp_radio = self._allocate_array(radios.shape[1:], "f", name="tmp_ccdcorr_radio")
        for i in range(radios.shape[0]):
            self.ccd_correction.median_clip_correction(radios[i], output=_tmp_radio)
            radios[i][:] = _tmp_radio[:]

    @pipeline_step("projs_rot", "Rotating projections")
    def _rotate_projections(self, radios=None):
        if radios is None:
            radios = self.radios
        tmp_radio = self._tmp_rotated_radio
        for i in range(radios.shape[0]):
            self.projs_rot.rotate(radios[i], output=tmp_radio)
            radios[i][:] = tmp_radio[:]

    @pipeline_step("phase_retrieval", "Performing phase retrieval")
    def _retrieve_phase(self):
        for i in range(self.radios.shape[0]):
            self.phase_retrieval.retrieve_phase(self.radios[i], output=self.radios[i])

    @pipeline_step("unsharp_mask", "Performing unsharp mask")
    def _apply_unsharp(self):
        for i in range(self.radios.shape[0]):
            self.radios[i] = self.unsharp_mask.unsharp(self.radios[i])

    @pipeline_step("mlog", "Taking logarithm")
    def _take_log(self):
        self.mlog.take_logarithm(self.radios)

    @pipeline_step("radios_movements", "Applying radios movements")
    def _radios_movements(self, radios=None):
        if radios is None:
            radios = self.radios
        self.radios_movements.apply_vertical_shifts(radios, list(range(radios.shape[0])))

    def _crop_radios(self):
        if self.use_margin:
            self._orig_radios = self.radios
            if self.processing_options.get("reconstruction", {}).get("method", None) == "cone":
                return
            ((U, D), (L, R)) = self.margin
            self.logger.debug(
                "Cropping radios from %s to %s" % (str(self.radios_shape), str(self.radios_cropped_shape))
            )
            U, D, L, R = U or None, -D or None, L or None, -R or None
            self.radios = self.radios[:, U:D, L:R]  # view

    @pipeline_step("sino_normalization", "Normalizing sinograms")
    def _normalize_sinos(self, radios=None):
        if radios is None:
            radios = self.radios
        sinos = radios.transpose((1, 0, 2))
        self.sino_normalization.normalize(sinos)

    def _dump_sinogram(self):
        if self.datadump_manager is not None:
            self.datadump_manager.dump_data_to_file("sinogram", self.radios)

    @pipeline_step("sino_deringer", "Removing rings on sinograms")
    def _destripe_sinos(self):
        sinos = np.rollaxis(self.radios, 1, 0)  # view
        self.sino_deringer.remove_rings(sinos)  # TODO check it works with non-contiguous view

    @pipeline_step("reconstruction", "Reconstruction")
    def _reconstruct(self):
        """
        Reconstruction for parallel geometry.
        For each target slice: get the corresponding sinogram, apply some processing, then reconstruct
        """
        if self.processing_options["reconstruction"]["method"] == "cone":
            self._reconstruct_cone()
            return
        for i in range(self.n_slices):
            sino = self.sino_builder.get_sino(self.radios, i)
            if self.sino_deringer is not None:
                self.sino_deringer.remove_rings(sino)
            self.reconstruction.fbp(sino, output=self.recs[i])

    def _reconstruct_cone(self):
        """
        This reconstructs the entire sinograms stack at once
        """
        self.sino_builder.get_sinos(self.radios, output=self.sinos)
        z_min, z_max = self.sub_region_xz[2:]
        n_z = self.process_config.radio_shape(binning=True)[0]

        self.reconstruction.reconstruct(  # pylint: disable=E1101
            self.sinos,
            output=self.recs,
            relative_z_position=((z_min + z_max) / self.process_config.binning_z / 2) - n_z / 2,
        )

    @pipeline_step("histogram", "Computing histogram")
    def _compute_histogram(self, data=None):
        if data is None:
            data = self.recs
        self.recs_histogram = self.histogram.compute_histogram(data)

    @pipeline_step("writer", "Saving data")
    def _write_data(self, data=None):
        if data is None:
            data = self.recs
        self.writer.write_data(data)
        self.logger.info("Wrote %s" % self.writer.fname)
        self._write_histogram()
        self.process_config.single_output_file_initialized = True

    def _write_histogram(self):
        if "histogram" not in self.processing_steps:
            return
        self.logger.info("Saving histogram")
        self.writer.write_histogram(
            hist_as_2Darray(self.recs_histogram),
            processing_index=1,
            config={
                "file": path.basename(self.writer.fname),
                "bins": self.processing_options["histogram"]["histogram_bins"],
            },
        )

    def _process_finalize(self):
        if self.use_margin:
            self.radios = self._orig_radios

    def __repr__(self):
        res = "%s(%s, margin=%s)" % (self.__class__.__name__, str(self.chunk_shape), str(self.margin))
        binning = self.process_config.binning
        subsampling = self.process_config.subsampling_factor

        if binning != (1, 1) or subsampling > 1:
            if binning != (1, 1):
                res += "\nImages binning: %s" % (str(binning))
            if subsampling:
                res += "\nAngles subsampling: %d" % subsampling
            res += "\nRadios chunk: %s ---> %s" % (self.chunk_shape, self.radios_shape)

        if self.use_margin:
            res += "\nMargin: %s" % (str(self.margin))
            res += "\nRadios chunk: %s ---> %s" % (str(self.radios_shape), str(self.radios_cropped_shape))

        res += "\nCurrent subregion: %s" % (str(self.sub_region))
        for step_name in self.processing_steps:
            res += "\n- %s" % (step_name)
        return res

    def _process_chunk(self):
        self._flatfield()
        self._double_flatfield()
        self._ccd_corrections()
        self._rotate_projections()
        self._retrieve_phase()
        self._apply_unsharp()
        self._take_log()
        self._radios_movements()
        self._crop_radios()
        self._normalize_sinos()
        self._dump_sinogram()
        # self._destripe_sinos()
        self._reconstruct()
        self._compute_histogram()
        self._write_data()
        self._process_finalize()

    def _reset_reader_subregion(self):
        if self._resume_from_step is None:
            # Normal mode - read data from raw radios
            self.chunk_reader._set_subregion(self.sub_region_xz)
            self.chunk_reader._init_reader()
            self.chunk_reader._loaded = False
        else:
            # Resume from a checkpoint. In this case, we have to re-initialize "datadump manager"
            # sooner to configure start_xyz, end_xyz
            self._init_data_dump()
            self.chunk_reader._set_subregion(self.datadump_manager.get_read_dump_subregion())
            self.chunk_reader._loaded = False
        if self._grouped_processing:
            self._update_reader_configuration()
            self.chunk_reader._set_files(self._read_options["files"])

    def _reset_sub_region(self, sub_region):
        self.set_subregion(sub_region)
        # When sub_region is changed, all components involving files reading have to be updated
        self._reset_reader_subregion()
        self._init_flatfield()  # reset flatfield
        self._init_writer()
        self._init_double_flatfield()
        self._init_data_dump()

    def process_chunk(self, sub_region):
        self._reset_sub_region(sub_region)
        self._read_data()
        self._process_chunk()
