"""Biomappings resources."""

from __future__ import annotations

import csv
import itertools as itt
import logging
from collections import defaultdict
from collections.abc import Iterable, Mapping, Sequence
from pathlib import Path
from typing import TYPE_CHECKING, Callable, NamedTuple

from bioregistry import NormalizedNamableReference, NormalizedNamedReference
from curies import NamableReference, NamedReference
from pydantic import BaseModel, ConfigDict, Field
from tqdm.auto import tqdm
from typing_extensions import Literal, Self

from biomappings.utils import (
    CURATORS_PATH,
    NEGATIVES_SSSOM_PATH,
    POSITIVES_SSSOM_PATH,
    PREDICTIONS_SSSOM_PATH,
    UNSURE_SSSOM_PATH,
    get_canonical_tuple,
)

if TYPE_CHECKING:
    import semra

__all__ = [
    "SemanticMapping",
    "append_false_mappings",
    "append_prediction_tuples",
    "append_predictions",
    "append_true_mapping_tuples",
    "append_true_mappings",
    "append_unsure_mappings",
    "filter_predictions",
    "get_curated_filter",
    "load_curators",
    "load_false_mappings",
    "load_mappings",
    "load_mappings_subset",
    "load_predictions",
    "load_unsure",
    "prediction_tuples_from_semra",
    "remove_mappings",
    "write_false_mappings",
    "write_predictions",
    "write_true_mappings",
    "write_unsure_mappings",
]

logger = logging.getLogger(__name__)


class _CuratedTuple(NamedTuple):
    """A tuple for writing manual curations to SSSOM TSV."""

    subject_id: str
    subject_label: str
    predicate_id: str
    object_id: str
    object_label: str
    mapping_justification: str
    author_id: str
    mapping_tool: str
    predicate_modifier: str


class _PredictedTuple(NamedTuple):
    """A tuple for writing predictions to SSSOM TSV."""

    subject_id: str
    subject_label: str
    predicate_id: str
    object_id: str
    object_label: str
    mapping_justification: str
    confidence: str
    mapping_tool: str


class SemanticMapping(BaseModel):
    """A class for semantic mappings."""

    model_config = ConfigDict(frozen=True)

    subject: NamableReference
    predicate: NamableReference
    object: NamableReference
    mapping_justification: NamableReference = Field(
        ...,
        description="""\
        A `semapv <https://bioregistry.io/registry/semapv>`_ term describing the mapping type.

        These are relatively high level, and can be any child of ``semapv:Matching``, including:

        1. ``semapv:LexicalMatching``
        2. ``semapv:LogicalReasoning``
        """,
    )
    author: NamableReference | None = None
    mapping_tool: str | None = Field(
        None,
        description="""\
            The script or process that generated this mapping.

            Most of these scripts are in https://github.com/biopragmatics/biomappings/tree/master/scripts,
            or can be based off of them.
        """,
    )
    predicate_modifier: Literal["Not"] | None = None
    confidence: float | None = Field(
        None,
        description="""\
    An assessment of the confidence of the mapping, reported by the method used to generate it.

    This means that confidence values aren't generally comparable, though they should follow
    the rough standard that closer to 1 is more confident and closer to 0 is less confident.

    Most of the lexical mappings already in Biomappings were generated with Gilda.
    Depending on the script, the score therefore refers to either:

    1. The Gilda match score, inspired by https://aclanthology.org/W15-3801/. Section 5.2 of the
       `supplementary material for the Gilda paper <https://doi.org/10.1093/bioadv/vbac034>`_
       describes this score in detail, where 1.0 is best and 0 is worst.
       https://github.com/biopragmatics/biomappings/blob/master/scripts/generate_agrovoc_mappings.py
       is an example that uses this variant.
    2. A high-level estimation of the precision of the scores generated by the given script.
       For example, the CL-MeSH mappings were estimated to be 90% correct, so all the mappings
       generated by https://github.com/biopragmatics/biomappings/blob/master/scripts/generate_cl_mesh_mappings.py
       are marked with 0.9 as its score.

    However, other variants are possible. For example, this confidence could reflect the loss function
    if a knowledge graph embedding model was used ot generate a mapping prediction.
    """,
    )

    def flip(self) -> Self:
        """Flip the mapping, if it's an exact match."""
        if self.predicate.curie != "skos:exactMatch":
            raise NotImplementedError
        return self.model_copy(
            update={
                "subject": self.object,
                "object": self.subject,
            }
        )

    def as_curated_row(self) -> _CuratedTuple:
        """Get a row for a curated SSSOM TSV."""
        return _CuratedTuple(
            self.subject.curie,
            self.subject.name or "",
            self.predicate.curie,
            self.object.curie,
            self.object.name or "",
            self.mapping_justification.curie,
            self.author.curie if self.author is not None else "",
            self.mapping_tool or "",
            str(self.predicate_modifier or ""),
        )

    @classmethod
    def from_row(cls, row: dict[str, str], reference_cls: type[NamableReference]) -> Self:
        """Construct from a dictionary."""
        for k in ["subject", "object", "predicate", "author"]:
            k1 = f"{k}_id"
            v1 = row.pop(k1, None)
            if v1 is None:
                continue
            if not isinstance(v1, str):
                raise TypeError
            k2 = f"{k}_label"
            # type ignore because properly typing the dict i/o is a pain
            row[k] = reference_cls.from_curie(v1, name=row.pop(k2, None))  # type:ignore[assignment]
        for k in ["mapping_justification"]:
            row[k] = reference_cls.from_curie(row.pop(k))  # type:ignore[assignment]
        return cls.model_validate(row)

    def as_predicted_row(self) -> _PredictedTuple:
        """Get a row for a predicted SSSOM TSV."""
        return _PredictedTuple(
            self.subject.curie,
            self.subject.name or "",
            self.predicate.curie,
            self.object.curie,
            self.object.name or "",
            self.mapping_justification.curie,
            str(self.confidence) if self.confidence is not None else "",
            self.mapping_tool or "",
        )


