from typing import Dict
from abc import ABC, abstractmethod
from analogai.foundation.settings import config
from analogai.foundation.common.utils.statement_serializers.helpers.relation_verbalization_map import RELATION_TO_VERBALIZATION


class VerbSelectionStrategy(ABC):
    @abstractmethod
    def select_verb(self, relation: str, probability: float, is_plural: bool) -> str:
        pass


class DefaultVerbSelectionStrategy(VerbSelectionStrategy):
    def __init__(self, verbalization_map: Dict):
        self.verbalization_map = verbalization_map
    
    def select_verb(self, relation: str, probability: float, is_plural: bool) -> str:
        verbalization = RELATION_TO_VERBALIZATION.get(relation)
        if verbalization is None:
            return relation
        try:
            forms = self.verbalization_map[verbalization.name]._forms(is_plural)
        except KeyError:
            return relation
        if probability < config.get("relation.pos.neg", 0.3):
            return forms["negative"]
        elif probability <= config.get("relation.pos.uncertain", 0.7):
            return forms["uncertain"]
        return forms["positive"]


class VerbSelector:
    def __init__(self, strategy: VerbSelectionStrategy):
        self.strategy = strategy
    
    def select(self, relation: str, probability: float, is_plural: bool) -> str:
        return self.strategy.select_verb(relation, probability, is_plural)


__all__ = ["VerbSelector", "DefaultVerbSelectionStrategy", "VerbSelectionStrategy"]
