#    Copyright 2023 SECTRA AB
#
#    Licensed under the Apache License, Version 2.0 (the "License");
#    you may not use this file except in compliance with the License.
#    You may obtain a copy of the License at
#
#        http://www.apache.org/licenses/LICENSE-2.0
#
#    Unless required by applicable law or agreed to in writing, software
#    distributed under the License is distributed on an "AS IS" BASIS,
#    WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
#    See the License for the specific language governing permissions and
#    limitations under the License.

"""Reader for bioformat supported files."""

import math
import os
from contextlib import contextmanager
from functools import cached_property, lru_cache
from pathlib import Path
from queue import Empty, SimpleQueue
from tempfile import TemporaryDirectory
from threading import Lock
from typing import Dict, Generator, List, Optional, Tuple, Type, Union

import jpype.imports  # Needed for loci import to work # noqa
import numpy as np
import ome_types
import scyjava
from jpype.types import JArray
from wsidicom.geometry import Region, Size, SizeMm

"""
Set version of bioformats jar to use with the environmental variable
"BIOFORMATS_VERSION". Note the bioformats jar either has a BSD or GPL-2
license.
"""
bioformats_version = os.getenv("BIOFORMATS_VERSION", "bsd:8.3.0")
scyjava.config.endpoints.append(f"ome:formats-{bioformats_version}")

if not scyjava.jvm_started():
    scyjava.start_jvm()

"""
Load bioformats modules using jpype and scyjava
"""
from loci.common.services import ServiceFactory  # type: ignore # noqa
from loci.formats import ImageReader, Memoizer  # type: ignore # noqa
from loci.formats.services import OMEXMLService  # type: ignore # noqa


class ReaderPool:
    def __init__(
        self,
        filepath: Path,
        max_readers: Optional[int] = None,
        cache_path: Optional[Union[Path, str]] = None,
    ):
        self._filepath = filepath
        if max_readers is None:
            cpu_count = os.cpu_count()
            if cpu_count is not None:
                max_readers = cpu_count
            else:
                max_readers = 1
        self._max_readers = max_readers
        self._current_count = 0
        if cache_path is None:
            self._tempdir = TemporaryDirectory()
            self._cache_path = Path(self._tempdir.name)
        else:
            self._tempdir = None
            self._cache_path = Path(cache_path)
        self._queue: SimpleQueue[Memoizer] = SimpleQueue()
        self._lock = Lock()

    @property
    def filepath(self) -> Path:
        return self._filepath

    @contextmanager
    def get_reader(self) -> Generator[Memoizer, None, None]:
        """Return a reader. Should be used as a context manager. Will block
        if no reader is available."""
        reader = self._get_reader()
        try:
            yield reader
        finally:
            self._return_reader(reader)

    @property
    def _is_full(self) -> bool:
        return self._current_count == self._max_readers

    @property
    def _is_empty(self) -> bool:
        return self._current_count == 0

    def _increment(self):
        if self._current_count < self._max_readers:
            self._current_count += 1
            return
        raise ValueError("Current count should never be higher than max readers.")

    def _decrement(self):
        if self._current_count > 0:
            self._current_count -= 1
            return
        raise ValueError("Current count should never be lower than 0.")

    def _get_reader(self) -> Memoizer:
        """Return a reader with no wait if one is available. Else if current
        reader count is less than maximum reader count return a new reader.
        Otherwise wait for an available readern."""
        try:
            return self._queue.get_nowait()
        except Empty:
            pass
        with self._lock:
            if not self._is_full:
                self._increment()
                return self._create_new_reader()

        return self._queue.get(block=True, timeout=None)

    def _create_new_reader(self) -> Memoizer:
        """Return a new reader."""
        # Create a reader using Memoizer to load file faster
        # See https://docs.openmicroscopy.org/bio-formats/6.11.0/developers/matlab-dev.html#reader-performance
        reader = Memoizer(ImageReader(), 0, self._cache_path)
        reader.setFlattenedResolutions(False)
        reader.setId(str(self._filepath))
        return reader

    def _return_reader(self, reader: Memoizer) -> None:
        """Release a used reader."""
        self._queue.put(reader)

    def close(self) -> None:
        """Close all readers and clean up cache directory if a temporary directory."""
        with self._lock:
            while not self._is_empty:
                reader = self._queue.get()
                reader.close()
                self._decrement()
            if self._tempdir is not None:
                self._tempdir.cleanup()


