from badic import *
from sage.combinat.words.word import Word
from sage.combinat.words.morphism import WordMorphism

class RandomSubstitution:
    """
    Random substitution, given as a dictionnary that maps to each letter a list of words.

    EXAMPLES::
        sage: from badic import *
        sage: rs = RandomSubstitution({'a':['ab', 'ba'], 'b':['ac'], 'c':['a']})
    """
    def __init__(self, d):
        """
        Constructor of RandomSubstitution.

         - d (dict) : dictionnary that maps to each letter a list of words.
        """
        self.d = d.copy()
        self.A = set()
        for a in d:
            for w in d[a]:
                for b in w:
                    self.A.add(b)
        self.A = list(self.A)
        self.A.sort()
        for a in d:
            self.d[a] = [Word(w, alphabet=self.A) for w in d[a]]
    def __repr__(self):
        return str(self.d)
    def __str__(self):
        return str(self.d)
    def is_compatible(self):
        """
        Determine whether the random substitution is compatible, i.e. if its matrix does not depends of probabilities.

        EXAMPLES::
            sage: from badic import *
            sage: rs = RandomSubstitution({'a':['ab', 'ba'], 'b':['ac'], 'c':['a']})
            sage: rs.is_compatible()
            True
        """
        for a in self.d:
            if len({tuple(w.abelian_vector()) for w in self.d[a]}) > 1:
                return False
        return True
    def incidence_matrix(self):
        """
        Return the incidence matrix, assuming that the random substitution is compatible.
        """
        if self.is_compatible():
            return WordMorphism({a:self.d[a][0] for a in self.d}).incidence_matrix()
        else:
            raise NotImplementedError("Sorry, implemented only for compatible substitutions.")
    def rauzy_fractal_projection(self):
        """
        Return the projection from Z^d to R^k associated to the incidence matrix, assuming that the random substitution is compatible.
        """
        if self.is_compatible():
            return rauzy_fractal_projection(WordMorphism({a:self.d[a][0] for a in self.d}))
        else:
            raise NotImplementedError("Sorry, implemented only for compatible substitutions.")
    def rauzy_fractal(self, prefix_aut = False):
        """
        Return a BetaAdicSet describing the Rauzy fractal.

         - prefix_aut (boolean) - if True, return the prefix automaton.

        EXAMPLES::
            sage: from badic import *
            sage: rs = RandomSubstitution({'a':['ab', 'ba'], 'b':['ac'], 'c':['a']})
            sage: rs.rauzy_fractal()
            b-adic set with b root of x^3 - x^2 - x - 1, and an automaton of 3 states and 3 letters
        """
        projection = self.rauzy_fractal_projection()
        b = projection([1 for i in self.A]).parent().gen()
        A = set()
        lt = []
        for l in self.d:
            for w in self.d[l]:
                for i in range(len(w)):
                    t = projection(w[:i].abelian_vector())
                    lt.append((self.A.index(l), t, self.A.index(w[i])))
                    A.add(t)
        A = list(A)
        A.sort()
        a = CAutomaton(DetAutomaton([], S=self.A, A=A, final_states=self.A))
        for (i,j,k) in lt:
            a.add_transition(i,j,k)
        for i in a.states:
            a.set_initial(i)
        if prefix_aut:
            return a
        return BetaAdicSet(b, a.determinize().mirror())