# 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.
"""Query data features published as Ebrains datasets with AtlasConcepts"""

from collections import defaultdict
import re
from packaging.version import Version
from tempfile import NamedTemporaryFile
from typing import TYPE_CHECKING

from . import query
from ..features.dataset import ebrains as _ebrains
from ..commons import logger, siibra_tqdm
from ..features import anchor as _anchor
from ..retrieval import requests, datasets, cache
from ..core.concept import get_registry

if TYPE_CHECKING:
    from ..core.region import Region


class EbrainsFeatureQuery(query.LiveQuery, args=[], FeatureType=_ebrains.EbrainsDataFeature):
    # in EBRAINS knowledge graph prior to v3, versions were modelled
    # in dataset names. Typically found formats are (v1.0) and [rat, v2.1]
    VERSION_PATTERN = re.compile(r"^(.*?) *[\[\(][^v]*?(v[0-9].*?)[\]\)]")
    COMPACT_FEATURE_LIST = True

    # datasets whose name contains any of these strings will be ignored
    _BLACKLIST = {
        "Whole-brain parcellation of the Julich-Brain Cytoarchitectonic Atlas",
        "whole-brain collections of cytoarchitectonic probabilistic maps",
        "DiFuMo atlas",
        "Automated Anatomical Labeling (AAL1) atlas",
    }

    loader = requests.MultiSourcedRequest(
        requests=[
            requests.GitlabProxy(
                flavour=requests.GitlabProxyEnum.PARCELLATIONREGION_V1,
            ),
        ]
    )

    parcellation_ids = None

    def __init__(self, **kwargs):
        query.LiveQuery.__init__(self, **kwargs)

        if self.__class__.parcellation_ids is None:
            self.__class__.parcellation_ids = [
                dset.id
                for parc in get_registry("Parcellation")
                for dset in parc.datasets
                if isinstance(dset, datasets.EbrainsV3DatasetVersion)
            ]

    def query(self, region: "Region"):
        versioned_datasets = defaultdict(dict)
        invalid_species_datasets = {}
        results = self.loader.data.get("results", [])

        for r in siibra_tqdm(results, total=len(results), unit='dataset', desc=f'Matching datasets to "{region}"'):
            regionname = r.get("name", None)
            alias = r.get("alias", None)
            for ds_spec in r.get("datasets", []):
                ds_name = ds_spec.get("name")
                ds_id = ds_spec.get("@id")
                if "dataset" not in ds_id:
                    continue

                try:
                    ds_species = _anchor.Species.decode(ds_spec)
                except ValueError:
                    logger.debug(f"Cannot decode {ds_spec}")
                    invalid_species_datasets[ds_id] = ds_name
                    continue

                if self.COMPACT_FEATURE_LIST:
                    if any(ds_id.endswith(i) for i in self.parcellation_ids):
                        continue
                    if any(e.lower() in ds_name.lower() for e in self._BLACKLIST):
                        continue

                dset = _ebrains.EbrainsDataFeature(
                    dataset_version_id=ds_id,
                    anchor=_anchor.AnatomicalAnchor(
                        region=alias or regionname,
                        species=ds_species,
                    ),
                )
                if not dset.matches(region):
                    continue

                version_match = self.VERSION_PATTERN.search(ds_name)
                if version_match is None or not self.COMPACT_FEATURE_LIST:
                    yield dset
                else:  # store version, add only the latest version after the loop
                    name, version = version_match.groups()
                    versioned_datasets[name][version] = dset

        if len(invalid_species_datasets) > 0:
            with NamedTemporaryFile(mode="w", suffix=".txt", delete=False, encoding="utf-8") as f:
                for dsid, dsname in invalid_species_datasets.items():
                    f.write(f"{dsid} {dsname}\n")
                logger.warning(
                    f"{len(invalid_species_datasets)} datasets have been ignored, "
                    "because siibra could not decode their species. "
                    f"See {f.name}"
                )

        # if versioned datasets have been recorded, register only
        # the newest one with older ones linked as a version history.
        for name, dsets in versioned_datasets.items():
            try:  # if possible, sort by version tag
                sorted_versions = sorted(dsets.keys(), key=Version)
            except TypeError:  # else sort lexicographically
                sorted_versions = sorted(dsets.keys())

            # chain the dataset versions
            prev = None
            for version in sorted_versions:
                curr = dsets[version]
                curr.version = version
                if prev is not None:
                    curr._prev = prev
                    prev._next = curr
                prev = curr

            logger.debug(
                f"Registered only version {version} of {', '.join(sorted_versions)} for {name}. "
                f"Its version history is: {curr.version_history}"
            )
            yield curr


cache.Warmup.register_warmup_fn(cache.WarmupLevel.DATA)(lambda: EbrainsFeatureQuery.loader.data)
