import os
from dataclasses import asdict, dataclass
from threading import Semaphore
from time import time

import fabio
import numpy as np
from nabu.io.reader import import_h5_to_dict
from nabu.io.writer import convert_dict_values
from nabu.resources.logger import Logger, LoggerOrPrint, PrinterLogger
from nabu.utils import check_supported, compare_dicts, copy_dict_items
from silx.io import get_data
from silx.io.dictdump import dicttoh5
from silx.io.url import DataUrl
from tomoscan.io import HDF5File

from .utils import make_azims, make_detector_atten, parse_poni_file

path = os.path


@dataclass
class AIConfiguration:
    n_pts: int
    unit: str
    polarization_factor: float
    correct_solid_angle: bool
    detector: str
    poni_file: str
    mask_file: str
    flatfield_file: str
    dark_file: str
    error_model: str
    azimuthal_range: tuple
    radial_range: tuple
    ai_method: str
    trim_method: str = None
    trim_n_pts: int = 0
    trim_bounds: tuple = None
    do_stack_mean: bool = False
    pixel_splitting: str = "no"
    average_xy: str = None
    correct_detector_attenuation: bool = False


def create_azimuthal_integrator(ai_config):
    """
    Create a AzimuthalIntegrator object from a configuration.

    Parameters
    ----------
    ai_config: AIConfiguration
        Data structure describing the azimuthal integration configuration

    Returns
    -------
    azimuthal_integrator: pyFAI.integrator.azimuthal.AzimuthalIntegrator
        Azimuthal integrator instance
    """
    import pyFAI.integrator.azimuthal  # import here to avoid multiple opencl contexts creation

    # Create azimuthal integrator.
    # pyFAI.load(ai_config.poni_file) cannot be used directly if the detector file therein
    # can't be accessed. So we have to use this trick instead
    poni = parse_poni_file(ai_config.poni_file)
    ai_kwargs = {
        name: float(poni[name])
        for name in ["poni1", "poni2", "rot1", "rot2", "rot3", "wavelength"]
    }
    ai_kwargs["dist"] = poni["distance"]
    # azimuthal_integrator = pyFAI.integrator.azimuthal.AzimuthalIntegrator(**ai_kwargs) # does not work either
    azimuthal_integrator = pyFAI.integrator.azimuthal.AzimuthalIntegrator()
    for param_name, value in ai_kwargs.items():
        setattr(azimuthal_integrator, param_name, value)

    # Set detector
    detector = pyFAI.detector_factory(ai_config.detector)
    azimuthal_integrator.detector = detector

    # Set mask
    mask_file = ai_config.mask_file
    if isinstance(mask_file, str):
        mask_file = DataUrl(file_path=mask_file, scheme="fabio")
    mask_array = get_data(mask_file)
    detector.mask = mask_array

    # Set flatfield and/or dark
    corrections = [
        {"fname": ai_config.flatfield_file, "function": detector.set_flatfield},
        {"fname": ai_config.dark_file, "function": detector.set_darkcurrent},
    ]
    for correction in corrections:
        fname = correction["fname"]
        if fname is None:
            continue
        try:
            with fabio.open(fname) as f:
                img = f.data
            correction["function"](
                img,
            )  # detector.set_flatfield(img) or detector.set_darkcurrent(img)
        except Exception as exc:
            print("ERROR: unable to read %s: %s" % (fname, str(exc)))

    return azimuthal_integrator


