import math
from functools import lru_cache
from typing import Optional

import numpy as np
import torch
import torch.nn.functional as F

from cellfinder.core.tools.array_operations import bin_mean_3d
from cellfinder.core.tools.geometry import make_sphere


class InvalidVolume(ValueError):
    """
    Raised when the volume passed to BallFilter is too small or does not meet
    requirements.
    """

    pass


@lru_cache(maxsize=50)
def get_kernel(ball_xy_size: int, ball_z_size: int) -> np.ndarray:
    """
    Create a spherical kernel.

    This is done by:
    1. Generating a binary sphere at a resolution *upscale_factor* larger
       than desired.
    2. Downscaling the binary sphere to get a 'fuzzy' sphere at the
       original intended scale
    """
    upscale_factor: int = 7
    upscaled_kernel_shape = (
        upscale_factor * ball_xy_size,
        upscale_factor * ball_xy_size,
        upscale_factor * ball_z_size,
    )
    upscaled_ball_centre_position = (
        np.floor(upscaled_kernel_shape[0] / 2),
        np.floor(upscaled_kernel_shape[1] / 2),
        np.floor(upscaled_kernel_shape[2] / 2),
    )
    upscaled_ball_radius = upscaled_kernel_shape[0] / 2.0

    sphere_kernel = make_sphere(
        upscaled_kernel_shape,
        upscaled_ball_radius,
        upscaled_ball_centre_position,
    )
    sphere_kernel = sphere_kernel.astype(np.float32)
    kernel = bin_mean_3d(
        sphere_kernel,
        bin_width=upscale_factor,
        bin_height=upscale_factor,
        bin_depth=upscale_factor,
    )

    assert (
        kernel.shape[2] == ball_z_size
    ), "Kernel z dimension should be {}, got {}".format(
        ball_z_size, kernel.shape[2]
    )

    return kernel


