from abc import ABC, abstractmethod
import warnings
from dataclasses import dataclass

class CohesionKernel(ABC):
	"""
	Defines the substitution weight distributions over target symbols
	given a source symbol.
	"""

	def _normalize(self, source : str, dist: dict[str, float]) -> dict[str, float]:
		wsum = sum(w for w in dist.values())
		if wsum <= 1e-15:
			warnings.warn( "Weight function sums to 0, implicitly mapping source with probability 1" )
			res = dist.copy()
			res[ source ] = 1.0
			return res
		return {s: w / wsum for s, w in dist.items()}

	@abstractmethod
	def cohesion_scores( self, source: str, symbols: set[str] ) -> dict[str, float]:
		"""
		Maps a source symbol to a distribution of cohesion scores over other symbols.
		"""

@dataclass
class Uniform(CohesionKernel):
	def cohesion_scores(self, source: str, symbols: set[str]) -> dict[str, float]:
		return { s: 1.0 for s in symbols }

@dataclass
class Marginal(CohesionKernel):
	"""Per-symbol weight distribution."""
	scores : dict[str, float]

	def cohesion_scores(self, source: str, symbols: set[str]) -> dict[str, float]:
		return {s: self.scores[s] for s in symbols }

@dataclass
class Symmetric(CohesionKernel):
	"""w(a->b) == w(b->a)"""
	scores : dict[frozenset, float]
	def cohesion_scores(self, source: str, symbols: set[str]) -> dict[str, float]:
		return {s: self.scores[frozenset({source, s})] for s in symbols }

@dataclass
class Asymmetric(CohesionKernel):
	"""Directed pairwise scores."""
	scores : dict[tuple[str,str], float]

	def cohesion_scores(self, source: str, symbols: set[str]) -> dict[str, float]:
		return {s: self.scores[(source, s)] for s in symbols }