from rich.console import Console
from rich.table import Table
from itertools import product


class InvalidInputException(Exception):
    pass

def JK_flop(Q, J, K):
    # Returns Q+ and Q+'
    if J == 0 and K == 0:
        return Q, Q^1
    if J == 0 and K == 1:
        return 0, 1
    if J == 1 and K == 0:
        return 1, 0
    if J == 1 and K == 1:
        return Q^1, Q
    
def SR_flop(Q, S, R):
    # Returns Q+ and Q+'
    if S == 0 and R == 0: 
        return Q, Q^1
    if S == 0 and R == 1:
        return 0, 1
    if S == 1 and R == 0:
        return 1, 0
    if S == 1 and R == 1:
        raise InvalidInputException("Unpredictable behaviour when S=1 and R=1.")

def D_flop(Q):
    # Returns Q+ and Q+'
    return Q, Q^1

def T_flop(Q, T):
    if T == 1:
        return Q^1, Q
    if T == 0:
        return Q, Q^1 

def main():
    table = Table(show_lines=True)
    for name in "A B C X".split(" "):
        table.add_column(name, style="cyan")
    for name in "A+ B+ C+ Z".split(" "):
        table.add_column(name, style="green")
    for name in "JA KA JB KB JC KC".split(" "):
        table.add_column(name, style="yellow")


    for A,B,C,X in product([0,1], repeat=4):
        # Use & for AND, | for OR, (x^1) for NOT
        # Make sure to bracket everything

        # Output
        Z = X & B & A 

        # Flip Flop inputs
        JA = X & B & (C^1) 
        KA = (X^1)

        JB = X 
        KB = (X^1) | (A^1)

        JC = 1 
        KC = 1

        A_plus, _ = JK_flop(A, JA, KA)
        B_plus, _ = JK_flop(B, JB, KB)
        C_plus, _ = JK_flop(C, JC, KC) 

        # Add to table in this order
        table.add_row(*[str(i) for i in [A,B,C,X,A_plus,B_plus,C_plus,Z,JA,KA,JB,KB,JC,KC]])
    
    console = Console()
    console.print(table)

if __name__ == "__main__":
    main()