import gradio as gr
import pandas as pd
import circuit_transformer as ct

circuit_transformer = ct.CircuitTransformer()
aiger_str = """aag 33 8 0 2 25
2\n4\n6\n8\n10\n12\n14\n16\n58\n67
18 13 16\n20 19 7\n22 21 15\n24 3 9\n26 25 11
28 27 17\n30 3 6\n32 29 31\n34 29 32\n36 23 35
38 7 36\n40 10 29\n42 41 32\n44 13 15\n46 42 45
48 47 21\n50 39 49\n52 4 45\n54 25 53\n56 54 5
58 51 57\n60 45 12\n62 18 61\n64 63 19\n66 48 64
"""

def predict(input_text, mcts_playouts):
    if input_text[-1] != '\n':
        input_text += '\n'
    try:
        aig, info = ct.read_aiger(aiger_str=input_text)
    except Exception as e:
        raise gr.Error("Aiger format error!")
    num_inputs, num_outputs = info[1], info[3]
    if not isinstance(aig, list):
        aig = [aig]
    if num_inputs > 8:
        raise gr.Error("Number of inputs should <= 8")
    if num_outputs > 2:
        raise gr.Error("Number of outputs should <= 2")
    mcts_playouts = int(mcts_playouts)
    optimized_aigs = circuit_transformer.optimize(
        aigs=[aig],
        num_mcts_steps=1 if mcts_playouts > 0 else 0,
        num_mcts_playouts_per_step=mcts_playouts
    )
    optimized_aig = optimized_aigs[0]
    input_tt = ct.compute_input_tt(num_inputs)
    aig_tts = ct.compute_tts(aig, input_tt=input_tt)
    aig_tts_str = "\n".join([tt.to01() for tt in aig_tts])
    optimized_aig_tts = ct.compute_tts(optimized_aig, input_tt=input_tt)
    optimized_aig_tts_str = "\n".join([tt.to01() for tt in optimized_aig_tts])
    num_ands = pd.DataFrame(data={
        "circuit": ["Original", "Optimized"],
        "size": [ct.count_num_ands(aig), ct.count_num_ands(optimized_aig)]
    })
    ct.plot_network(aig, view=False, filename="aig.png")
    ct.plot_network(optimized_aig, view=False, filename="optimized_aig.png")
    return ct.write_aiger(optimized_aig), num_ands, "aig.png", "optimized_aig.png", aig_tts_str, optimized_aig_tts_str, str(ct.cec(aig, optimized_aig))

gradio_app = gr.Interface(
    predict,
    inputs=[gr.Code(value=aiger_str, label="Input logic circuit (in Aiger format, #(inputs) <= 8, #(outputs) <= 2)"),
            gr.Slider(minimum=0, maximum=8, value=0, step=1, label="Monte-Carlo tree search steps")],
    outputs=[gr.Code(label="Output logic circuit (size-optimized, in Aiger format)"),
             gr.BarPlot(label="Circuit size comparison", x="circuit", y="size", sort="-x"),
             gr.Image(label="Original logic circuit (And-Inverter Graph)"),
             gr.Image(label="Optimized logic circuit (And-Inverter Graph)"),
             gr.Text(label="Truth table of the original circuit"),
             gr.Text(label="Truth table of the optimized circuit"),
             gr.Text(label="Equivalence"),],
    title="Circuit Transformer for Size Minimization of Logic Circuits",
    description="""This is a demo to show how a [Circuit Transformer](https://openreview.net/forum?id=kpnW12Lm9p) 
minimize the size of a logic circuit by next token prediction, while strictly preserving logical equivalence.
[[GitHub Repo]](https://github.com/snowkylin/circuit-transformer)"""
)

if __name__ == "__main__":
    gradio_app.launch()