class BallFilter:
    """
    A 3D ball filter.

    This runs a spherical kernel across the 2d planar dimensions
    of a *ball_z_size* stack of planes, and marks pixels in the middle
    plane of the stack that have a high enough intensity over the
    the spherical kernel.

    Parameters
    ----------
    plane_height, plane_width : int
        Height/width of the planes.
    ball_xy_size : int
        Diameter of the spherical kernel in the x/y dimensions.
    ball_z_size : int
        Diameter of the spherical kernel in the z dimension.
        Equal to the number of planes stacked to filter
        the central plane of the stack.
    overlap_fraction : float
        The fraction of pixels within the spherical kernel that
        have to be over *threshold_value* for a pixel to be marked
        as having a high intensity.
    threshold_value : int
        Value above which an individual pixel is considered to have
        a high intensity.
    soma_centre_value : int
        Value used to mark pixels with a high enough intensity.
    tile_height, tile_width : int
        Width/height of individual tiles in the mask generated by
        2D filtering.
    dtype : str
        The data-type of the input planes and the type to use internally.
        E.g. "float32".
    batch_size: int
        The number of planes that will be typically passed in a single batch to
        `append`. This is only used to calculate `num_batches_before_ready`.
        Defaults to 1.
    torch_device: str
        The device on which the data and processing occurs on. Can be e.g.
        "cpu", "cuda" etc. Defaults to "cpu". Any data passed to the filter
        must be on this device. Returned data will also be on this device.
    use_mask : bool
        Whether tiling masks will be used in `append`. If False, tile masks
        won't be passed in and/or will be ignored. Defaults to True.
    """

    num_batches_before_ready: int
    """
    The number of batches of size `batch_size` passed to `append`
    before `ready` would return True.
    """

    # the inside brain tiled mask, if tiles are used (use_mask is True)
    inside_brain_tiles: Optional[torch.Tensor] = None

    def __init__(
        self,
        plane_height: int,
        plane_width: int,
        ball_xy_size: int,
        ball_z_size: int,
        overlap_fraction: float,
        threshold_value: int,
        soma_centre_value: int,
        tile_height: int,
        tile_width: int,
        dtype: str,
        batch_size: int = 1,
        torch_device: str = "cpu",
        use_mask: bool = True,
    ):
        self.ball_xy_size = ball_xy_size
        self.ball_z_size = ball_z_size
        self.overlap_fraction = overlap_fraction
        self.tile_step_height = tile_height
        self.tile_step_width = tile_width

        d1 = plane_height
        d2 = plane_width
        ball_xy_size = self.ball_xy_size
        if d1 < ball_xy_size or d2 < ball_xy_size:
            raise InvalidVolume(
                f"Invalid plane size {d1}x{d2}. Needs to be at least "
                f"{ball_xy_size} in each dimension"
            )

        self.THRESHOLD_VALUE = threshold_value
        self.SOMA_CENTRE_VALUE = soma_centre_value

        # kernel comes in as XYZ, change to ZYX so it aligns with data
        kernel = np.moveaxis(get_kernel(ball_xy_size, self.ball_z_size), 2, 0)
        self.overlap_threshold = np.sum(self.overlap_fraction * kernel)
        self.kernel_xy_size = kernel.shape[-2:]
        self.kernel_z_size = self.ball_z_size

        # convert to right type and pin for faster copying
        kernel = torch.from_numpy(kernel).type(getattr(torch, dtype))
        if torch_device != "cpu":
            # torch at one point threw a cuda memory error when splitting cells
            # on cpu due to pinning. It's best to only pin on using cuda
            kernel.pin_memory()
        # add 2 dimensions at the start so we have 11ZYX We need this shape in
        # the conv step
        self.kernel = (
            kernel.unsqueeze(0)
            .unsqueeze(0)
            .to(device=torch_device, non_blocking=True)
        )

        self.num_batches_before_ready = int(
            math.ceil(self.ball_z_size / batch_size)
        )
        # Stores the current planes that are being filtered. Start with no data
        self.volume = torch.empty(
            (0, plane_height, plane_width),
            dtype=getattr(torch, dtype),
        )
        # Index of the middle plane in the volume
        self.middle_z_idx = int(np.floor(self.ball_z_size / 2))

        if not use_mask:
            return
        # first axis is z
        n_vertical_tiles = int(np.ceil(plane_height / self.tile_step_height))
        n_horizontal_tiles = int(np.ceil(plane_width / self.tile_step_width))
        # Stores tile masks. We start with no data
        self.inside_brain_tiles = torch.empty(
            (
                0,
                n_vertical_tiles,
                n_horizontal_tiles,
            ),
            dtype=torch.bool,
        )

    @property
    def first_valid_plane(self) -> int:
        """
        The index in `self.volume` (or the planes passed in) that will be the
        first plane returned from `get_processed_planes`.

        E.g. if `ball_z_size` is 3, then this may return 1. Meaning the second
        plane passed to `append` (index 1), will be the first returned plane
        by `get_processed_planes`.
        """
        return int(math.floor(self.ball_z_size / 2))

    @property
    def remaining_planes(self) -> int:
        """
        The number of planes in `self.volume` (or the planes passed in) that
        will remain unprocessed after all the planes have been `walk`ed
        and `get_processed_planes` called.

        E.g. if `ball_z_size` is 3, then this may return 1. Meaning the last
        plane passed to `append`, will never be returned by
        `get_processed_planes` because the filter "center" never overlapped
        with it.
        """
        return self.ball_z_size - self.first_valid_plane - 1

    @property
    def ready(self) -> bool:
        """
        Return whether enough planes have been appended to run the filter
        using `walk`.
        """
        return self.volume.shape[0] >= self.kernel_z_size

    def append(
        self, planes: torch.Tensor, masks: Optional[torch.Tensor] = None
    ) -> None:
        """
        Add a new z-stack to the filter.

        Previous stacks passed to `append` are removed, except enough planes
        at the top of the previous z-stack to provide padding so we can filter
        starting from the first plane in `planes`. The first time we call
        `append`, `first_valid_plane` is the first plane to actually be
        filtered in the z-stack due to lack of padding.

        So make sure to call `walk`/`get_processed_planes` before calling
        `append` again.

        Parameters
        ----------
        planes : torch.Tensor
            The z-stack. There can be one or more planes in the stack, but it
            must have 3 dimensions. Input data is not modified.
        masks : torch.Tensor
            A z-stack tile mask, indicating for each tile whether it's in or
            outside the brain. If the latter it's excluded.

            If `use_mask` was True, this must be provided. If False, this
            parameter will be ignored.

            Input data is not modified.
        """
        if self.volume.shape[0]:
            if self.volume.shape[0] < self.kernel_z_size:
                num_remaining_with_padding = 0
            else:
                num_remaining = self.kernel_z_size - (self.middle_z_idx + 1)
                num_remaining_with_padding = num_remaining + self.middle_z_idx

            self.volume = torch.cat(
                [self.volume[-num_remaining_with_padding:, :, :], planes],
                dim=0,
            )

            if self.inside_brain_tiles is not None:
                self.inside_brain_tiles = torch.cat(
                    [
                        self.inside_brain_tiles[
                            -num_remaining_with_padding:, :, :
                        ],
                        masks,
                    ],
                    dim=0,
                )
        else:
            self.volume = planes.clone()
            if self.inside_brain_tiles is not None:
                self.inside_brain_tiles = masks.clone()

    def get_processed_planes(self) -> np.ndarray:
        """
        After passing enough planes to `append`, and after `walk`, this returns
        a copy of the processed planes as a numpy z-stack.

        It only starts returning planes corresponding to plane
        `first_valid_plane` relative to the first planes passed to `append`.
        E.g. if `ball_z_size` is 3 and `first_valid_plane` is 1, and you passed
        5 planes total to `append`, then this will have returned planes [1, 3].

        Notice the last plane was not included, because we return only "middle"
        planes - planes that can correspond to the center of a ball.
        """
        if not self.ready:
            raise TypeError("Not enough planes were appended")

        num_processed = self.volume.shape[0] - self.kernel_z_size + 1
        assert num_processed
        middle = self.middle_z_idx
        planes = (
            self.volume[middle : middle + num_processed, :, :]
            .cpu()
            .numpy()
            .copy()
        )
        return planes

    def walk(self) -> None:
        """
        Applies the filter to all the planes passed to `append`.

        May only be called if `ready` was True.

        You can get the processed planes from `get_processed_planes`.
        """
        if not self.ready:
            raise TypeError("Called walk before enough planes were appended")

        _walk(
            self.kernel_z_size,
            self.middle_z_idx,
            self.tile_step_height,
            self.tile_step_width,
            self.overlap_threshold,
            self.THRESHOLD_VALUE,
            self.SOMA_CENTRE_VALUE,
            self.kernel,
            self.volume,
            self.inside_brain_tiles,
        )


