# Copyright 2018-2025
# 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 os import path, makedirs
from typing import Dict, List

import numpy as np

from . import parcellationmap, volume as _volume
from .providers import provider
from ..commons import MapIndex, logger, connected_components, siibra_tqdm, resample_img_to_img
from ..locations import boundingbox
from ..retrieval.cache import CACHE
from ..retrieval.requests import HttpRequest, FileLoader


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.
        """
        from nibabel import Nifti1Image

        queryimg = queryvolume.fetch()
        assert isinstance(queryimg, Nifti1Image)
        assignments = []

        # resample query image into this image's voxel space, if required (nilearn checks)
        queryimg = resample_img_to_img(
            source_img=queryimg,
            target_img=Nifti1Image(
                np.zeros(self.shape), affine=self.affine, dtype=queryimg.dataobj.dtype
            )
        )
        self.space.get_template()
        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
