from abc import ABC, abstractmethod
from dataclasses import dataclass
from functools import cached_property

@dataclass(frozen=True)
class SimplexMap:
  """A sSet map symbol, mathematically defined as a
  monotone map of ordered finite sets. Represented as a
  tuple of natural numbers, where `l[i]` represents the
  number of elements mapped to `i`."""
  data : tuple[int]
  def __repr__(self):
    return f"⟨Simplex Map [{self.dom}] -> [{self.cod}]: {self.data}⟩"

  def __getitem__(self, k):
    return self.data[k]

  @cached_property
  def dom(self):
    return sum(self.data)-1
  
  @cached_property
  def cod(self):
    return len(self.data)-1

  @staticmethod
  def identity(n : int):
    return SimplexMap(tuple(1 for _ in range(n+1)))
  
  def __mul__(p, q):
    """Given [m] -> [n] -> [s], computes [m] -> [s]."""
    ix = 0
    result = []
    for i in q:
      result.append(sum(p[ix:ix+i]))
      ix += i
    return SimplexMap(tuple(result))
  
  @cached_property
  def isDegeneracy(self):
    return all(i > 0 for i in self.data)

  @cached_property
  def isIdentity(self):
    return all(i == 1 for i in self.data)

  @cached_property
  def factor(self) -> tuple["SimplexMap", "SimplexMap"]:
    """Factors itself into a degeneracy `d` followed by a face map `f`.
    So it returns (d, f) such that `self = f * d`.
    It returns f by its zero indices, which is easier to use."""
    return (SimplexMap(tuple(i for i in self.data if i > 0)),
            tuple(ix for (ix, i) in enumerate(self.data) if i == 0))

Degeneracy = SimplexMap

class SimplicialSet(ABC):
  """A simplicial set. The elements of this type should be
  non-degenerate simplices. The `dim` function assigns the
  dimension of each simplex, and the `face` function calculates
  the face of a simplex (which may be degenerate). If this
  sSet is of finite type, further implement an `enum` function.
  """
  decidableEq : bool = False
  """Whether this complex has decidable equality of simplices.
  This is called "locally effective" in Kenzo."""

  @staticmethod
  def enum(dim: int):
    """If the specified dimension is known to have finitely many
    simplices, yield all of them. Otherwise return None. This is
    called "effective" in Kenzo.
    """
    raise NotImplementedError("This sSet is not of finite type.")

  @property
  @abstractmethod
  def dim(self) -> int:
    """A natural number, the dimension of the simplex."""
    raise NotImplementedError

  @abstractmethod
  def face(self, vertex: int) -> "Simplex":
    """Returns a possibly degenerate simplex.
    Requires 0 ≤ vertex ≤ dim(simplex)"""
    raise NotImplementedError
  
  def toSimplex(self) -> "Simplex":
    return Simplex(self, Degeneracy.identity(self.dim))

@dataclass(frozen=True)
class Simplex:
  simplex : SimplicialSet
  degen : Degeneracy
  def __repr__(self):
    return f"⟨Simplex {self.simplex} with {self.degen}⟩"
  def __iter__(self):
    return iter((self.simplex, self.degen))

  def apply(self, p: SimplexMap):
    """Apply a simplex map to a general simplex."""
    # Room for optimization
    if hasattr(self.simplex, "apply"):
      s, d0 = self.simplex.apply(p)
      return Simplex(s, d * d0)
    # If no optimized route, do the hard way
    d, f = (p * self.degen).factor
    s = self.simplex
    for i in f[::-1]:
      s, d0 = s.face(i)
      d = d * d0
    return Simplex(s, d)

if __name__ == "__main__":
  # Example implementation: the circle
  class Circle(SimplicialSet):
    def __init__(self, s):  # s is "Base" or "Loop"
      self.type = s
    def __repr__(self): return self.type
    @property
    def dim(self):
      if self.type == "Base": return 0
      if self.type == "Loop": return 1
      raise ValueError
    def face(self, vertex):
      if self.type == "Base":
        raise ValueError  # impossible
      if self.type == "Loop":
        return Circle("Base").toSimplex()
    @staticmethod
    def enum(dim):
      if dim == 0:
        yield Circle("Base")
      elif dim == 1:
        yield Circle("Loop")
      else:
        return

  f = Degeneracy((1,2,1,0))
  print(f, f[1], f * f)

  loop = Circle("Loop").toSimplex()
  print(loop, loop.apply(SimplexMap((2,2))).apply(SimplexMap((0,0,1,1))))