def _load_table(path: str | Path, *, standardize: bool) -> list[SemanticMapping]:
    reference_cls: type[NamableReference]
    if standardize:
        reference_cls = NormalizedNamableReference
    else:
        reference_cls = NamableReference

    with Path(path).expanduser().resolve().open() as file:
        return [
            SemanticMapping.from_row(
                {
                    key: value_stripped
                    for key, value in record.items()
                    if value and (value_stripped := value.strip()) and value_stripped != "."
                },
                reference_cls,
            )
            for record in csv.DictReader(file, delimiter="\t")
        ]


def _write_helper(
    mappings: Iterable[SemanticMapping],
    path: str | Path,
    mode: Literal["w", "a"],
    t: Literal["curated", "predicted"],
) -> None:
    mappings = sorted(set(mappings), key=mapping_sort_key)
    header: Sequence[str]
    to_row: Callable[[SemanticMapping], _PredictedTuple | _CuratedTuple]
    if t == "curated":
        header = _CuratedTuple._fields
        to_row = SemanticMapping.as_curated_row
    else:
        header = _PredictedTuple._fields
        to_row = SemanticMapping.as_predicted_row
    path = Path(path).expanduser().resolve()
    with path.open(mode) as file:
        if mode == "w":
            print(*header, sep="\t", file=file)
        for mapping in mappings:
            print(*to_row(mapping), sep="\t", file=file)


def mapping_sort_key(mapping: SemanticMapping) -> tuple[str, ...]:
    """Return a tuple for sorting mapping dictionaries."""
    return (
        mapping.subject.curie,
        mapping.predicate.curie,
        mapping.object.curie,
        mapping.mapping_justification.curie,
        mapping.mapping_tool or "",
    )


def load_mappings(
    *, path: str | Path | None = None, standardize: bool = False
) -> list[SemanticMapping]:
    """Load the mappings table."""
    return _load_table(path or POSITIVES_SSSOM_PATH, standardize=standardize)


def load_mappings_subset(source: str, target: str) -> Mapping[str, str]:
    """Get a dictionary of 1-1 mappings from the source prefix to the target prefix."""
    # TODO replace with SeMRA functionality?
    return {
        mapping.subject.identifier: mapping.object.identifier
        for mapping in load_mappings()
        if mapping.subject.prefix == source and mapping.object.prefix == target
    }


def append_true_mappings(
    mappings: Iterable[SemanticMapping],
    *,
    sort: bool = True,
    path: Path | None = None,
    standardize: bool = False,
) -> None:
    """Append new lines to the mappings table."""
    if path is None:
        path = POSITIVES_SSSOM_PATH
    _write_helper(mappings, path=path, mode="a", t="curated")
    if sort:
        lint_true_mappings(path=path, standardize=standardize)


def append_true_mapping_tuples(mappings: Iterable[SemanticMapping]) -> None:
    """Append new lines to the mappings table."""
    append_true_mappings(mappings)


def write_true_mappings(mappings: Iterable[SemanticMapping], *, path: Path | None = None) -> None:
    """Write mappigns to the true mappings file."""
    _write_helper(mappings, path=path or POSITIVES_SSSOM_PATH, mode="w", t="curated")


def lint_true_mappings(*, path: Path | None = None, standardize: bool) -> None:
    """Lint the true mappings file."""
    _lint_curated_mappings(path=path or POSITIVES_SSSOM_PATH, standardize=standardize)


