"""
Embedding similarity engine.

This module provides functionality for calculating the similarity between entities
using embeddings and cosine similarity. It includes classes and methods for
searching for similar entities and retrieving textual representations of entities.

Classes:
    EmbeddingSimilarity: A dataclass representing the similarity between two entities.
    SimilarityEngine: A class for calculating similarity between entities.

Methods:
    SimilarityEngine.similarity(entity1: str, entity2: str) -> EmbeddingSimilarity:
        Get similarity between two entities.
    SimilarityEngine.search(entity1: str, entities: List[str]) -> Iterable[EmbeddingSimilarity]:
        Search for similar entities.
    SimilarityEngine.entity_text(entity: str) -> str:
        Get text representation for an entity.
"""

import logging
from dataclasses import dataclass
from typing import Iterable, List

from oaklib import BasicOntologyInterface
from oaklib.datamodels.vocabulary import IS_A

from ontogpt.clients import LLMClient
from ontogpt.engines.knowledge_engine import KnowledgeEngine

logger = logging.getLogger(__name__)


@dataclass
class EmbeddingSimilarity:
    subject_id: str = ""
    subject_label: str = ""
    object_id: str = ""
    object_label: str = ""
    embedding_cosine_similarity: float = None
    object_rank_for_subject: int = None


@dataclass
class SimilarityEngine(KnowledgeEngine):
    """Engine for finding embedding similarity."""

    adapter: BasicOntologyInterface = None
    autolabel: bool = True
    definitions: bool = True
    parents: bool = True
    ancestors: bool = True
    synonyms: bool = True
    logical_definitions: bool = False

    def similarity(self, entity1: str, entity2: str) -> EmbeddingSimilarity:
        """Get similarity."""
        t1 = self.entity_text(entity1)
        t2 = self.entity_text(entity2)
        client = LLMClient()
        score = client.similarity(t1, t2)
        obj = EmbeddingSimilarity(
            subject_id=entity1, object_id=entity2, embedding_cosine_similarity=score
        )
        if self.autolabel:
            obj.subject_label = self.adapter.label(entity1)
            obj.object_label = self.adapter.label(entity2)
        return obj

    def search(self, entity1: str, entities: List[str]) -> Iterable[EmbeddingSimilarity]:
        """Get similarity."""
        client = LLMClient()
        t1 = self.entity_text(entity1)
        sims = []
        for entity2 in entities:
            t2 = self.entity_text(entity2)
            score = client.similarity(t1, t2)
            sim = EmbeddingSimilarity(
                subject_id=entity1, object_id=entity2, embedding_cosine_similarity=score
            )
            if self.autolabel:
                sim.subject_label = self.adapter.label(entity1)
                sim.object_label = self.adapter.label(entity2)
            sims.append(sim)
        sims = sorted(sims, key=lambda x: x.embedding_cosine_similarity, reverse=True)
        for i, sim in enumerate(sims):
            sim.object_rank_for_subject = i
        yield from sims

    def entity_text(self, entity: str) -> str:
        """Get text for an entity."""
        adapter = self.adapter
        s = f"{adapter.label(entity)}"
        if self.definitions:
            s += f"\ndefinition: {adapter.definition(entity)}"
        if self.parents:
            parent_labels = [
                adapter.label(o) for _s, _p, o in adapter.relationships([entity], [IS_A])
            ]
            s += f"\nparents: {'; '.join(parent_labels)}"
        if self.ancestors:
            ancestor_labels = [
                adapter.label(a) for a in adapter.ancestors([entity], [IS_A], reflexive=False)
            ]
            s += f"\nancestors: {'; '.join(ancestor_labels)}"
        if self.synonyms:
            s += f"\nsynonyms: {'; '.join(adapter.entity_aliases(entity))}"
        if self.logical_definitions:
            for ldef in adapter.logical_definitions(entity):
                genus_labels = [adapter.label(g) for g in ldef.genusIds]
                restriction_labels = [
                    f"{adapter.label(r.propertyId)} {adapter.label(r.valueId)}"
                    for r in ldef.restrictions
                ]
                s += f"\nlogical definition: A {', '.join(genus_labels)} that\
                    {' and '.join(restriction_labels)}"
            s += f"\nlogical definitions: {'; '.join(adapter.logical_definitions(entity))}"
        logger.info(f"Entity text for {entity}: {s}")
        return s
