#!/usr/bin/env python3

"""
pylib_dfa.py

Usage: pylib_dfa.py

Class representing deterministic automata.

Options:

-h
--help  Print this help.

-d
--dot   Print a description of the automaton in the Graphviz language
        suitable for processing with dot to get a graphical
        representation of the automaton.

-p
--print Print back the automaton in tabular form.

Copyright 2015 Stephan Schulz, schulz@eprover.org

This program is free software; you can redistribute it and/or modify
it under the terms of the GNU General Public License as published by
the Free Software Foundation; either version 2 of the License, or
(at your option) any later version.

This program is distributed in the hope that it will be useful,
but WITHOUT ANY WARRANTY; without even the implied warranty of
MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
GNU General Public License for more details.

You should have received a copy of the GNU General Public License
along with this program ; if not, write to the Free Software
Foundation, Inc., 59 Temple Place, Suite 330, Boston,
MA  02111-1307 USA

The original copyright holder can be contacted as

Stephan Schulz
DHBW Stuttgart
Informatik
Postfach 10 05 63
D-70004 Stuttgart
Germany

or via email (address above).
"""

import sys
import re
import unittest

start_re = re.compile("dummy_x23z *-> *([A-Za-z0-9]+)")
trans_re = re.compile("([A-Za-z0-9]+) *-> *([A-Za-z0-9]+) *\[label *= *([^]]*)\]")

accept_re = re.compile("node *\[shape *= *doublecircle\];")
naccept_re = re.compile("node *\[shape *= *circle\]")
state_re = re.compile(" *([A-Za-z0-9]+)( *\[label=[^]]*\])?;?")

subscr_re=re.compile("[0-9]+$")

dfa_verbose = False

def Verbose(msg):
    if dfa_verbose:
        print(msg)



def prettyState(str):
    mo = subscr_re.search(str)
    if mo:
        init = str[0:mo.start()]
        subscr = str[mo.start():]
        str = "%s [label=<%s<sub>%s</sub>>]"%(str,init,subscr)
    return str

class SymRel(object):
    def __init__(self, tuples=None):
        self.relation = set()
        if tuples:
            for t in tuples:
                self.addTuple(t)

    def addTuple(self, tuple):
        e1,e2 = tuple
        self.addRel(e1, e2)

    def addRel(self, e1, e2):
        if e1 <= e2:
            self.relation.add((e1, e2))
        else:
            self.relation.add((e2, e1))

    def inRelTuple(self, tuple):
        e1,e2 = tuple
        return self.inRel(e1, e2)

    def inRel(self, e1, e2):
        if e1 <= e2:
            return (e1, e2) in self.relation
        else:
            return (e2, e1) in self.relation

    def getPairs(self):
        return self.relation.copy()

    def __str__(self):
        return str(self.relation)


