import pprint

# translated from https://algs4.cs.princeton.edu/53substring/KMP.java.html

def KMP(pat):
    R=256
    m=len(pat)

    # build DFA from pattern
    # https://stackoverflow.com/questions/2397141/how-to-initialize-a-two-dimensional-array-in-python
    dfa=[[0]*m for r in range(R)]
    dfa[ord(pat[0])][0]=1
    x=0
    for j in range(1, m):
        for c in range(R):
            dfa[c][j] = dfa[c][x]      # Copy mismatch cases. 
        dfa[ord(pat[j])][j] = j+1      # Set match case. 
        x = dfa[ord(pat[j])][x]        # Update restart state. 

    return dfa

# https://stackoverflow.com/questions/3525953/check-if-all-values-of-iterable-are-zero
def all_elements_in_list_are_zeros (l):
    return all(v==0 for v in l)

def export_dfa_to_graphviz(dfa, fname):
    m=len(dfa[0]) # len of pattern
    f=open(fname,"w")
    f.write("digraph finite_state_machine {\n")
    f.write("rankdir=LR;\n")
    f.write("\n")
    f.write("size=\"8,5\"\n")
    f.write(f"node [shape = doublecircle]; S_0 S_{m};\n")
    f.write("node [shape = circle];\n")

    for state in range(m):
        exits=[]
        for R in range(256):
            next=dfa[R][state]
            if next!=0:
                exits.append((R,next))
        for exit in exits:
            next_state=exit[1]
            label="'"+chr(exit[0])+"'"
            s=f"S_{state} -> S_{next_state} [ label = \"{label}\" ];\n"
            f.write (s)
        s=f"S_{state} -> S_0 [ label = \"other\" ];\n"
        f.write (s)
    f.write("}\n")
    f.close()
    print (f"{fname} written")

def search(dfa, txt):
    # simulate operation of DFA on text
    m=len(dfa[0]) # len of pattern
    n=len(txt)
    j=0 # FA state
    i=0
    while i<n and j<m:
        j = dfa[ord(txt[i])][j]
        i=i+1
    if j == m:
        return i - m # found
    return n         # not found

def search_eel(txt):
    # simulate operation of DFA on text
    m=3 # len of pattern
    n=len(txt)
    j=0 # FA state
    i=0 # iterator for txt[]
    while i<n and j<m:
        ch=txt[i]
        if j==0 and ch=='e':
            j=1
        elif j==1 and ch=='e':
            j=2
        elif j==2 and ch=='e':
            j=2
        elif j==2 and ch=='l':
            j=3
        else:
            j=0 # reset
        i=i+1
    if j == m:
        return i - m # found
    return n         # not found

def search_egg(txt):
    # simulate operation of DFA on text
    m=3 # len of pattern
    n=len(txt)
    j=0 # FA state
    i=0 # iterator for txt[]
    while i<n and j<m:
        ch=txt[i]
        if j==0 and ch=='e':
            j=1
        elif j==1 and ch=='e':
            j=1
        elif j==1 and ch=='g':
            j=2
        elif j==2 and ch=='e':
            j=1
        elif j==2 and ch=='g':
            j=3
        else:
            j=0 # reset
        i=i+1
    if j == m:
        return i - m # found
    return n         # not found

def search_oops(txt):
    # simulate operation of DFA on text
    m=4 # len of pattern
    n=len(txt)
    j=0 # FA state
    i=0 # iterator for txt[]
    while i<n and j<m:
        ch=txt[i]
        if j==0 and ch=='o':
            j=1
        elif j==1 and ch=='o':
            j=2
        elif j==2 and ch=='o':
            j=2
        elif j==2 and ch=='p':
            j=3
        elif j==3 and ch=='o':
            j=1
        elif j==3 and ch=='s':
            j=4
        else:
            j=0 # reset
        i=i+1
    if j == m:
        return i - m # found
    return n         # not found

def search_ok(s):
    seen_o=False
    for i in range(len(s)):
        ch=s[i]           # this is single read operation
        if ch=='o':
            seen_o=True
        elif seen_o and ch=='k':
            return i-1    # found
        else:
            seen_o=False  # reset
    return len(s)         # not found

def dump_FA (dfa):
    for r in range(256):
        if not(all_elements_in_list_are_zeros(dfa[r])):
            print (chr(r), dfa[r])

def test_ok():
    dfa=KMP("ok")

    assert search(dfa, "ok")==search_ok("ok")
    assert search(dfa, "oook")==search_ok("oook")
    assert search(dfa, "okk")==search_ok("okk")
    assert search(dfa, "ookk")==search_ok("ookk")
    assert search(dfa, "hello, ok")==search_ok("hello, ok")

def test_hello():
    dfa=KMP("hello")

    assert search(dfa, "oh hello ha")==3
    assert search(dfa, "oh hell ha")==10
    assert search(dfa, " hello world ha")==1
    assert search(dfa, "hello world ha")==0

def test_oops():
    dfa=KMP("oops")

    assert search(dfa, "ops")==3 # not found
    assert search(dfa, "oops")==0
    assert search(dfa, "ooops")==1
    assert search(dfa, "oooops")==2
    assert search(dfa, "ooooops")==3
    assert search(dfa, "oooooops")==4
    assert search(dfa, "hello ooops")==7
    assert search(dfa, "hello oops")==6
    assert search(dfa, "hello oozing oops")==13
    assert search(dfa, "ooze oops")==5

    assert search_oops("ops")==3 # not found
    assert search_oops("oops")==0
    assert search_oops("ooops")==1
    assert search_oops("oooops")==2
    assert search_oops("ooooops")==3
    assert search_oops("oooooops")==4
    assert search_oops("hello ooops")==7
    assert search_oops("hello oops")==6
    assert search_oops("hello oozing oops")==13
    assert search_oops("ooze oops")==5

    export_dfa_to_graphviz(dfa, "oops.gv")

def test_eel():
    dfa=KMP("eel")
    
    assert search(dfa, "eel hello")==0
    assert search(dfa, "hello, eel")==7
    assert search(dfa, "hello, eeel")==8

    export_dfa_to_graphviz(dfa, "eel.gv")

def test_egg():
    dfa=KMP("egg")
    
    assert search(dfa, "eggs, hello")==0
    assert search(dfa, "hello, egg")==7
    assert search(dfa, "hello, eggs")==7

    assert search_egg("eggs, hello")==0
    assert search_egg("hello, egg")==7
    assert search_egg("hello, eggs")==7

    export_dfa_to_graphviz(dfa, "egg.gv")

def test_test():
    dfa=KMP("test")
    
    assert search(dfa, "test, hello")==0
    assert search(dfa, "a test")==2
    assert search(dfa, "tessa test")==6

    export_dfa_to_graphviz(dfa, "test.gv")

test_ok()
test_hello()
test_oops()
test_eel()
test_egg()
test_test()

dfa=KMP("abracadabra")
export_dfa_to_graphviz(dfa, "abracadabra.gv")
print ("FA for 'abracadabra'")
dump_FA(dfa)

dfa=KMP("cocoon")
export_dfa_to_graphviz(dfa, "cocoon.gv")
print ("FA for 'cocoon'")
dump_FA(dfa)

dfa=KMP("cocos")
export_dfa_to_graphviz(dfa, "cocos.gv")
print ("FA for 'cocos'")
dump_FA(dfa)

dfa=KMP("banana")
export_dfa_to_graphviz(dfa, "banana.gv")
print ("FA for 'banana'")
dump_FA(dfa)