namespace.py

# Standard Library
import copy
import gzip
import json
from typing import IO

# Third Party Imports
import timy
from arango import ArangoError
from structlog import get_logger

# Local Imports
import bel.db.arangodb as arangodb
import bel.db.elasticsearch as elasticsearch
import bel.utils
from bel.Config import config

# import structlog
# import logging
# log = logging.getLogger(__name__)


log = get_logger()

terms_alias = "terms"


def load_terms(fo: IO, metadata: dict, forceupdate: bool):
    """Load terms into Elasticsearch and ArangoDB

    Forceupdate will create a new index in Elasticsearch regardless of whether
    an index with the resource version already exists.

    Args:
        fo: file obj - terminology file
        metadata: dict containing the metadata for terminology
        forceupdate: force full update - e.g. don't leave Elasticsearch indexes
            alone if their version ID matches
    """

    version = metadata["metadata"]["version"]

    # LOAD TERMS INTO Elasticsearch
    with timy.Timer("Load Terms") as timer:
        es = bel.db.elasticsearch.get_client()

        es_version = version.replace("T", "").replace("-", "").replace(":", "")
        index_prefix = f"terms_{metadata['metadata']['namespace'].lower()}"
        index_name = f"{index_prefix}_{es_version}"

        # Create index with mapping
        if not elasticsearch.index_exists(es, index_name):
            elasticsearch.create_terms_index(es, index_name)
        elif forceupdate:  # force an update to the index
            index_name += "_alt"
            elasticsearch.create_terms_index(es, index_name)
        else:
            return  # Skip loading if not forced and not a new namespace

        terms_iterator = terms_iterator_for_elasticsearch(fo, index_name)
        elasticsearch.bulk_load_docs(es, terms_iterator)

        # Remove old namespace index
        index_names = elasticsearch.get_all_index_names(es)
        for name in index_names:
            if name != index_name and index_prefix in name:
                elasticsearch.delete_index(es, name)

        # Add terms_alias to this index
        elasticsearch.add_index_alias(es, index_name, terms_alias)

        log.info(
            "Load namespace terms",
            elapsed=timer.elapsed,
            namespace=metadata["metadata"]["namespace"],
        )

    # LOAD EQUIVALENCES INTO ArangoDB
    with timy.Timer("Load Term Equivalences") as timer:
        arango_client = arangodb.get_client()
        if not arango_client:
            print("Cannot load terms without ArangoDB access")
            quit()
        belns_db = arangodb.get_belns_handle(arango_client)
        arangodb.batch_load_docs(
            belns_db, terms_iterator_for_arangodb(fo, version), on_duplicate="update"
        )

        log.info(
            "Loaded namespace equivalences",
            elapsed=timer.elapsed,
            namespace=metadata["metadata"]["namespace"],
        )

        # Clean up old entries
        remove_old_equivalence_edges = f"""
            FOR edge in equivalence_edges
                FILTER edge.source == "{metadata["metadata"]["namespace"]}"
                FILTER edge.version != "{version}"
                REMOVE edge IN equivalence_edges
        """
        remove_old_equivalence_nodes = f"""
            FOR node in equivalence_nodes
                FILTER node.source == "{metadata["metadata"]["namespace"]}"
                FILTER node.version != "{version}"
                REMOVE node IN equivalence_nodes
        """
        arangodb.aql_query(belns_db, remove_old_equivalence_edges)
        arangodb.aql_query(belns_db, remove_old_equivalence_nodes)

    # Add metadata to resource metadata collection
    metadata["_key"] = f"Namespace_{metadata['metadata']['namespace']}"
    try:
        belns_db.collection(arangodb.belns_metadata_name).insert(metadata)
    except ArangoError as ae:
        belns_db.collection(arangodb.belns_metadata_name).replace(metadata)