class DFA(object):
    """
    Object representing a (deterministic) finite automaton.
    """

    def __init__(self, spec=None):
        self.states = set()
        self.sigma  = []
        self.delta  = {}
        self.start  = None
        self.accept = set()
        self.pretty = prettyState
        if spec:
            self.parse(spec)

    def delta_fun(self, state, letter):
        return self.delta[(state, letter)]

    def proc_string(self, string):
        print("Processing: ", string)

        state = self.start
        for c in string:
            newstate = self.delta_fun(state, c)
            print(state, "::", c, "->", newstate)
            state=newstate

        if state in self.accept:
            print("Accepted\n")
        else:
            print("Rejected\n")

    def get_accepting(self):
        return self.accept

    def get_nonaccepting(self):
        return self.states - self.accept

    def set_start(self, start):
        # Will overwrite previous start state, if any
        self.states.add(start)
        self.start = start

    def set_accepting(self, state):
        self.states.add(state)
        self.accept.add(state)

    def add_trans(self, fr, char, to):
        if (fr, char) in self.delta:
            print("Error: Duplicate entry for (%s, %s) -> %s)"%(fr,char,to))
        self.add_state(fr)
        self.add_state(to)
        if not char in self.sigma:
            self.sigma.append(char)
        self.delta[(fr,char)]=to

    def add_state(self,state):
        self.states.add(state)


    def merge_states(self, s1, s2):
        """
        Merge state s2 into s1.
        """
        assert s1!=s2

        if not s1 in self.states:
            print("# %s already removed"%(s1,))
            return
        if not s2 in self.states:
            print("# %s already removed"%(s2,))
            return
        print("# Merging %s into %s"%(s2, s1))

        for d in self.delta:
            if self.delta[d] == s2:
                self.delta[d] = s1

        delta_rem = set()
        for s,c in self.delta:
            if s==s2:
                delta_rem.add((s,c))
        for d in delta_rem:
            del self.delta[d]

        self.states.remove(s2)
        try:
            self.accept.remove(s2)
        except KeyError:
            pass

    def parse(self, spec):

        lines = spec.split("\n")
        # Find sigma
        while True:
            i = lines.pop(0)
            l=i.strip()
            if l and l[0] == '#':
                continue
            if l:
                sigmastr = l.split("|")[1]
                self.sigma = sigmastr.split()
                break

        # Skip ---------
        while True:
            i = lines.pop(0)
            l=i.strip()
            if l and l[0] == '#':
                continue
            if l:
                break
        # Process the rest - "(->)?(*)? state | state1 ... staten"
        while lines:
            i = lines.pop(0)
            l=i.strip()
            if l and l[0] == '#':
                continue
            if l:
                state, values = l.split("|")
                state = state.strip()

                start  = False
                accept = False
                if state.startswith("->"):
                    start = True
                    state = state[2:].strip()
                if state.startswith("*"):
                    accept = True
                    state = state[1:].strip()

                self.states.add(state)
                if start:
                    self.start = state
                if accept:
                    self.accept.add(state)


                dvals = values.split()

                for i in range(len(self.sigma)):
                    self.delta[(state, self.sigma[i])] = dvals[i]

    def parse_gv(self, inp):
        lines = inp.split("\n")
        accepting = False
        parseStates = False
        for line in lines:
            mo = accept_re.search(line)
            if mo:
                Verbose("# Accept")
                accepting = True
                parseStates = True
                line = line[mo.end():]
            else:
                mo = naccept_re.search(line)
                if mo:
                    accepting = False
                    parseStates = True
                    line = line[mo.end():]

            mo = start_re.search(line)
            if mo:
                Verbose("# Start state is "+str(mo.groups()[0]))
                self.set_start(mo.groups()[0])
                continue
            else:
                Verbose("# Scanning line "+line)
                mo = trans_re.search(line)
                if mo:
                    Verbose("# Found")
                    fr = mo.groups()[0]
                    to = mo.groups()[1]
                    transstr =  mo.groups()[2].strip("\"")
                    chars = transstr.split(",")
                    for c in chars:
                        c=c.strip()
                        char = ""
                        if c != "<&epsilon;>":
                            char = c
                        Verbose("# Transition %s -- %s -> %s"% (fr, char, to))
                        self.add_trans(fr, char, to)
                    continue
            if line.find("#")==0 or line.find("}")!=-1 or line.find("{")!=-1:
                continue

            if parseStates:
                mo = state_re.search(line)
                while mo:
                    Verbose("# Found state "+str(mo.groups()[0]))
                    if accepting:
                        self.set_accepting(mo.groups()[0])
                    self.add_state(mo.groups()[0])
                    line = line[mo.end():]
                    mo = state_re.search(line)


    def __str__(self):
        res = list()
        sigma = sorted(self.sigma)
        l1 = "         | "+" ".join(["%4s"%(i,) for i in sigma])
        res.append(l1)
        l2 = "-"*(len(l1)+1)
        res.append(l2)

        for i in sorted(self.states):
            if i == self.start:
                start = "->"
            else:
                start = "  "
            if i in self.accept:
                accept = "*"
            else:
                accept = " "

            xstate = "%2s %1s %3s | " %(start, accept, i)
            l = xstate+" ".join(["%4s"% (self.delta[(i,s)],) for s in sigma])
            res.append(l)

        return "\n".join(res)


    dfa_dot_template = """
digraph finite_state_machine {
   rankdir=LR;
   node [shape = plaintext]; dummy_x23z [label = ""];
   node [shape = doublecircle]; %s
   node [shape = circle]; %s
%s
}
"""

    def dotify(self):

        accepts_str = "; ".join(map(self.pretty,self.accept))
        if accepts_str:
            accepts_str += ";"

        other_states = [state for state in self.states if not state in self.accept]
        other_str = "; ".join(map(self.pretty,other_states))
        if other_str:
            other_str += ";"

        transitions = []
        if self.start:
            transitions.append("   dummy_x23z -> %s\n"%(self.start,))

        for i in self.delta.keys():
            source, label = i
            to = self.delta_fun(source, label)
            transitions.append("   %s -> %s [label = \"%s\"];\n"%
                               (source, to, label))

        transitions_str = "".join(transitions)

        return self.dfa_dot_template%(accepts_str,other_str,transitions_str)

    def minimize(self):
        """
        Minimize the automaton.
        """
        v = SymRel()
        for rel in [(x,y) for x in self.get_accepting() for y in self.get_nonaccepting()]:
            v.addTuple(rel)

        all_pairs = SymRel([(x,y) for x in self.states for y in self.states if x!=y])
        open_pairs = all_pairs.getPairs() - v.getPairs()

        incomplete = True
        while incomplete:
            incomplete = False
            new_v = set()
            for pair in open_pairs:
                s1, s2 = pair
                for c in self.sigma:
                    d1 = self.delta_fun(s1, c)
                    d2 = self.delta_fun(s2, c)
                    if v.inRel(d1,d2):
                        new_v.add(pair)
                        v.addTuple(pair)
                        incomplete = True
                        print("# Adding (%s, %s) to V because delta(%s,%s)=%s, delta(%s,%s)=%s, and (%s,%s) in V"%(s1,s2,s1,c,d1,s2,c,d2,d1,d2))
                        break
            open_pairs -= new_v

        table = []
        line=["# +---+"]
        s_states = sorted(self.states)
        for j in s_states:
            line.append("%3s+"%("---"),)
        sep="".join(line)
        table.append(sep)
        line=["# |   |"]
        for j in s_states:
            line.append("%3s|"%(j,),)
        table.append("".join(line))
        table.append(sep)
        for i in s_states:
            line=["# |%3s|"%(i,)]
            for j in s_states:
                if v.inRel(i,j):
                    line.append(" X |")
                else:
                    line.append(" o |")
            table.append("".join(line))
            table.append(sep)
        res = "\n".join(table)

        for s1, s2 in open_pairs:
            self.merge_states(s1,s2)

        return res

