import logging
import os
import re
from pathlib import Path
import warnings

import h5py
import numpy as np
from PIL import Image
from ewokscore import Task
from tomoscan.esrf.scan.nxtomoscan import NXtomoScan
from nabu.preproc.flatfield import FlatField


logger = logging.getLogger(__name__)


def _auto_intensity_bounds(image: np.ndarray) -> tuple[float, float]:
    """Compute robust lower/upper bounds for scaling to 8-bit."""
    finite = np.asarray(image, dtype=np.float32)
    finite = finite[np.isfinite(finite)]
    if finite.size == 0:
        return 0.0, 255.0

    upper_candidates = finite[finite < 1e9]
    if upper_candidates.size == 0:
        upper_candidates = finite

    lower = float(np.percentile(finite, 0.01))
    upper = float(np.percentile(upper_candidates, 99.99))
    if not np.isfinite(lower) or not np.isfinite(upper):
        return 0.0, 255.0
    if lower == upper:
        upper = lower + 1.0
    return lower, upper


def clean_angle_key(angle_key):
    """Convert angle key like '90.00000009(1)' to float, or leave float as is."""
    if isinstance(angle_key, float):
        return angle_key  # already clean
    cleaned = re.sub(r"\(.*?\)", "", angle_key)  # remove '(1)' etc.
    return float(cleaned)