def terms_iterator_for_arangodb(fo, version):

    species_list = config["bel_resources"].get("species_list", [])

    fo.seek(0)
    with gzip.open(fo, "rt") as f:
        for line in f:
            term = json.loads(line)
            # skip if not term record (e.g. is a metadata record)
            if "term" not in term:
                continue
            term = term["term"]

            species_id = term.get("species_id", None)
            # Skip if species not listed in species_list
            if species_list and species_id and species_id not in species_list:
                continue

            source = term["namespace"]
            term_id = term["id"]
            term_key = arangodb.arango_id_to_key(term_id)

            (ns, val) = term_id.split(":", maxsplit=1)

            # Add primary ID node
            yield (
                arangodb.equiv_nodes_name,
                {
                    "_key": term_key,
                    "name": term_id,
                    "primary": True,
                    "namespace": ns,
                    "source": source,
                    "version": version,
                },
            )

            # Create Alt ID nodes/equivalences (to support other database equivalences using non-preferred Namespace IDs)
            if "alt_ids" in term:
                for alt_id in term["alt_ids"]:
                    # log.info(f'Added {alt_id} equivalence')
                    alt_id_key = arangodb.arango_id_to_key(alt_id)
                    yield (
                        arangodb.equiv_nodes_name,
                        {
                            "_key": alt_id_key,
                            "name": alt_id,
                            "namespace": ns,
                            "source": source,
                            "version": version,
                        },
                    )

                    arango_edge = {
                        "_from": f"{arangodb.equiv_nodes_name}/{term_key}",
                        "_to": f"{arangodb.equiv_nodes_name}/{alt_id_key}",
                        "_key": bel.utils._create_hash(f"{term_id}>>{alt_id}"),
                        "type": "equivalent_to",
                        "source": source,
                        "version": version,
                    }
                    yield (arangodb.equiv_edges_name, arango_edge)

            # Cross-DB equivalences
            if "equivalences" in term:
                for eqv in term["equivalences"]:
                    (ns, val) = eqv.split(":", maxsplit=1)
                    eqv_key = arangodb.arango_id_to_key(eqv)

                    yield (
                        arangodb.equiv_nodes_name,
                        {
                            "_key": eqv_key,
                            "name": eqv,
                            "namespace": ns,
                            "source": source,
                            "version": version,
                        },
                    )

                    arango_edge = {
                        "_from": f"{arangodb.equiv_nodes_name}/{term_key}",
                        "_to": f"{arangodb.equiv_nodes_name}/{eqv_key}",
                        "_key": bel.utils._create_hash(f"{term_id}>>{eqv}"),
                        "type": "equivalent_to",
                        "source": source,
                        "version": version,
                    }
                    yield (arangodb.equiv_edges_name, arango_edge)


def terms_iterator_for_elasticsearch(fo: IO, index_name: str):
    """Add index_name to term documents for bulk load"""

    species_list = config["bel_resources"].get("species_list", [])

    fo.seek(0)  # Seek back to beginning of file
    with gzip.open(fo, "rt") as f:
        for line in f:
            term = json.loads(line)
            # skip if not term record (e.g. is a metadata record)
            if "term" not in term:
                continue
            term = term["term"]

            # Filter species if enabled in config
            species_id = term.get("species_id", None)
            if species_list and species_id and species_id not in species_list:
                continue

            all_term_ids = set()
            for term_id in [term["id"]] + term.get("alt_ids", []):
                all_term_ids.add(term_id)
                all_term_ids.add(lowercase_term_id(term_id))

            term["alt_ids"] = copy.copy(list(all_term_ids))

            yield {
                "_op_type": "index",
                "_index": index_name,
                "_type": "term",
                "_id": term["id"],
                "_source": copy.deepcopy(term),
            }


def lowercase_term_id(term_id: str) -> str:
    """Lowercase the term value (not the namespace prefix)

    Args:
        term_id (str): term identifier with namespace prefix, e.g. MESH:Atherosclerosis

    Returns:
        str: lowercased, e.g. MESH:atherosclerosis
    """
    (ns, val) = term_id.split(":", maxsplit=1)
    term_id = f"{ns}:{val.lower()}"

    return term_id