class BioformatsReader:
    def __init__(
        self,
        filepath: Path,
        max_readers: Optional[int] = None,
        cache_path: Optional[Union[Path, str]] = None,
    ):
        """Reader for image data and metadata from file using Bio-Formats api.

        Parameters
        ----------
        filepath: Path
            Path to file to open.
        max_readers: int
            Maximum number of readers to use. If not specified the number
            of cpus will be used.
        cache_path: Optional[Union[Path, str]] = None
            Path to store cache file quicker opening of new readers. If None
            a temporary directory will be used.
        """
        self._reader_pool = ReaderPool(filepath, max_readers, cache_path)

    @staticmethod
    def is_supported(filepath: Path) -> bool:
        try:
            reader = ImageReader()
            reader.setId(str(filepath))
            reader.close()
        except Exception:
            return False
        return True

    @property
    def filepath(self) -> Path:
        """Return filepath of the opened file."""
        return self._reader_pool.filepath

    @cached_property
    def metadata(self) -> ome_types.OME:
        """Return parsed metadata."""
        metadata = self._read_metadata()
        return ome_types.from_xml(str(metadata), parser="lxml")

    @property
    def images_count(self) -> int:
        """Return number of images in file."""
        return len(self.metadata.images)

    def image_name(self, image_index: int) -> Optional[str]:
        """Return name of image."""
        return self.metadata.images[image_index].name

    @lru_cache
    def dtype(self, image_index: int) -> np.dtype:
        """Return the numpy datatype for image in file."""
        NUMPY_DATA_TYPES: Dict[str, Type] = {
            "bit": np.bool_,
            "double-complex": np.cdouble,
            "complex": np.csingle,
            "double": np.float64,
            "float": np.float32,
            "int16": np.int16,
            "int32": np.int32,
            "int8": np.int8,
            "uint16": np.uint16,
            "uint32": np.uint32,
            "uint8": np.uint8,
        }
        if self.metadata.images[image_index].pixels.big_endian:
            byte_order = ">"
        else:
            byte_order = "<"
        data_type = self.metadata.images[image_index].pixels.type.value
        try:
            numpy_data_type = np.dtype(NUMPY_DATA_TYPES[data_type])
        except KeyError:
            raise ValueError(f"Unkown data type {data_type}")
        return numpy_data_type.newbyteorder(byte_order)

    @lru_cache
    def samples_per_pixel(self, image_index: int) -> int:
        """Return the samples per pixel for image in file."""
        pixels = self.metadata.images[image_index].pixels
        samples_per_pixel = pixels.channels[0].samples_per_pixel
        assert samples_per_pixel is not None
        return int(samples_per_pixel)

    @lru_cache
    def size(self, image_index: int, resolution_index: int = 0) -> Size:
        """Return the image size for image in file."""
        with self._reader_pool.get_reader() as reader:
            reader.setSeries(image_index)
            reader.setResolution(resolution_index)
            width = reader.getSizeX()
            height = reader.getSizeY()
            return Size(int(width), int(height))

    @lru_cache
    def pixel_spacing(
        self, image_index: int, resolution_index: int = 0
    ) -> Optional[SizeMm]:
        """Return the size of the pixels in mm/pixel for image in file."""
        pixels = self.metadata.images[image_index].pixels
        if pixels.physical_size_x is None or pixels.physical_size_y is None:
            return None
        scale = self._resolution_scales(image_index)[resolution_index]
        return (
            SizeMm(float(pixels.physical_size_x), float(pixels.physical_size_y))
            * scale
            / 1000
        )

    @lru_cache
    def is_interleaved(self, image_index: int) -> bool:
        """Return true if image data is interleaved."""
        interleaved = self.metadata.images[image_index].pixels.interleaved
        assert interleaved is not None
        return interleaved

    @lru_cache
    def pyramid_levels(self, image_index: int) -> Dict[Tuple[int, float, str], int]:
        """Return dictionary of dyadic scaling, focal plane, and optical path as key
        and resolution index as value for resolutions in image.

        Focal planes and optical paths are not implemented.
        """
        TOLERANCE = 1e-2
        float_pyramid_levels = (
            math.log2(scale) for scale in self._resolution_scales(image_index)
        )
        return {
            (round(float_level), 0.0, "0"): resolution_index
            for resolution_index, float_level in enumerate(float_pyramid_levels)
            if math.isclose(float_level, round(float_level), abs_tol=TOLERANCE)
        }

    @lru_cache
    def _resolution_scales(self, image_index: int) -> List[float]:
        """Return resolution scales for image as resolution width divided by image width."""
        with self._reader_pool.get_reader() as reader:
            reader.setSeries(image_index)
            reader.setResolution(0)
            base_width = reader.getSizeX()
            return [
                self._get_resolution_scale(reader, resolution_index, base_width)
                for resolution_index in range(reader.getResolutionCount())
            ]

    @staticmethod
    def _get_resolution_scale(
        reader: Memoizer, resolution_index: int, base_width: int
    ) -> float:
        """Return resolution scale for resolution as rounded int of resolution width
        divided by base widht."""
        reader.setResolution(resolution_index)
        width = reader.getSizeX()
        return base_width / width

    @contextmanager
    def read_image(
        self,
        image_index: int,
        resolution_index: int,
        region: Region,
        output_size: Optional[Size] = None,
        index: int = 0,
    ) -> Generator[np.ndarray, None, None]:
        """Read image data from file. Preferably used as a context manager.
        Data is returned as a memoryview that should be released after use.

        Parameters
        ----------
        image_index: int
            The image to read data from.
        region: Region
            The region to read data from.
        output_size: Optional[Size] = None
            Optional size to resize image data to. Must be equal to or larger than
            region size.
        index: int
            The index in image to read data from.

        Returns
        ----------
            Generator[np.ndarray, None, None]

        """
        try:
            raw_data = memoryview(
                self._read(
                    image_index,
                    resolution_index,
                    index,
                    region.start.x,
                    region.start.y,
                    region.size.width,
                    region.size.height,
                )
            )
        except Exception as exception:
            raise Exception(
                f"Failed to read image data from image {image_index}, resolution {resolution_index}, index {index} at region {region.box}."
            ) from exception
        try:
            data: np.ndarray = np.frombuffer(raw_data, self.dtype(image_index))
            if self.is_interleaved(image_index):
                data = data.reshape(
                    region.size.height,
                    region.size.width,
                    self.samples_per_pixel(image_index),
                )
            else:
                data = data.reshape(
                    self.samples_per_pixel(image_index),
                    region.size.height,
                    region.size.width,
                )
                data = np.moveaxis(data, 0, 2).copy()
            if output_size is not None and data.shape[0:2] != output_size.to_tuple():
                # Pad with zeros to get requsted output size.
                if not output_size.all_greater_than_or_equal(region.size):
                    raise ValueError(
                        "Output size should be equal to or larger than region size."
                    )
                padding_width = output_size.width - data.shape[0]
                padding_height = output_size.height - data.shape[1]
                data = np.pad(data, ((0, padding_width), (0, padding_height), (0, 0)))
            yield data
        finally:
            raw_data.release()

    def close(self) -> None:
        """Close reader pool."""
        self._reader_pool.close()

    def _read(
        self,
        image_index: int,
        resolution_index: int,
        index: int,
        start_x: int,
        start_y: int,
        end_x: int,
        end_y: int,
    ) -> JArray:
        """Read image data from file."""
        with self._reader_pool.get_reader() as reader:
            reader.setSeries(image_index)
            reader.setResolution(resolution_index)
            return reader.openBytes(index, start_x, start_y, end_x, end_y)

    def _read_metadata(self) -> str:
        """Read metadata from file."""
        with self._reader_pool.get_reader() as reader:
            metadata_store = reader.getMetadataStore()
            return str(metadata_store.dumpXML())
