# Copyright 2018-2021
# 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.
"""A specific mesh or 3D array."""
from .. import logger
from ..retrieval import requests
from ..core import space

import nibabel as nib
from abc import ABC, abstractmethod
from typing import List, Dict, Union, Set
import json


class ColorVolumeNotSupported(NotImplementedError):
    pass


class Volume:
    """
    A volume is a specific mesh or 3D array,
    which can be accessible via multiple providers in different formats.
    """

    IMAGE_FORMATS = [
        "nii",
        "zip/nii",
        "neuroglancer/precomputed"
    ]

    MESH_FORMATS = [
        "neuroglancer/precompmesh",
        "neuroglancer/precompmesh/surface",
        "gii-mesh",
        "gii-label"
    ]

    SUPPORTED_FORMATS = IMAGE_FORMATS + MESH_FORMATS

    def __init__(
        self,
        space_spec: dict,
        providers: List['VolumeProvider'],
        name: str = "",
        variant: str = None,
    ):
        self._name_cached = name  # see lazy implementation below
        self._space_spec = space_spec
        self.variant = variant
        self._providers: Dict[str, 'VolumeProvider'] = {}
        for provider in providers:
            srctype = provider.srctype
            assert srctype not in self._providers
            self._providers[srctype] = provider
        if len(self._providers) == 0:
            logger.debug(f"No provider for volume {self}")

    @property
    def name(self):
        """ allows derived classes to implemente a lazy name specification."""
        return self._name_cached

    @property
    def providers(self):
        return {
            srctype: prov._url for srctype, prov in self._providers.items()
        }

    @property
    def formats(self) -> Set[str]:
        result = set()
        for fmt in self._providers:
            result.add(fmt)
            result.add('mesh' if fmt in self.MESH_FORMATS else 'image')
        return result

    @property
    def provides_mesh(self):
        return any(f in self.MESH_FORMATS for f in self.formats)

    @property
    def provides_image(self):
        return any(f in self.IMAGE_FORMATS for f in self.formats)

    @property
    def fragments(self) -> Dict[str, List[str]]:
        result = {}
        for srctype, p in self._providers.items():
            t = 'mesh' if srctype in self.MESH_FORMATS else 'image'
            for fragment_name in p.fragments:
                if t in result:
                    result[t].append(fragment_name)
                else:
                    result[t] = [fragment_name]
        return result

    @property
    def space(self):
        for key in ["@id", "name"]:
            if key in self._space_spec:
                return space.Space.get_instance(self._space_spec[key])
        return space.Space(None, "Unspecified space", species=space.Species.UNSPECIFIED_SPECIES)

    def __str__(self):
        if self.space is None:
            return f"{self.__class__.__name__} '{self.name}'"
        else:
            return f"{self.__class__.__name__} '{self.name}' in space '{self.space.name}'"

    def __repr__(self):
        return self.__str__()

    def fetch(
        self,
        format: str = None,
        **kwargs
    ):
        """
        Fetch a volumetric or surface representation from one of the providers.

        Parameters
        ----------
        format: str
            Requested format. Per default, several formats are tried,
            starting with volumetric formats. You may explicitly specify:
            - 'surface' or 'mesh' to fetch a surface format
            - 'volumetric' or 'voxel' to fetch a volumetric format
            - supported format types, see :data:`Volume.SUPPORTED_FORMATS`. This includes:
                :class:`'nii'<siibra.volumes.nifti.NiftiProvider()>`,
                :class:`'zip/nii'<siibra.volumes.nifti.ZipContainedNiftiProvider()>`,
                :class:`'neuroglancer/precomputed'<siibra.volumes.neuroglancer.NeuroglancerProvider()>`,
                :class:`'gii-mesh'<siibra.volumes.gifti.GiftiMesh()>`,
                :class:`'neuroglancer/precompmesh'<siibra.volumes.neuroglancer.NeuroglancerMesh()>`,
                :class:`'gii-label'<siibra.volumes.gifti.GiftiSurfaceLabeling()>`
        """

        if format is None:
            requested_formats = self.SUPPORTED_FORMATS
        elif format == 'mesh':
            requested_formats = self.MESH_FORMATS
        elif format == 'image':
            requested_formats = set(self.SUPPORTED_FORMATS) - set(self.MESH_FORMATS)
        elif format in self.SUPPORTED_FORMATS:
            requested_formats = [format]
        else:
            raise ValueError(f"Invalid format requested: {format}")

        for fmt in requested_formats:
            if fmt in self.formats:
                try:
                    if fmt == "gii-label":
                        tpl = self.space.get_template(variant=kwargs.get('variant'), format=format)
                        mesh = tpl.fetch(**kwargs)
                        labels = self._providers[fmt].fetch(**kwargs)
                        return dict(**mesh, **labels)
                    else:
                        return self._providers[fmt].fetch(**kwargs)
                except requests.SiibraHttpRequestError as e:
                    logger.error(f"Cannot access {self._providers[fmt]}")
                    print(str(e))
                    continue
        return None