class BuildProjectionsGallery(
    Task,
    input_names=["nx_path", "reduced_darks_path", "reduced_flats_path"],
    optional_input_names=[
        "bounds",
        "angle_step",
        "output_binning",
        "output_format",
        "overwrite",
    ],
    output_names=["processed_data_dir", "gallery_path"],
):
    def run(self):
        """
        Creates a gallery of images from the NXtomoScan object.
        """

        self.gallery_output_format = str(
            self.get_input_value("output_format", "png")
        ).lower()
        self.gallery_overwrite = self.get_input_value("overwrite", True)
        self.gallery_output_binning = self.get_input_value("output_binning", 2)
        bounds = self.get_input_value("bounds", None)
        angle_step = self.get_input_value("angle_step", 90)

        # Use the directory of the output file as the processed data directory.
        nx_path = Path(self.inputs.nx_path)
        processed_data_dir = nx_path.parent
        gallery_dir = self.get_gallery_dir(processed_data_dir)
        os.makedirs(gallery_dir, exist_ok=True)

        # Open the NXtomoScan object.
        self.nxtomoscan = NXtomoScan(str(nx_path), entry="entry0000")

        angles, slices = self.get_slices_by_angle_step(angle_step)
        corrected_slices = self.flat_field_correction(slices)

        for angle, slice in zip(angles, corrected_slices):
            gallery_file_path = self.get_gallery_file_path(gallery_dir, nx_path, angle)
            Path(gallery_file_path).parent.mkdir(parents=True, exist_ok=True)

            # Process the image and save it in the gallery.
            self._save_to_gallery(gallery_file_path, slice, bounds)

        self.outputs.processed_data_dir = str(processed_data_dir)
        self.outputs.gallery_path = str(gallery_dir)

    def get_flats_from_h5(
        self, reduced_flat_path: str, data_path: str = "entry0000/flats"
    ) -> np.ndarray:
        """
        Loads the data from an HDF5 file.
        """
        with h5py.File(reduced_flat_path, "r") as h5f:
            for idx in h5f[data_path]:
                data = h5f[data_path][idx]
                flats_idx = int(idx)
                flats_data = data[()]
        return {flats_idx: flats_data}

    def get_darks_from_h5(
        self, reduced_dark_path: str, data_path: str = "entry0000/darks"
    ) -> np.ndarray:
        """
        Loads the data from an HDF5 file.
        """
        with h5py.File(reduced_dark_path, "r") as h5f:
            for idx in h5f[data_path]:
                data = h5f[data_path][idx]
                darks_idx = int(idx)
                darks_data = data[()]
        return {darks_idx: darks_data}

    def flat_field_correction(self, slices):
        """
        Applies flat field correction to the slices.
        """
        slices = np.asarray(slices, dtype=np.float32)
        reduced_darks = self.get_darks_from_h5(self.inputs.reduced_darks_path)
        reduced_flats = self.get_flats_from_h5(self.inputs.reduced_flats_path)
        x, y = slices[0].shape
        radios_shape = (len(slices), x, y)
        flat_field = FlatField(
            radios_shape=radios_shape, flats=reduced_flats, darks=reduced_darks
        )
        warnings.filterwarnings(
            "ignore",
            message=".*encountered in divide",
            category=RuntimeWarning,
            module="nabu.preproc.flatfield",
        )
        normalized_slices = flat_field.normalize_radios(slices)
        return normalized_slices

    def get_gallery_dir(self, processed_data_dir: Path) -> str:
        return processed_data_dir / "gallery"

    def get_gallery_file_path(self, gallery_dir, nx_path: Path, angle: float) -> str:
        filename = f"{nx_path.stem}_{angle:.2f}deg.{self.gallery_output_format}"
        gallery_path = gallery_dir / filename
        return str(gallery_path)

    def get_proj_from_data_url(self, data_url) -> np.ndarray:
        """Load the data from a DataUrl object."""
        with h5py.File(data_url.file_path(), "r") as h5f:
            data = h5f[data_url.data_path()]
            if data_url.data_slice() is not None:
                return data[data_url.data_slice()].astype(np.float32)

    def get_slices_by_angle_step(self, angle_step=90) -> list:
        """
        Returns the slices of the image to be processed.
        """
        # Get all angles
        angles_dict = self.nxtomoscan.get_proj_angle_url()
        angles_dict = {clean_angle_key(k): v for k, v in angles_dict.items()}
        all_angles = np.array(list(angles_dict.keys()))

        # Determine all 90° targets within full range
        min_angle = np.min(all_angles)
        max_angle = np.max(all_angles)
        target_angles = np.arange(min_angle, max_angle + angle_step, angle_step)

        # For each target angle, find the closest available
        selected_angles = []
        used_indices = set()
        for target in target_angles:
            diffs = np.abs(all_angles - target)
            idx = np.argmin(diffs)
            if idx not in used_indices:  # avoid duplicates
                used_indices.add(idx)
                selected_angles.append(all_angles[idx])

        selected_slices = [
            self.get_proj_from_data_url(angles_dict[angle]) for angle in selected_angles
        ]
        return selected_angles, selected_slices

    def _bin_data(self, data: np.ndarray, binning: int) -> np.ndarray:
        """
        Bins a 2D array by the specified binning factor.
        If binning <= 1, returns the original data.
        """
        if binning <= 1:
            return data
        h, w = data.shape
        new_h = h // binning
        new_w = w // binning
        # Crop the image if necessary so dimensions are divisible by the binning factor.
        data_cropped = data[: new_h * binning, : new_w * binning]
        # Reshape and compute the mean over each bin.
        binned = data_cropped.reshape(new_h, binning, new_w, binning).mean(axis=(1, 3))
        return binned

    def _save_to_gallery(
        self,
        output_file_name: str,
        image: np.ndarray,
        bounds: tuple[float, float] | None = None,
    ) -> None:
        """
        Processes and saves two images to the gallery folder:
          - If the image is 3D with a singleton first dimension, reshapes it to 2D.
          - Normalizes the image to 8-bit grayscale using the provided bounds if available.
            If no bounds are provided, lower_bound defaults to the 0.01st percentile of finite values and upper_bound to the 99.99th percentile of
            finite values below 1e9. This prevents negative slices from being clipped to zero while still ignoring saturated pixels.
          - Applies binning based on gallery_output_binning.
          - Saves the full-size result as `<name>_large` and a 200x200 resized version with the standard name.
        """
        overwrite = self.gallery_overwrite
        binning = self.gallery_output_binning

        # Ensure the image is 2D. If it's 3D with a single channel, squeeze it.
        if image.ndim == 3 and image.shape[0] == 1:
            image = image.reshape(image.shape[1:])
        elif image.ndim != 2:
            raise ValueError(f"Only 2D grayscale images are handled. Got {image.shape}")

        # Check if bounds is a valid tuple; otherwise derive robust defaults.
        if not isinstance(bounds, tuple):
            lower_bound, upper_bound = _auto_intensity_bounds(image)
        else:
            lower_bound = float(bounds[0])
            upper_bound = float(bounds[1])

        # Apply clamping and normalization.
        image = np.clip(image, lower_bound, upper_bound)
        image = image - lower_bound
        if upper_bound != lower_bound:
            image = image * (255.0 / (upper_bound - lower_bound))

        # Apply binning if necessary.
        image = self._bin_data(data=image, binning=binning)

        # Convert the image to a PIL Image.
        output_path = Path(output_file_name)
        img = Image.fromarray(image.astype(np.uint8), mode="L")

        # Prepare save options (no compression).
        save_kwargs: dict[str, object] = {}
        gallery_output_format = getattr(self, "gallery_output_format", "png")
        if "png" in gallery_output_format:
            save_kwargs["compress_level"] = 0

        # Build both output paths (full-size gets a `_large` suffix).
        large_output_path = output_path.with_name(
            f"{output_path.stem}_large{output_path.suffix}"
        )
        if not overwrite and (output_path.exists() or large_output_path.exists()):
            raise OSError(f"File already exists ({output_path})")

        # Save full-size (`_large`) image.
        img.save(str(large_output_path), **save_kwargs)

        # Save resized 200x200 version with the standard name.
        img_small = img.resize((200, 200), resample=Image.Resampling.LANCZOS)
        img_small.save(str(output_path), **save_kwargs)