@torch.jit.script
def _walk(
    kernel_z_size: int,
    middle: int,
    tile_step_height: int,
    tile_step_width: int,
    overlap_threshold: float,
    threshold_value: int,
    soma_centre_value: int,
    kernel: torch.Tensor,
    volume: torch.Tensor,
    inside_brain_tiles: Optional[torch.Tensor],
):
    num_process = volume.shape[0] - kernel_z_size + 1
    height, width = volume.shape[1:]
    num_z = kernel.shape[2]

    # threshold volume so it's zero/one. And add two dims at start so
    # it's 11ZYX
    volume_tresh = (
        (volume >= threshold_value)
        .unsqueeze(0)
        .unsqueeze(0)
        .type(kernel.dtype)
    )

    # we do a plane at a time, volume: i:i+num_z, for plane i+middle
    for i in range(num_process):
        # spherical kernel is symmetric so convolution=correlation. Use
        # binary threshold mask over the kernel to sum the value of the
        # kernel at voxels that are bright
        overlap = F.conv3d(
            volume_tresh[:, :, i : i + num_z, :, :],
            kernel,
            stride=1,
            padding="valid",
        )[0, 0, 0, :, :]
        overlaps = overlap > overlap_threshold

        # only edit the volume that is valid - conv excludes edges so we
        # only edit the plane parts returned by conv3d
        height_valid, width_valid = overlaps.shape
        height_offset = (height - height_valid) // 2
        width_offset = (width - width_valid) // 2
        sub_volume = volume[
            i + middle,
            height_offset : height_offset + height_valid,
            width_offset : width_offset + width_valid,
        ]

        # do we use tile masks or just overlapping?
        if inside_brain_tiles is not None:
            # unfold tiles to cover the full voxel area each tile covers
            inside = (
                inside_brain_tiles[i + middle, :, :]
                .repeat_interleave(tile_step_height, dim=0)
                .repeat_interleave(tile_step_width, dim=1)
            )
            # again only process pixels in the valid area
            inside = inside[
                height_offset : height_offset + height_valid,
                width_offset : width_offset + width_valid,
            ]

            # must have enough ball overlap to be bright/tile is in brain
            sub_volume[torch.logical_and(overlaps, inside)] = soma_centre_value

        else:
            # must have enough ball overlap to be bright
            sub_volume[overlaps] = soma_centre_value
