import logging
import numpy as np

from .batch_filter import BatchFilter
from gunpowder.array import Array
from gunpowder.batch_request import BatchRequest
from gunpowder.batch import Batch
from gunpowder.coordinate import Coordinate

logger = logging.getLogger(__name__)


def seg_to_affgraph(seg, nhood):
    nhood = np.array(nhood)

    # constructs an affinity graph from a segmentation
    # assume affinity graph is represented as:
    # shape = (e, z, y, x)
    # nhood.shape = (edges, 3)
    shape = seg.shape
    nEdge = nhood.shape[0]
    dims = nhood.shape[1]
    aff = np.zeros((nEdge,) + shape, dtype=np.int32)

    if dims == 2:
        for e in range(nEdge):
            aff[
                e,
                max(0, -nhood[e, 0]) : min(shape[0], shape[0] - nhood[e, 0]),
                max(0, -nhood[e, 1]) : min(shape[1], shape[1] - nhood[e, 1]),
            ] = (
                (
                    seg[
                        max(0, -nhood[e, 0]) : min(shape[0], shape[0] - nhood[e, 0]),
                        max(0, -nhood[e, 1]) : min(shape[1], shape[1] - nhood[e, 1]),
                    ]
                    == seg[
                        max(0, nhood[e, 0]) : min(shape[0], shape[0] + nhood[e, 0]),
                        max(0, nhood[e, 1]) : min(shape[1], shape[1] + nhood[e, 1]),
                    ]
                )
                * (
                    seg[
                        max(0, -nhood[e, 0]) : min(shape[0], shape[0] - nhood[e, 0]),
                        max(0, -nhood[e, 1]) : min(shape[1], shape[1] - nhood[e, 1]),
                    ]
                    > 0
                )
                * (
                    seg[
                        max(0, nhood[e, 0]) : min(shape[0], shape[0] + nhood[e, 0]),
                        max(0, nhood[e, 1]) : min(shape[1], shape[1] + nhood[e, 1]),
                    ]
                    > 0
                )
            )

    elif dims == 3:
        for e in range(nEdge):
            aff[
                e,
                max(0, -nhood[e, 0]) : min(shape[0], shape[0] - nhood[e, 0]),
                max(0, -nhood[e, 1]) : min(shape[1], shape[1] - nhood[e, 1]),
                max(0, -nhood[e, 2]) : min(shape[2], shape[2] - nhood[e, 2]),
            ] = (
                (
                    seg[
                        max(0, -nhood[e, 0]) : min(shape[0], shape[0] - nhood[e, 0]),
                        max(0, -nhood[e, 1]) : min(shape[1], shape[1] - nhood[e, 1]),
                        max(0, -nhood[e, 2]) : min(shape[2], shape[2] - nhood[e, 2]),
                    ]
                    == seg[
                        max(0, nhood[e, 0]) : min(shape[0], shape[0] + nhood[e, 0]),
                        max(0, nhood[e, 1]) : min(shape[1], shape[1] + nhood[e, 1]),
                        max(0, nhood[e, 2]) : min(shape[2], shape[2] + nhood[e, 2]),
                    ]
                )
                * (
                    seg[
                        max(0, -nhood[e, 0]) : min(shape[0], shape[0] - nhood[e, 0]),
                        max(0, -nhood[e, 1]) : min(shape[1], shape[1] - nhood[e, 1]),
                        max(0, -nhood[e, 2]) : min(shape[2], shape[2] - nhood[e, 2]),
                    ]
                    > 0
                )
                * (
                    seg[
                        max(0, nhood[e, 0]) : min(shape[0], shape[0] + nhood[e, 0]),
                        max(0, nhood[e, 1]) : min(shape[1], shape[1] + nhood[e, 1]),
                        max(0, nhood[e, 2]) : min(shape[2], shape[2] + nhood[e, 2]),
                    ]
                    > 0
                )
            )

    else:
        raise RuntimeError(f"AddAffinities works only in 2 or 3 dimensions, not {dims}")

    return aff