class BuildSlicesGallery(
    Task,
    input_names=["reconstructed_slice_path"],
    optional_input_names=[
        "bounds",
        "output_binning",
        "output_format",
        "overwrite",
    ],
    output_names=["processed_data_dir", "gallery_path", "gallery_image_path"],
):
    """Create two gallery images from a reconstructed slice (full-size suffixed `_large`, and a 200x200 resized version).
    The resized output file keeps the input basename, with `_bin{binning}` suffix and the chosen extension.
    """

    def run(self):
        """Read the slice, normalize/optionally bin, and save to <processed>/gallery."""
        fmt = str(self.get_input_value("output_format", "png")).lower()
        overwrite = bool(self.get_input_value("overwrite", True))
        binning = int(self.get_input_value("output_binning", 2))
        bounds = self.get_input_value("bounds", None)

        slice_path = Path(self.inputs.reconstructed_slice_path)
        if not slice_path.exists():
            raise FileNotFoundError(f"Reconstructed slice not found: {slice_path}")

        slices_dir = slice_path.parent
        processed_data_dir = slices_dir.parent
        gallery_dir = Path(processed_data_dir) / "gallery"
        os.makedirs(gallery_dir, exist_ok=True)

        arr = self._load_slice(slice_path)
        out_name = self.get_gallery_file_path(gallery_dir, slice_path, fmt)
        out_path = gallery_dir / out_name
        self._save_to_gallery(out_path, arr, bounds, overwrite, binning)

        self.outputs.processed_data_dir = str(processed_data_dir)
        self.outputs.gallery_path = str(gallery_dir)
        self.outputs.gallery_image_path = str(out_path)

    def get_gallery_dir(self, processed_data_dir: Path | str) -> str:
        """Return the fixed gallery directory path."""
        return str(Path(processed_data_dir) / "gallery")

    def get_gallery_file_path(self, gallery_dir, reconstructed_slice_path, fmt) -> str:
        filename = f"{reconstructed_slice_path.stem}.{fmt}"
        gallery_path = gallery_dir / filename
        return str(gallery_path)

    @staticmethod
    def _load_slice(img_path: Path) -> np.ndarray:
        """Load a 2D float32 slice (HDF5 at entry0000/reconstruction/results/data, EDF, or image)."""
        ext = img_path.suffix.lower()
        if ext in (".h5", ".hdf5"):
            with h5py.File(img_path, "r") as h5in:
                img = h5in["entry0000/reconstruction/results/data"][:]
            return np.squeeze(img).astype(np.float32)
        if ext == ".edf":
            try:
                import fabio  # type: ignore
            except Exception as exc:
                raise RuntimeError(
                    "EDF support requires 'fabio' (pip install fabio)."
                ) from exc
            return fabio.open(str(img_path)).data.astype(np.float32)
        with Image.open(img_path) as im:
            arr = np.array(im, dtype=np.float32)
        if arr.ndim == 3 and arr.shape[-1] in (3, 4):
            arr = arr[..., :3].mean(axis=-1)
        return arr

    @staticmethod
    def _bin(data: np.ndarray, binning: int) -> np.ndarray:
        """Downsample by integer binning using mean."""
        if binning <= 1:
            return data
        h, w = data.shape
        nh, nw = h // binning, w // binning
        data = data[: nh * binning, : nw * binning]
        return data.reshape(nh, binning, nw, binning).mean(axis=(1, 3))

    def _save_to_gallery(
        self,
        output_path: Path,
        image: np.ndarray,
        bounds: tuple[float, float] | None,
        overwrite: bool,
        binning: int,
    ) -> None:
        """Clamp to bounds, scale to 8-bit, apply optional binning, then save both large_ and resized 200x200 images."""
        if image.ndim != 2:
            raise ValueError(
                f"Only 2D grayscale images are handled. Got shape={image.shape}"
            )
        if not isinstance(bounds, tuple):
            lower, upper = _auto_intensity_bounds(image)
        else:
            lower, upper = float(bounds[0]), float(bounds[1])
        img = np.clip(image, lower, upper) - lower
        scale = 255.0 / (upper - lower) if upper != lower else 1.0
        img = (img * scale).astype(np.float32)
        if binning > 1:
            img = self._bin(img, binning)
        pil = Image.fromarray(np.clip(img, 0, 255).astype(np.uint8), mode="L")
        save_kwargs = {}
        if output_path.suffix.lower() == ".png":
            save_kwargs["compress_level"] = 0

        large_output_path = output_path.with_name(
            f"{output_path.stem}_large{output_path.suffix}"
        )
        if not overwrite and (output_path.exists() or large_output_path.exists()):
            raise OSError(f"File already exists ({output_path})")

        # Save full-size (`_large`) image.
        pil.save(str(large_output_path), **save_kwargs)

        # Save resized 200x200 version with the standard name.
        pil_small = pil.resize((200, 200), resample=Image.Resampling.LANCZOS)
        pil_small.save(str(output_path), **save_kwargs)