def _lint_curated_mappings(path: Path, *, standardize: bool) -> None:
    """Lint the true mappings file."""
    mapping_list = _load_table(path, standardize=standardize)
    mappings = _remove_redundant(mapping_list)
    _write_helper(mappings, path, mode="w", t="curated")


def load_false_mappings(
    *, path: Path | None = None, standardize: bool = False
) -> list[SemanticMapping]:
    """Load the false mappings table."""
    return _load_table(path or NEGATIVES_SSSOM_PATH, standardize=standardize)


def append_false_mappings(
    mappings: Iterable[SemanticMapping],
    *,
    sort: bool = True,
    path: Path | None = None,
    standardize: bool = False,
) -> None:
    """Append new lines to the false mappings table."""
    if path is None:
        path = NEGATIVES_SSSOM_PATH
    _write_helper(mappings=mappings, path=path, mode="a", t="curated")
    if sort:
        lint_false_mappings(path=path, standardize=standardize)


def write_false_mappings(mappings: Iterable[SemanticMapping], *, path: Path | None = None) -> None:
    """Write mappings to the false mappings file."""
    _write_helper(mappings, path or NEGATIVES_SSSOM_PATH, mode="w", t="curated")


def lint_false_mappings(*, path: Path | None = None, standardize: bool) -> None:
    """Lint the false mappings file."""
    _lint_curated_mappings(path=path or NEGATIVES_SSSOM_PATH, standardize=standardize)


def load_unsure(*, path: Path | None = None, standardize: bool = False) -> list[SemanticMapping]:
    """Load the unsure table."""
    return _load_table(path or UNSURE_SSSOM_PATH, standardize=standardize)


def append_unsure_mappings(
    mappings: Iterable[SemanticMapping],
    *,
    sort: bool = True,
    path: Path | None = None,
    standardize: bool = False,
) -> None:
    """Append new lines to the "unsure" mappings table."""
    if path is None:
        path = UNSURE_SSSOM_PATH
    _write_helper(mappings, path=path, mode="a", t="curated")
    if sort:
        lint_unsure_mappings(path=path, standardize=standardize)


def write_unsure_mappings(mappings: Iterable[SemanticMapping], *, path: Path | None = None) -> None:
    """Write mappings to the unsure mappings file."""
    _write_helper(mappings, path or UNSURE_SSSOM_PATH, mode="w", t="curated")


def lint_unsure_mappings(*, standardize: bool, path: Path | None = None) -> None:
    """Lint the unsure mappings file."""
    _lint_curated_mappings(path=path or UNSURE_SSSOM_PATH, standardize=standardize)


def load_predictions(
    *, path: str | Path | None = None, standardize: bool = False
) -> list[SemanticMapping]:
    """Load the predictions table."""
    return _load_table(path or PREDICTIONS_SSSOM_PATH, standardize=standardize)


def write_predictions(mappings: Iterable[SemanticMapping], *, path: Path | None = None) -> None:
    """Write new content to the predictions table."""
    _write_helper(mappings, path or PREDICTIONS_SSSOM_PATH, mode="w", t="predicted")


def append_prediction_tuples(
    prediction_tuples: Iterable[SemanticMapping],
    *,
    deduplicate: bool = True,
    sort: bool = True,
    path: Path | None = None,
    standardize: bool = False,
) -> None:
    """Append new lines to the predictions table that come as canonical tuples."""
    append_predictions(
        prediction_tuples, deduplicate=deduplicate, sort=sort, path=path, standardize=standardize
    )


def append_predictions(
    mappings: Iterable[SemanticMapping],
    *,
    deduplicate: bool = True,
    sort: bool = True,
    path: Path | None = None,
    standardize: bool,
) -> None:
    """Append new lines to the predictions table."""
    if deduplicate:
        existing_mappings = {
            get_canonical_tuple(existing_mapping)
            for existing_mapping in itt.chain(
                load_mappings(),
                load_false_mappings(),
                load_unsure(),
                load_predictions(),
            )
        }
        mappings = (
            mapping for mapping in mappings if get_canonical_tuple(mapping) not in existing_mappings
        )

    if path is None:
        path = PREDICTIONS_SSSOM_PATH
    _write_helper(mappings, path, mode="a", t="predicted")
    if sort:
        lint_predictions(path=path, standardize=standardize)