class Subvolume(Volume):
    """
    Wrapper class for exposing a z level of a 4D volume to be used like a 3D volume.
    """

    def __init__(self, parent_volume: Volume, z: int):
        Volume.__init__(
            self,
            space_spec=parent_volume._space_spec,
            providers=[
                SubvolumeProvider(p, z=z)
                for p in parent_volume._providers.values()
            ]
        )


# TODO add mesh primitive. Check nibabel implementation? Use trimesh? Do we want to add yet another dependency?
VolumeData = Union[nib.Nifti1Image, Dict]


class VolumeProvider(ABC):

    def __init_subclass__(cls, srctype: str) -> None:
        cls.srctype = srctype
        return super().__init_subclass__()

    @property
    def fragments(self) -> List[str]:
        return []

    @abstractmethod
    def fetch(self, *args, **kwargs) -> VolumeData:
        raise NotImplementedError

    @property
    @abstractmethod
    def _url(self) -> Union[str, Dict[str, str]]:
        """
        This is needed to provide urls to applications that can utilise such resources directly.
        e.g. siibra-api
        """
        return {}


class SubvolumeProvider(VolumeProvider, srctype="subvolume"):
    """
    This provider wraps around an existing volume provider,
    but is preconfigured to always fetch a fixed subvolume.
    The primary use is to provide a fixed z coordinate
    of a 4D volume provider as a 3D volume under the
    interface of a normal volume provider.
    """

    _USE_CACHING = False
    _FETCHED_VOLUMES = {}

    class UseCaching:
        def __enter__(self):
            SubvolumeProvider._USE_CACHING = True

        def __exit__(self, et, ev, tb):
            SubvolumeProvider._USE_CACHING = False
            SubvolumeProvider._FETCHED_VOLUMES = {}

    def __init__(self, parent_provider: VolumeProvider, z: int):
        VolumeProvider.__init__(self)
        self.provider = parent_provider
        self.srctype = parent_provider.srctype
        self.z = z

    def fetch(self, **kwargs):
        # activate caching at the caller using "with SubvolumeProvider.UseCaching():""
        if self.__class__._USE_CACHING:
            data_key = json.dumps(self.provider._url, sort_keys=True) \
                + json.dumps(kwargs, sort_keys=True)
            if data_key not in self.__class__._FETCHED_VOLUMES:
                vol = self.provider.fetch(**kwargs)
                self.__class__._FETCHED_VOLUMES[data_key] = vol
            vol = self.__class__._FETCHED_VOLUMES[data_key]
        else:
            vol = self.provider.fetch(**kwargs)
        return vol.slicer[:, :, :, self.z]

    def __getattr__(self, attr):
        return self.provider.__getattribute__(attr)

    @property
    def _url(self) -> Union[str, Dict[str, str]]:
        return super()._url
