# Copyright 2018-2024
# Institute of Neuroscience and Medicine (INM-1), Forschungszentrum Jülich GmbH

# 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.
"""Represents lists of probabilistic brain region maps."""
from . import parcellationmap, volume as _volume

from .providers import provider
from ..commons import MapIndex, logger, connected_components, siibra_tqdm
from ..locations import boundingbox
from ..retrieval.cache import CACHE
from ..retrieval.requests import HttpRequest, FileLoader

from os import path, makedirs
from typing import Dict, List
from nilearn import image
import numpy as np


class SparseIndex:

    # Precomputed sparse indices are stored in an EBRAINS data proxy
    _BUCKET = "https://data-proxy.ebrains.eu/api/v1/buckets/reference-atlas-data/"
    _DATAPROXY_BASEURL = _BUCKET + "sparse-indices/sparse-indices-siibra_python_v1.0/"

    _SUFFIXES = {
        "probs": ".sparseindex_v1.probs.txt.gz",
        "bboxes": ".sparseindex_v1.bboxes.txt.gz",
        "voxels": ".sparseindex_v1.voxels.nii.gz"
    }

    def __init__(self):
        self.probs = []
        self.bboxes = []

        # these are initialized when adding the first volume, see below
        self.affine: np.ndarray = None
        self.shape = None
        self.voxels: np.ndarray = None

    def add_img(self, imgdata: np.ndarray, affine: np.ndarray):

        if self.num_volumes == 0:
            self.affine = affine
            self.shape = imgdata.shape
            self.voxels = np.zeros(imgdata.shape, dtype=np.int32) - 1
        else:
            if (imgdata.shape != self.shape) or ((affine - self.affine).sum() != 0):
                raise RuntimeError(
                    "Building sparse maps from volumes with different voxel spaces is not yet supported in siibra."
                )

        volume = self.num_volumes
        X, Y, Z = [v.astype("int32") for v in np.where(imgdata > 0)]
        for x, y, z, prob in zip(X, Y, Z, imgdata[X, Y, Z]):
            coord_id = self.voxels[x, y, z]
            if coord_id >= 0:
                # Coordinate already seen. Add observed value.
                assert volume not in self.probs[coord_id]
                assert len(self.probs) > coord_id
                self.probs[coord_id][volume] = prob
            else:
                # New coordinate. Append entry with observed value.
                coord_id = len(self.probs)
                self.voxels[x, y, z] = coord_id
                self.probs.append({volume: prob})

        self.bboxes.append(
            {
                "minpoint": (X.min(), Y.min(), Z.min()),
                "maxpoint": (X.max(), Y.max(), Z.max()),
            }
        )

    @property
    def num_volumes(self):
        return len(self.bboxes)

    def max(self):
        return self.voxels.max()

    def coords(self, volume: int):
        # Nx3 array with x/y/z coordinates of the N nonzero values of the given mapindex
        assert volume in range(self.num_volumes)
        coord_ids = [i for i, l in enumerate(self.probs) if volume in l]
        x0, y0, z0 = self.bboxes[volume]["minpoint"]
        x1, y1, z1 = self.bboxes[volume]["maxpoint"]
        return (
            np.array(
                np.where(
                    np.isin(
                        self.voxels[x0: x1 + 1, y0: y1 + 1, z0: z1 + 1],
                        coord_ids,
                    )
                )
            ).T
            + (x0, y0, z0)
        ).T

    def mapped_voxels(self, volume: int):
        # returns the x, y, and z coordinates of nonzero voxels for the map
        # with the given index, together with their corresponding values v.
        assert volume in range(self.num_volumes)
        x, y, z = [v.squeeze() for v in np.split(self.coords(volume), 3)]
        v = [self.probs[i][volume] for i in self.voxels[x, y, z]]
        return x, y, z, v

    @classmethod
    def load(cls, filepath_or_url: str) -> "SparseIndex":
        """
        Loads a precomputed SparseIndex to the memory.

        Parameters
        ----------
        filepath_or_url: str
            Path/url to the SparseIndex files
            (eg. https://url_to_files/basefilename):
            - basefilename.sparseindex_v1.probs.txt.gz
            - basefilename.sparseindex_v1.bboxes.txt.gz
            - basefilename.sparseindex_v1.voxels.nii.gz

        Returns
        -------
        SparseIndex
        """
        from gzip import decompress
        spindtxt_decoder = lambda b: decompress(b).decode('utf-8').strip().splitlines()

        probsfile = filepath_or_url + SparseIndex._SUFFIXES["probs"]
        bboxfile = filepath_or_url + SparseIndex._SUFFIXES["bboxes"]
        voxelfile = filepath_or_url + SparseIndex._SUFFIXES["voxels"]
        if all(path.isfile(f) for f in [probsfile, bboxfile, voxelfile]):
            request = FileLoader
        else:
            request = HttpRequest

        result = cls()

        voxels = request(voxelfile).get()
        result.voxels = np.asanyarray(voxels.dataobj)
        result.affine = voxels.affine
        result.shape = voxels.shape

        lines_probs = request(probsfile, func=spindtxt_decoder).get()
        for line in siibra_tqdm(
            lines_probs,
            total=len(lines_probs),
            desc="Loading sparse index",
            unit="voxels"
        ):
            fields = line.strip().split(" ")
            mapindices = list(map(int, fields[0::2]))
            values = list(map(float, fields[1::2]))
            D = dict(zip(mapindices, values))
            result.probs.append(D)

        lines_bboxes = request(bboxfile, func=spindtxt_decoder).get()
        for line in lines_bboxes:
            fields = line.strip().split(" ")
            result.bboxes.append({
                "minpoint": tuple(map(int, fields[:3])),
                "maxpoint": tuple(map(int, fields[3:])),
            })

        return result

    def save(self, base_filename: str, folder: str = ""):
        """
        Save SparseIndex (3x) files to under the folder `folder`
        with base_filename. If SparseIndex is not cached, siibra will first
        create it first.

        Parameters
        ----------
        base_filename: str
            The files that will be created as:
            - base_filename.sparseindex_v1.probs.txt.gz
            - base_filename.sparseindex_v1.bboxes.txt.gz
            - base_filename.sparseindex_v1.voxels.nii.gz

        folder: str, default=""
        """
        from nibabel import Nifti1Image
        import gzip

        fullpath = path.join(folder, base_filename)
        logger.info(f"Saving SparseIndex to '{base_filename}' with suffixes {SparseIndex._SUFFIXES}")

        if folder and not path.isdir(folder):
            makedirs(folder)

        Nifti1Image(self.voxels, self.affine).to_filename(
            fullpath + SparseIndex._SUFFIXES["voxels"]
        )
        with gzip.open(fullpath + SparseIndex._SUFFIXES["probs"], 'wt') as f:
            for D in self.probs:
                f.write(
                    "{}\n".format(
                        " ".join(f"{i} {p}" for i, p in D.items())
                    )
                )
        with gzip.open(fullpath + SparseIndex._SUFFIXES["bboxes"], "wt") as f:
            for bbox in self.bboxes:
                f.write(
                    "{} {}\n".format(
                        " ".join(map(str, bbox["minpoint"])),
                        " ".join(map(str, bbox["maxpoint"])),
                    )
                )
        logger.info(f"SparseIndex is saved to {fullpath}.")

    @classmethod
    def from_sparsemap(cls, sparsemap: "SparseMap") -> "SparseIndex":
        with provider.SubvolumeProvider.UseCaching():
            spind = cls()
            for img in siibra_tqdm(
                sparsemap.fetch_iter(), total=len(sparsemap), unit="maps",
                desc="Fetching volumetric maps and computing SparseIndex"
            ):
                spind.add_img(np.asanyarray(img.dataobj), img.affine)
        return spind