def lint_predictions(
    *,
    path: Path | None = None,
    additional_curated_mappings: Iterable[SemanticMapping] | None = None,
    standardize: bool,
) -> None:
    """Lint the predictions file.

    1. Make sure there are no redundant rows
    2. Make sure no rows in predictions match a row in the curated files
    3. Make sure it's sorted

    :param path: The path to the predicted mappings
    :param additional_curated_mappings: A list of additional mappings
    """
    mappings = remove_mappings(
        load_predictions(path=path, standardize=standardize),
        itt.chain(
            load_mappings(standardize=standardize),
            load_false_mappings(standardize=standardize),
            load_unsure(standardize=standardize),
            additional_curated_mappings or [],
        ),
    )
    mappings = _remove_redundant(mappings)
    mappings = sorted(mappings, key=mapping_sort_key)
    write_predictions(mappings, path=path)


def remove_mappings(
    mappings: Iterable[SemanticMapping], mappings_to_remove: Iterable[SemanticMapping]
) -> Iterable[SemanticMapping]:
    """Remove the first set of mappings from the second."""
    skip_tuples = {get_canonical_tuple(mtr) for mtr in mappings_to_remove}
    return (mapping for mapping in mappings if get_canonical_tuple(mapping) not in skip_tuples)


def _remove_redundant(mappings: Iterable[SemanticMapping]) -> Iterable[SemanticMapping]:
    dd = defaultdict(list)
    for mapping in mappings:
        dd[get_canonical_tuple(mapping)].append(mapping)
    return (max(mappings, key=_pick_best) for mappings in dd.values())


def _pick_best(mapping: SemanticMapping) -> int:
    """Assign a value for this mapping.

    :param mapping: A mapping dictionary
    :returns: An integer, where higher means a better choice.

    This function is currently simple, but can later be extended to
    account for several other things including:

    - confidence in the curator
    - prediction methodology
    - date of prediction/curation (to keep the earliest)
    """
    if mapping.author and mapping.author.prefix == "orcid":
        return 1
    return 0


def load_curators() -> dict[str, NamedReference]:
    """Load the curators table."""
    with CURATORS_PATH.open() as file:
        return {
            record["user"]: NormalizedNamedReference(
                prefix="orcid", identifier=record["orcid"], name=record["name"]
            )
            for record in csv.DictReader(file, delimiter="\t")
        }


def filter_predictions(custom_filter: Mapping[str, Mapping[str, Mapping[str, str]]]) -> None:
    """Filter all the predictions by removing what's in the custom filter then re-write.

    :param custom_filter: A filter 3-dictionary of source prefix to target prefix
        to source identifier to target identifier
    """
    predictions = load_predictions()
    predictions = [
        prediction for prediction in predictions if _check_filter(prediction, custom_filter)
    ]
    write_predictions(predictions)


def _check_filter(
    prediction: SemanticMapping,
    custom_filter: Mapping[str, Mapping[str, Mapping[str, str]]],
) -> bool:
    v = (
        custom_filter.get(prediction.subject.prefix, {})
        .get(prediction.object.prefix, {})
        .get(prediction.subject.identifier)
    )
    return prediction.object.identifier != v


def get_curated_filter() -> Mapping[str, Mapping[str, Mapping[str, str]]]:
    """Get a filter over all curated mappings."""
    d: defaultdict[str, defaultdict[str, dict[str, str]]] = defaultdict(lambda: defaultdict(dict))
    for m in itt.chain(load_mappings(), load_false_mappings(), load_unsure()):
        d[m.subject.prefix][m.object.prefix][m.subject.identifier] = m.object.identifier
    return {k: dict(v) for k, v in d.items()}


def prediction_tuples_from_semra(
    mappings: Iterable[semra.Mapping],
    *,
    confidence: float,
) -> list[SemanticMapping]:
    """Get prediction tuples from SeMRA mappings."""
    rows = []
    for mapping in mappings:
        try:
            row = _mapping_from_semra(mapping, confidence)
        except KeyError as e:
            tqdm.write(str(e))
            continue
        rows.append(row)
    return rows


def _mapping_from_semra(mapping: semra.Mapping, confidence: float) -> SemanticMapping:
    """Instantiate from a SeMRA mapping."""
    import pyobo
    import semra

    s_name = mapping.s.name or pyobo.get_name(mapping.s)
    if not s_name:
        raise KeyError(f"could not look up name for {mapping.s.curie}")
    o_name = mapping.o.name or pyobo.get_name(mapping.o)
    if not o_name:
        raise KeyError(f"could not look up name for {mapping.o.curie}")
    # Assume that each mapping has a single simple evidence with a mapping set annotation
    if len(mapping.evidence) != 1:
        raise ValueError
    evidence = mapping.evidence[0]
    if not isinstance(evidence, semra.SimpleEvidence):
        raise TypeError
    if evidence.mapping_set is None:
        raise ValueError
    return SemanticMapping(  # type:ignore
        subject=mapping.subject,
        predicate=mapping.predicate,
        object=mapping.object,
        mapping_justification=evidence.justification,
        confidence=confidence,
        mapping_tool=evidence.mapping_set.name,
    )