def prod_encode(state):
    s1,s2 = state
    return s1+s2

def product(a1, a2, union=False):
    res = DFA()
    res.pretty = lambda x:x
    sigma = set(a1.sigma)
    sigma.union(set(a2.sigma))
    res.sigma = list(sigma)
    res.sigma.sort()
    states = [(s1,s2) for s1 in a1.states for s2 in a2.states]
    res.states = set(map(prod_encode,states))
    for state in states:
        s1,s2 = state
        for l in res.sigma:
            r1 = a1.delta_fun(s1, l)
            r2 = a2.delta_fun(s2, l)
            res.add_trans(prod_encode(state), l, prod_encode((r1, r2)))

    res.set_start(prod_encode((a1.start,a2.start)))

    if union:
        acc1 = [(c1, c2) for c1 in a1.get_accepting() for c2 in a2.states]
        acc2 = [(c1, c2) for c1 in a1.states for c2 in a2.get_accepting()]
        res.accept=set(map(prod_encode, acc1))
        res.accept.union(set(map(prod_encode, acc2)))
    else:
        accepting = [(c1, c2) for c1 in a1.get_accepting() for c2 in a2.get_accepting()]
        res.accept = set(map(prod_encode, accepting))

    return res

testdfa = """
 A2     |  0  1  2
--------------------
->   q0 | q0  q1 q2
     q1 | q0  q1 q2
   * q2 | q0  q1 q2
"""


testdfagv="""
digraph finite_state_machine {
   rankdir=LR;
   size="8,5"
   node [shape = plaintext]; dummy_x23z [label = ""];
   node [shape = doublecircle]; q2;
   node [shape = circle]; q1 q0;
   dummy_x23z -> q0
   q0 -> q1 [label = "1"];
   q0 -> q0 [label = "0"];
   q1 -> q1 [label = "1"];
   q1 -> q0 [label = "0"];
   q0 -> q2 [label = "2"];
   q2 -> q2 [label = "2"];
   q2 -> q1 [label = "1"];
   q2 -> q0 [label = "0"];
   q1 -> q2 [label = "2"];

}
"""

class TestSymRel(unittest.TestCase):
    """
    Test symmetric relation feature.
    """
    def testSymRel(self):
        rel = SymRel()
        self.assertTrue(not rel.inRel("a","b"))
        rel.addRel("a","b")
        self.assertTrue(rel.inRel("a","b"))
        self.assertTrue(rel.inRel("b","a"))
        rel.addTuple(("b","c"))
        self.assertTrue(rel.inRel("b","c"))
        self.assertTrue(rel.inRelTuple(("c","b")))
        print(rel)


class TestDFA(unittest.TestCase):
    """
    Test basic functionality of the DFA.
    """

    def testDFA(self):
        """
        Test signature object.
        """
        dfa1 = DFA()
        dfa1.parse(testdfa)

        self.assertEqual(dfa1.get_accepting(), set(["q2"]))
        self.assertEqual(dfa1.get_nonaccepting(), set(["q1","q0"]))

        dfa2 = DFA(testdfa)

        print(dfa1)
        print(dfa1.dotify())

        self.assertEqual(str(dfa1), str(dfa2))
        self.assertEqual(dfa1.dotify(), dfa2.dotify())

        dfa3 = DFA()
        dfa3.parse_gv(testdfagv)
        print(dfa3)

        self.assertEqual(str(dfa1), str(dfa3))

        print("Product Automaton\n==============")
        prod = product(dfa1, dfa2)
        print(prod)
        prod = product(dfa1, dfa2, True)
        print(prod.dotify())

if __name__ == '__main__':
    unittest.main()