class StackIntegrator:
    _default_extra_options = {
        "target_device": None,
        "use_gpu_decompression": False,
    }

    def __init__(
        self,
        ai_config,
        logger=None,
        existing_output="skip",
        extra_options=None,
    ):
        """
        Initialize an Integrator object.

        Parameters
        ----------
        ai_config: AIConfiguration
            Azimuthal Integration configuration
        logger: Logger, optional
            Logger object. If not provided, messages are shown with the print() function
        existing_output: str, optional
            What do do in the case where the output file already exists:
               - "raise": raise an error
               - "overwrite": overwrite the existing file
               - "skip": do not process the current stack


        Other Parameters
        ----------------
        extra_options: dict, optional
            Dictionary of advanced options. Current values are:
                - "scan_num_as_h5_entry": False
                    Whether to use the current scan number as HDF5 entry.
        """
        self._create_logger(logger)
        self._semaphore = Semaphore()  # TODO is it still needed ?
        self._configure_extra_options(extra_options)
        self._set_ai_config(ai_config)
        self._initialize_ai_engine()
        self._init_data_reader_and_writer(existing_output)

    #
    # Initialization
    #
    def _configure_extra_options(self, user_extra_options):
        self.extra_options = self._default_extra_options.copy()
        self.extra_options.update(user_extra_options or {})

    def _create_logger(self, logger):
        if logger is None:
            self.logger = PrinterLogger()
        elif isinstance(logger, str):
            self.logger = Logger(logger, level="DEBUG", logfile=logger + ".log")
        else:
            self.logger = logger

    def _load_hdf5plugin(self):
        # Don't do it earlier - OpenMP messes with threads affinity
        import hdf5plugin

    def _set_ai_config(self, ai_config):
        check_supported(
            ai_config.trim_method,
            [None, "median", "sigma_clip"],
            "trim_method",
        )
        self.ai_method = ai_config.ai_method

        if isinstance(self.ai_method, str):
            from platform import processor

            import pyFAI.integrator.azimuthal  # have to do it to actually get available integration methods
            from pyFAI.method_registry import IntegrationMethod

            ai_methods = IntegrationMethod.select_method(
                dim=1,
                split=ai_config.pixel_splitting,
                algo="csr",
                impl=ai_config.ai_method,
                target=self.extra_options["target_device"],
            )
            pocl_str = "Portable Computing Language"
            if processor() == "ppc64le":
                # pocl is the only working Opencl platform on powerpc
                ai_methods = [
                    method
                    for method in ai_methods
                    if pocl_str in method.target_name and method.target_type == "gpu"
                ]
            else:
                # pocl is the only non-working Opencl platform on x86_64
                ai_methods = [
                    method
                    for method in ai_methods
                    if pocl_str not in method.target_name
                ]
            self.ai_method = ai_methods[0]
            self.logger.debug("Using AI method %s" % str(self.ai_method))

        self.ai_config = AIConfiguration(**asdict(ai_config))  # shallow copy
        self.do_error_propagation = self.ai_config.error_model is not None
        self.azimuthal_ranges = make_azims(*self.ai_config.azimuthal_range)
        self.n_azim_slices = int(self.ai_config.azimuthal_range[-1])
        self.logger.debug(
            "Using %d azimuthal slices: %s"
            % (self.n_azim_slices, self.azimuthal_ranges),
        )

    def _get_ai_radial_positions(self):
        # warmup + get radial positions by integrating first frame
        # This assumes that the radial pos are the same for all images
        """
        When the radial range limits are not set in the input arguments, they \
        are extracted as the minimum and the maximum of the Q or 2theta array \
        (depending on the choice of radial units) known to the integrator. \
        These are both in the final units (Q: A^-1, 2theta: deg) and in standard \
        units (Q: nm^-1, 2theta:rad) so that calls to function 'setup_sparse_integrator' \
        will impose this radial range (pos0_range), work in standard units (by setting scale=False) \
        and finally converting to the selected units.
        """
        if self.ai_config.unit == "q_A^-1":
            sivals = np.ravel(self.ai.qArray())  # qArray in nm^-1 (??): convert to A^-1
            radialvals = 0.1 * sivals
        elif self.ai_config.unit == "2th_deg":
            sivals = np.ravel(
                self.ai.center_array(unit="2th_rad"),
            )  # 2theta in radians: convert to deg
            radialvals = np.degrees(sivals)
        elif self.ai_config.unit == "r_mm":
            sivals = np.ravel(
                self.ai.rArray(),
            )  # rArray gives radius in meters: convert to mm
            radialvals = 1e3 * sivals
        if self.ai_config.radial_range is None:
            self.ai_config.radial_range = (min(radialvals), max(radialvals))
            # self.si_xrange = (min(sivals), max(sivals))
            self.logger.warning(
                f"Radial range reset from None to {self.ai_config.radial_range} {self.ai_config.unit}",
            )
            self.integration_kwargs["radial_range"] = self.ai_config.radial_range
        # else:
        #     self.si_xrange = (sivals[np.argmin(abs(self.ai_config.radial_range[0]-radialvals))],
        #                       sivals[np.argmin(abs(self.ai_config.radial_range[1]-radialvals))])

        first_frame = np.zeros(
            self.azimuthal_integrator.detector.shape,
            dtype=np.float32,
        )
        ai_res = self.integrate_image(first_frame)
        self.radial_pos = ai_res.radial

    def _initialize_ai_engine(self):
        """
        Initialize a pyFAI engine based on a AIConfiguration object.
        """
        ai_config = self.ai_config

        self.azimuthal_integrator = create_azimuthal_integrator(ai_config)
        self.ai = self.azimuthal_integrator  # shorthand

        # Configure integrator call
        self._current_azimuthal_range = None
        self.integration_kwargs = {
            "unit": ai_config.unit,
            "polarization_factor": ai_config.polarization_factor,
            "correctSolidAngle": ai_config.correct_solid_angle,
            "method": self.ai_method,
            "radial_range": ai_config.radial_range,
        }
        if ai_config.trim_method == "median":
            self.integration_kwargs.update(
                {
                    "npt_rad": ai_config.n_pts,
                    "npt_azim": ai_config.trim_n_pts or ai_config.n_pts,
                    "percentile": ai_config.trim_bounds,
                    "dummy": 0.0,  # see https://github.com/silx-kit/pyFAI/issues/2187
                },
            )
            self.do_error_propagation = False
        elif ai_config.trim_method == "sigma_clip":
            self.integration_kwargs.update(
                {
                    "npt": ai_config.n_pts,
                    # "npt_azim": ai_config.trim_n_pts, # not in sigma_clip_ng
                    "thres": ai_config.trim_bounds[0],  # scalar only
                    "error_model": ai_config.error_model,  # TODO allow for "azimuthal" error model
                    # TODO: cutoff = 0 : Chauvenet criterion (drop-cutoff)
                },
            )
        else:
            self.integration_kwargs.update(
                {
                    "error_model": ai_config.error_model,
                },
            )
        if self.ai_config.flatfield_file is not None:
            self.integration_kwargs["flat"] = (
                self.ai.flatfield
            )  # not sure this is automatically propagated!

        # Warmup, create AI engine
        self.ai_intensities = None
        self.errors = None
        self._get_ai_radial_positions()
        self._get_ai_engine()
        self._setup_attenuation_correction()

    def _setup_attenuation_correction(self):
        if (
            not (self.ai_config.correct_detector_attenuation)
            or self.ai_config.trim_method == "median"
        ):
            return
        ai_ones = self.ai.integrate1d_ng(
            np.ones(self.ai.detector.shape),
            self.ai_config.n_pts,
            unit="2th_rad",
            polarization_factor=self.ai_config.polarization_factor,
            error_model=self.ai_config.error_model,
            method=self.ai_method,
        )
        twotheta_rad = ai_ones[0]

        self._detector_atten = make_detector_atten(self.ai.wavelength, twotheta_rad)
        # Not sure if actually different from self._detector_atten ?
        self.integration_kwargs["absorption"] = make_detector_atten(
            self.ai.wavelength,
            self.ai.center_array(shape=self.ai.detector.shape, unit="2th_rad"),
        )

    def _get_ai_engine(self):
        from pyFAI.method_registry import IntegrationMethod

        engines_methods = list(self.ai.engines.keys())
        opencl_engines = []
        for engine_method in engines_methods:
            # when using integrate1d_legacy, names are str
            if isinstance(engine_method, str) and "ocl" in engine_method:
                opencl_engines.append(engine_method)
            # when using _ng methods, keys are IntegrationMethod
            if (
                isinstance(engine_method, IntegrationMethod)
                and engine_method.impl_lower == "opencl"
            ):
                opencl_engines.append(engine_method)
        if opencl_engines == []:
            self.logger.warning(
                "No OpenCL engine found. Processing will likely be slow!",
            )
            engine_name = engines_methods[0]  # ?
        else:
            engine_name = opencl_engines[0]  # ?
        engine = self.ai.engines[engine_name].engine
        self.logger.debug(
            "Using pyFAI engine %s - OpenCL context: %s"
            % (str(engine), str(getattr(engine, "ctx", None))),
        )
        self._ai_engine = engine

    def _init_data_reader_and_writer(self, existing_output):
        self._load_hdf5plugin()
        self.existing_output_behavior = existing_output
        if self.extra_options["use_gpu_decompression"]:
            ...  # TODO

    #
    # Getters
    #

    def get_current_config(self, data_location=None):
        conf = {
            "ai_config": asdict(self.ai_config),
            "pixel_size": (
                self.ai.pixel1,
                self.ai.pixel2,
            ),  # (vertical, horizontal), in meter
            "wavelength": self.ai.wavelength,  # meter
            "pyfai_method": str(self.ai_method),
            "distance": self.ai.dist,  # sample-detector distance in meter
            "radial_positions": self.radial_pos,
        }
        if data_location is not None:
            conf["data_location"] = data_location
        return conf

    def _get_results_arrays(self, data):
        n_images = data.shape[0]
        results_shape = (self.n_azim_slices, n_images, self.ai_config.n_pts)

        if self.ai_intensities is None or self.ai_intensities.shape != results_shape:
            self.ai_intensities = np.zeros(results_shape, dtype="f")
            if self.ai_config.error_model is not None:
                self.ai_errors = np.zeros_like(self.ai_intensities)

    #
    # Execution
    #
    def _load_data(self, h5_url, start_idx=None, end_idx=None):
        self.logger.debug("Loading data")
        self._current_data_location = str(h5_url)
        with HDF5File(h5_url.file_path(), "r") as f:
            t0 = time()
            # Read the entire dataset slice. This works for both 3D and 4D arrays.
            data = f[h5_url.data_path()][start_idx:end_idx]
            el = time() - t0
        self.logger.debug(
            "Read %s elements in %.1f s (%.3f GB/s - %.1f FPS)"
            % (str(data.shape), el, data.nbytes / 1e9 / el, data.shape[0] / el),
        )
        # If data is 4D and the new axis (axis 1) is of size 1, squeeze it out.
        if data.ndim == 4 and data.shape[1] == 1:
            data = np.squeeze(data, axis=1)
            self.logger.debug(
                "Squeezed singleton axis from data. New shape: %s" % (data.shape,),
            )
        return data

    def integrate_image(self, image, azimuth_range=None):
        if self.ai_config.trim_method == "median":
            xye = self.ai.medfilt1d(image, **self.integration_kwargs)
        elif self.ai_config.trim_method == "sigma_clip":
            xye = self.ai.sigma_clip_ng(image, **self.integration_kwargs)
        else:
            xye = self.ai.integrate1d_ng(
                image,
                self.ai_config.n_pts,
                azimuth_range=azimuth_range,
                **self.integration_kwargs,
            )
        return xye

    def _integrate_stack(self, data):
        self.logger.debug("Integrating data")
        self._get_results_arrays(data)

        for az_idx, azim_range in enumerate(self.azimuthal_ranges):
            if self.n_azim_slices > 1:
                self.logger.info("azimuthal range %s" % self._current_azimuthal_range)

            t0 = time()
            for i in range(data.shape[0]):
                ai_result = self.integrate_image(data[i], azimuth_range=azim_range)
                self.ai_intensities[az_idx, i] = ai_result.intensity
                if self.do_error_propagation:
                    self.ai_errors[az_idx, i] = ai_result.sigma
            if self.ai_config.do_stack_mean:
                average_image = np.average(data, axis=0)
                xye = self.integrate_image(average_image, azimuth_range=azim_range)
                exptime = self.dataset.exposure_time
                if len(xye) == 3:  # self.do_error_propagation ?
                    self._avexy_res = np.transpose([xye[0], xye[1] / exptime, xye[2]])
                elif len(xye) == 2:
                    self._avexy_res = np.transpose([xye[0], xye[1] / exptime])

            el = time() - t0
            self.logger.debug(
                "Integrated %d images in %.1f s (%.1f FPS)" % (i + 1, el, (i + 1) / el),
            )

    def _save_result(self, output_file):

        result_dict = {
            "azimuthal integration": {
                "results": {
                    str("azim_%02d" % (i + 1)): {
                        "intensities": self.ai_intensities[i],
                        "azimuthal_range": self.azimuthal_ranges[i],
                    }
                    for i in range(self.n_azim_slices)
                },
                "configuration": self.get_current_config(self._current_data_location),
            },
        }

        if self.do_error_propagation:
            for i in range(self.n_azim_slices):
                result_dict["azimuthal integration"]["results"][
                    str("azim_%02d" % (i + 1))
                ]["error"] = self.ai_errors[i]
        if self.ai_config.do_stack_mean:
            result_dict["azimuthal integration"]["results"]["averaged stack"] = (
                self._avexy_res
            )
        result_dict = convert_dict_values(
            result_dict,
            {None: "None"},
        )  # patch silx dictdump

        # TODO more defensive - try/except
        dicttoh5(result_dict, output_file)

        self.logger.info("Wrote file %s" % output_file)

    def process_data(self, data, output):
        """
        Perform azimuthal integration on data.

        Parameters
        ----------
        data: DataUrl or array
            Data to process (stack of images)
        output: str or array
            Output file or array

        Notes
        -----
        Processing from/to in-memory arrays is not supported yet
        """
        if not isinstance(data, DataUrl):
            raise TypeError("Please provide a DataUrl for data")
        if not isinstance(output, str):
            raise TypeError("Please provide a file path for output")

        if result_already_exists(
            output,
            "azimuthal integration/configuration",
            self.existing_output_behavior,
            self.get_current_config(data_location=str(data)),
            sections_to_ignore=None,
            logger=self.logger,
        ):
            return output

        data_array = self._load_data(data)
        self._integrate_stack(data_array)
        self._save_result(output)
        return output