class SparseMap(parcellationmap.Map):
    """
    A sparse representation of list of statistical (e.g. probabilistic) brain
    region maps.

    It represents the 3D statistical maps of N brain regions by two data structures:

    1) 'spatial_index', a 3D volume where non-negative values represent unique indices into a list of region assignments
    2) 'probs', a list of region assignments where each entry is a dict

    More precisely, given ``i = sparse_index.voxels[x, y, z]`` we define that

    - if `i<0`, no brain region is assigned at this location
    - if `i>=0`, ``probs[i]`` defines the probabilities of brain regions.

    Each entry in probs is a dictionary that represents the region assignments for
    the unique voxel where ``spatial_index == i``. The assignment maps from a MapIndex
    to the actual (probability) value.
    """

    def __init__(
        self,
        identifier: str,
        name: str,
        space_spec: dict,
        parcellation_spec: dict,
        indices: Dict[str, MapIndex],
        volumes: list = [],
        shortname: str = "",
        description: str = "",
        modality: str = None,
        publications: list = [],
        datasets: list = [],
        prerelease: bool = False,
    ):
        parcellationmap.Map.__init__(
            self,
            identifier=identifier,
            name=name,
            space_spec=space_spec,
            parcellation_spec=parcellation_spec,
            indices=indices,
            shortname=shortname,
            description=description,
            modality=modality,
            publications=publications,
            datasets=datasets,
            volumes=volumes,
            prerelease=prerelease,
        )
        self._sparse_index_cached = None

    @property
    def _cache_prefix(self):
        return CACHE.build_filename(f"{self.parcellation.id}_{self.space.id}_{self.maptype}_{self.key.lower()}_index")

    @property
    def sparse_index(self):
        if self._sparse_index_cached is None:
            try:  # try loading from cache on disk
                spind = SparseIndex.load(self._cache_prefix)
            except Exception:
                spind = None
            if spind is None:  # try from precomputed source
                try:
                    logger.info("Downloading and loading precomputed SparseIndex...")
                    spind = SparseIndex.load(SparseIndex._DATAPROXY_BASEURL + self.key.lower())
                except Exception:
                    logger.error("Failed to download precomputed SparseIndex.", exc_info=1)
            if spind is None:  # Download each map and compute the SparseIndex
                spind = SparseIndex.from_sparsemap(self)
                spind.save(self._cache_prefix, folder=CACHE.folder)
            self._sparse_index_cached = spind
        assert self._sparse_index_cached.max() == len(self._sparse_index_cached.probs) - 1
        return self._sparse_index_cached

    @property
    def affine(self):
        return self.sparse_index.affine

    @property
    def shape(self):
        return self.sparse_index.shape

    def _read_voxel(self, x, y, z):
        spind = self.sparse_index
        vx = spind.voxels[x, y, z]
        if isinstance(vx, int):
            return list(
                (None, volume, None, value)
                for volume, value in spind.probs[vx].items()
            )
        else:
            return list(
                (pointindex, volume, None, value)
                for pointindex, voxel in enumerate(vx)
                for volume, value in spind.probs[voxel].items()
            )

    def _assign_volume(
        self,
        queryvolume: "_volume.Volume",
        minsize_voxel: int,
        lower_threshold: float,
        split_components: bool = True
    ) -> List[parcellationmap.AssignImageResult]:
        """
        Assign an image volume to this sparse map.

        Parameters
        -----------
        queryvolume: Volume
            the volume to be compared with maps
        minsize_voxel: int, default: 1
            Minimum voxel size of image components to be taken into account.
        lower_threshold: float, default: 0
            Lower threshold on values in the statistical map. Values smaller than
            this threshold will be excluded from the assignment computation.
        split_components: bool, default: True
            Whether to split the query volume into disjoint components.
        """
        queryimg = queryvolume.fetch()
        imgdata = np.asanyarray(queryimg.dataobj)
        imgaffine = queryimg.affine
        assignments = []

        # resample query image into this image's voxel space, if required
        if (imgaffine - self.affine).sum() == 0:
            querydata = imgdata.squeeze()
        else:
            if issubclass(imgdata.dtype.type, np.integer):
                interp = "nearest"
            else:
                interp = "linear"
            from nibabel import Nifti1Image
            queryimg = image.resample_img(
                Nifti1Image(imgdata, imgaffine),
                target_affine=self.affine,
                target_shape=self.shape,
                interpolation=interp,
            )
            querydata = np.asanyarray(queryimg.dataobj).squeeze()

        iter_func = connected_components if split_components \
            else lambda img: [(1, img)]

        for mode, modemask in iter_func(querydata):

            # determine bounding box of the mode
            XYZ2 = np.array(np.where(modemask)).T
            position = np.dot(self.affine, np.r_[XYZ2.mean(0), 1])[:3]
            if XYZ2.shape[0] <= minsize_voxel:
                continue
            X2, Y2, Z2 = [v.squeeze() for v in np.split(XYZ2, 3, axis=1)]

            bbox2 = boundingbox.BoundingBox(XYZ2.min(0), XYZ2.max(0) + 1, space=None)
            if bbox2.volume == 0:
                continue

            spind = self.sparse_index

            for volume in siibra_tqdm(
                range(len(self)),
                desc=f"Assigning structure #{mode} to {len(self)} sparse maps",
                total=len(self),
                unit=" map"
            ):
                bbox1 = boundingbox.BoundingBox(
                    self.sparse_index.bboxes[volume]["minpoint"],
                    self.sparse_index.bboxes[volume]["maxpoint"],
                    space=None,
                )
                if bbox1.intersection(bbox2) is None:
                    continue

                # compute union of voxel space bounding boxes
                bbox = bbox1.union(bbox2)
                bbshape = np.array(bbox.shape, dtype="int") + 1
                x0, y0, z0 = map(int, bbox.minpoint)

                # build flattened vector of map values
                v1 = np.zeros(np.prod(bbshape))
                XYZ1 = spind.coords(volume).T
                X1, Y1, Z1 = [v.squeeze() for v in np.split(XYZ1, 3, axis=1)]
                indices1 = np.ravel_multi_index(
                    (X1 - x0, Y1 - y0, Z1 - z0), bbshape
                )
                v1[indices1] = [spind.probs[i][volume] for i in spind.voxels[X1, Y1, Z1]]
                v1[v1 < lower_threshold] = 0

                # build flattened vector of input image mode
                v2 = np.zeros(np.prod(bbshape))
                indices2 = np.ravel_multi_index(
                    (X2 - x0, Y2 - y0, Z2 - z0), bbshape
                )
                v2[indices2] = querydata[X2, Y2, Z2]

                assert v1.shape == v2.shape

                intersection = np.sum(
                    (v1 > 0) & (v2 > 0)
                )  # np.minimum(v1, v2).sum()
                if intersection == 0:
                    continue
                iou = intersection / np.sum(
                    (v1 > 0) | (v2 > 0)
                )  # np.maximum(v1, v2).sum()

                v1d = v1 - v1.mean()
                v2d = v2 - v2.mean()
                rho = (
                    (v1d * v2d).sum()
                    / np.sqrt((v1d ** 2).sum())
                    / np.sqrt((v2d ** 2).sum())
                )

                maxval = v1.max()

                assignments.append(
                    parcellationmap.AssignImageResult(
                        input_structure=mode,
                        centroid=tuple(position.round(2)),
                        volume=volume,
                        fragment=None,
                        map_value=maxval,
                        intersection_over_union=iou,
                        intersection_over_first=intersection / (v1 > 0).sum(),
                        intersection_over_second=intersection / (v2 > 0).sum(),
                        correlation=rho,
                        weighted_mean_of_first=np.sum(v1 * v2) / np.sum(v2),
                        weighted_mean_of_second=np.sum(v1 * v2) / np.sum(v1)
                    )
                )

        return assignments