class AddAffinities(BatchFilter):
    """Add an array with affinities for a given label array and neighborhood to
    the batch. Affinity values are created one for each voxel and entry in the
    neighborhood list, i.e., for each voxel and each neighbor of this voxel.
    Values are 1 iff both labels (of the voxel and the neighbor) are equal and
    non-zero.

    Args:

        affinity_neighborhood (``list`` of array-like):

            List of offsets for the affinities to consider for each voxel.

        labels (:class:`ArrayKey`):

            The array to read the labels from.

        affinities (:class:`ArrayKey`):

            The array to generate containing the affinities.

        labels_mask (:class:`ArrayKey`, optional):

            The array to use as a mask for ``labels``. Affinities connecting at
            least one masked out label will be masked out in
            ``affinities_mask``. If not given, ``affinities_mask`` will contain
            ones everywhere (if requested).

        unlabelled (:class:`ArrayKey`, optional):

            A binary array to indicate unlabelled areas with 0. Affinities from
            labelled to unlabelled voxels are set to 0, affinities between
            unlabelled voxels are masked out (they will not be used for
            training).

        affinities_mask (:class:`ArrayKey`, optional):

            The array to generate containing the affinitiy mask, as derived
            from parameter ``labels_mask``.
    """

    def __init__(
        self,
        affinity_neighborhood,
        labels,
        affinities,
        labels_mask=None,
        unlabelled=None,
        affinities_mask=None,
        dtype=np.uint8,
    ):
        self.affinity_neighborhood = np.array(affinity_neighborhood)
        self.labels = labels
        self.unlabelled = unlabelled
        self.labels_mask = labels_mask
        self.affinities = affinities
        self.affinities_mask = affinities_mask
        self.dtype = dtype

    def setup(self):
        assert self.labels in self.spec, (
            "Upstream does not provide %s needed by " "AddAffinities" % self.labels
        )

        voxel_size = self.spec[self.labels].voxel_size

        dims = self.affinity_neighborhood.shape[1]
        self.padding_neg = (
            Coordinate(
                min([0] + [a[d] for a in self.affinity_neighborhood])
                for d in range(dims)
            )
            * voxel_size
        )

        self.padding_pos = (
            Coordinate(
                max([0] + [a[d] for a in self.affinity_neighborhood])
                for d in range(dims)
            )
            * voxel_size
        )

        logger.debug("padding neg: " + str(self.padding_neg))
        logger.debug("padding pos: " + str(self.padding_pos))

        spec = self.spec[self.labels].copy()
        if spec.roi is not None:
            spec.roi = spec.roi.grow(self.padding_neg, -self.padding_pos)
        spec.dtype = self.dtype

        self.provides(self.affinities, spec)
        if self.affinities_mask:
            self.provides(self.affinities_mask, spec)
        self.enable_autoskip()

    def prepare(self, request):
        deps = BatchRequest()

        # grow labels ROI to accomodate padding
        labels_roi = request[self.affinities].roi.grow(
            -self.padding_neg, self.padding_pos
        )
        deps[self.labels] = request[self.affinities].copy()
        deps[self.labels].dtype = None
        deps[self.labels].roi = labels_roi

        if self.labels_mask:
            deps[self.labels_mask] = deps[self.labels].copy()
        if self.unlabelled:
            deps[self.unlabelled] = deps[self.labels].copy()

        return deps

    def process(self, batch, request):
        outputs = Batch()

        affinities_roi = request[self.affinities].roi

        logger.debug("computing ground-truth affinities from labels")

        affinities = seg_to_affgraph(
            batch.arrays[self.labels].data.astype(np.int32), self.affinity_neighborhood
        ).astype(self.dtype)

        # crop affinities to requested ROI
        offset = affinities_roi.offset
        shift = -offset - self.padding_neg
        crop_roi = affinities_roi.shift(shift)
        crop_roi /= self.spec[self.labels].voxel_size
        crop = crop_roi.get_bounding_box()

        logger.debug("cropping with " + str(crop))
        affinities = affinities[(slice(None),) + crop]

        spec = self.spec[self.affinities].copy()
        spec.roi = affinities_roi
        outputs.arrays[self.affinities] = Array(affinities, spec)

        if self.affinities_mask and self.affinities_mask in request:
            if self.labels_mask:
                logger.debug(
                    "computing ground-truth affinities mask from " "labels mask"
                )
                affinities_mask = seg_to_affgraph(
                    batch.arrays[self.labels_mask].data.astype(np.int32),
                    self.affinity_neighborhood,
                )
                affinities_mask = affinities_mask[(slice(None),) + crop]

            else:
                affinities_mask = np.ones_like(affinities)

            if self.unlabelled:
                # 1 for all affinities between unlabelled voxels
                unlabelled = 1 - batch.arrays[self.unlabelled].data
                unlabelled_mask = seg_to_affgraph(
                    unlabelled.astype(np.int32), self.affinity_neighborhood
                )
                unlabelled_mask = unlabelled_mask[(slice(None),) + crop]

                # 0 for all affinities between unlabelled voxels
                unlabelled_mask = 1 - unlabelled_mask

                # combine with mask
                affinities_mask = affinities_mask * unlabelled_mask

            affinities_mask = affinities_mask.astype(affinities.dtype)
            outputs.arrays[self.affinities_mask] = Array(affinities_mask, spec)

        else:
            if self.labels_mask is not None:
                logger.warning(
                    "GT labels does have a mask, but affinities "
                    "mask is not requested."
                )

        # Should probably have a better way of handling arbitrary batch attributes
        batch.affinity_neighborhood = self.affinity_neighborhood

        return outputs