def result_already_exists(
    out_fname,
    config_path,
    existing_output_behavior,
    current_config,
    sections_to_ignore=None,
    logger=None,
):
    """
    Return true if a given filename contains integration result with the same configuration.
    Basically: "was this data already integrated ?"

    Parameters
    ----------
    out_fname: str
        Path to the output file
    config_path: str
        Configuration within the output file (eg. "azimuthal integration/configuration")
    existing_output_behavior: str
        What to do if the results already exist (possibly with the same configuration). Possible values are:
          - "skip": don't overwrite the existing results, proceed to the next stack
          - "overwrite": discard existing results, re-do processing and overwrite the existing file
          - "raise": raise an error
    current_config: dict
        Dictionary with the current azimuthal integration configuration
    sections_to_ignore: list, optional
        Which sections to ignore when comparing configurations
    logger: logging object, optional
        Logging object. If not provided, messages will be sent in stdout
    """
    sections_to_ignore = sections_to_ignore or []
    logger = LoggerOrPrint(logger)
    check_supported(
        existing_output_behavior,
        ["skip", "overwrite", "raise"],
        "existing_output_behavior",
    )

    if not path.exists(out_fname):
        return None
    msg = "File %s already exists" % out_fname
    if existing_output_behavior == "raise":
        logger.fatal(msg)
        raise ValueError(msg)
    elif existing_output_behavior == "overwrite":
        logger.warning(msg + " - it will be overwritten as requested")
        return False
    elif existing_output_behavior == "skip":
        try:
            existing_file_config = import_h5_to_dict(out_fname, config_path)
            # posixpath.join(self._current_data_entry, "azimuthal integration", "configuration")
        except KeyError:
            logger.debug(msg + " but no entry '%s' within. Proceeding." % config_path)
            existing_file_config = None
        except (RuntimeError, OSError):
            logger.exception(msg + " but file seems to be invalid. Deleting the file!")
            os.unlink(out_fname)
            existing_file_config = None
        if existing_file_config is None:
            return False

        my_conf = copy_dict_items(
            current_config,
            list(set(current_config.keys()) - set(sections_to_ignore)),
        )

        # Ignore "units"
        for key, val in my_conf.items():
            if isinstance(key, tuple) and key[1] == "unit":
                existing_file_config[key] = val
        diff = compare_dicts(existing_file_config, my_conf)
        if diff is not None:
            msg = msg + ", but cannot skip as configuration is not the same!"
            logger.warning(msg)
            return False
        else:
            logger.warning(msg + " - configuration is the same, skipping")
            return True
    return False
