app.py

import re
import io
from collections import defaultdict
from PIL import Image
from graphviz import Digraph
import gradio as gr

def parse_process_description(description):

    text = description.strip()

    # Split on commas, "then", "and"
    segments = re.split(r',|\band\b|\bthen\b', text, flags=re.IGNORECASE)
    segments = [s.strip() for s in segments if s.strip()]

    edges = []
    recycle_edges = []
    extra_inputs = set()
    final_output = None

    # We'll track the last_block for normal flows (so we can connect 'then to X')
    last_block = None

    # Regex patterns
    re_first_to         = re.compile(r'^first\s+(.*?)\s+to\s+(.*)$', re.IGNORECASE)
    re_to_with          = re.compile(r'^to\s+(.*?)\s+with\s+(.*)$', re.IGNORECASE)
    re_from_to_with     = re.compile(r'^from\s+(.*?)\s+to\s+(.*?)\s+with\s+(.*)$', re.IGNORECASE)
    re_to_only          = re.compile(r'^to\s+(.*)$', re.IGNORECASE)
    re_from_to_only     = re.compile(r'^from\s+(.*?)\s+to\s+(.*)$', re.IGNORECASE)
    re_final_output     = re.compile(r'^final\s+output\s+(.*)$', re.IGNORECASE)
    re_recycle          = re.compile(r'^recycle\s+from\s+(.*?)\s+to\s+(.*)$', re.IGNORECASE)

    for seg in segments:
        # 1) final output
        m_final = re_final_output.search(seg)
        if m_final:
            final_output = m_final.group(1).strip()
            continue

        # 2) recycle from X to Y
        m_rec = re_recycle.search(seg)
        if m_rec:
            src = m_rec.group(1).strip()
            dst = m_rec.group(2).strip()
            # We'll add this to recycle_edges
            recycle_edges.append((src, dst))
            # Importantly, we do NOT update last_block,
            # because the user wants final output to remain from the prior block.
            continue

        # 3) "First X to Y"
        m_first = re_first_to.search(seg)
        if m_first:
            inp = m_first.group(1).strip()
            blk = m_first.group(2).strip()
            edges.append((inp, blk))
            extra_inputs.add(inp)
            last_block = blk
            continue

        # 4) "from X to Y with Z"
        m_ftw = re_from_to_with.search(seg)
        if m_ftw:
            src = m_ftw.group(1).strip()
            dst = m_ftw.group(2).strip()
            extra_inp = m_ftw.group(3).strip()
            edges.append((src, dst))
            edges.append((extra_inp, dst))
            extra_inputs.add(extra_inp)
            last_block = dst
            continue

        # 5) "to Y with Z" (no 'from')
        m_tw = re_to_with.search(seg)
        if m_tw:
            block_candidate = m_tw.group(1).strip()
            extra_inp = m_tw.group(2).strip()
            if last_block:
                edges.append((last_block, block_candidate))
            edges.append((extra_inp, block_candidate))
            extra_inputs.add(extra_inp)
            last_block = block_candidate
            continue

        # 6) "from X to Y" (no 'with')
        m_ft = re_from_to_only.search(seg)
        if m_ft:
            src = m_ft.group(1).strip()
            dst = m_ft.group(2).strip()
            edges.append((src, dst))
            last_block = dst
            continue

        # 7) "to X" (no 'from', no 'with')
        m_t = re_to_only.search(seg)
        if m_t:
            blk = m_t.group(1).strip()
            if last_block:
                edges.append((last_block, blk))
            last_block = blk
            continue

        # If unmatched, ignore or debug:
        # print("Unrecognized segment:", seg)

    # If there's a final output & we have a last_block, connect them:
    if final_output and last_block:
        edges.append((last_block, final_output))

    return edges, recycle_edges, extra_inputs, final_output


def build_flowchart(edges, recycle_edges, extra_inputs, final_output):

    all_nodes = set()
    for s, t in edges:
        all_nodes.add(s)
        all_nodes.add(t)
    for s, t in recycle_edges:
        all_nodes.add(s)
        all_nodes.add(t)

    if final_output and final_output not in all_nodes:
        all_nodes.add(final_output)

    # Build in/out degrees for normal edges + recycle edges
    in_degree = defaultdict(int)
    out_degree = defaultdict(int)

    for s, t in edges:
        out_degree[s] += 1
        in_degree[t]  += 1

    for s, t in recycle_edges:
        out_degree[s] += 1
        in_degree[t]  += 1

    dot = Digraph(name="Flowchart", format="png")
    dot.attr(rankdir='LR')

    # Create each node
    for node in all_nodes:
        shape = "box"
        style = "rounded,filled"
        fillcolor = "lightgoldenrod1"

        # Circle if it's an extra input
        if node in extra_inputs:
            shape = "circle"
            style = "filled"
            fillcolor = "lightblue"

        # Double circle if final output
        if final_output and node == final_output:
            shape = "doublecircle"
            fillcolor = "lightgreen"
            style = "filled"

        # Tee/Mixer logic (skip if final output)
        indeg = in_degree[node]
        outdeg = out_degree[node]
        if node != final_output:
            if indeg > 1 and outdeg > 1:
                shape = "box"
                fillcolor = "lightgoldenrod1"
            elif indeg > 1:
                shape = "box"
                fillcolor = "lightgoldenrod1"
            elif outdeg > 1:
                shape = "box"
                fillcolor = "lightgoldenrod1"

        dot.node(node, label=node, shape=shape, style=style, fillcolor=fillcolor)

    # Add normal edges
    for s, t in edges:
        dot.edge(s, t)

    # Add recycle edges with dashed style or different color
    for s, t in recycle_edges:
        dot.edge(s, t, style="dashed", color="gray", label="recycle")

    return dot.pipe(format='png')


def flowchart_to_image(text_input):
    if not text_input.strip():
        return None

    edges, recycle_edges, extra_inputs, final_output = parse_process_description(text_input)
    png_data = build_flowchart(edges, recycle_edges, extra_inputs, final_output)
    return Image.open(io.BytesIO(png_data))

# Gradio interface
iface = gr.Interface(
    fn=flowchart_to_image,
    inputs=gr.Textbox(lines=7),
    outputs="image",
    title="Process Flow with Recycle Lines",
    description=(
        "Example:\n"
        "First Solid Sulfur to Heater,\n"
        "then to Primary Oxidizer with Air,\n"
        "then from Primary Oxidizer to Reactor,\n"
        "then from Primary Oxidizer to Secondary Oxidizer,\n"
        "from Secondary Oxidizer to Heater,\n"
        "then from Secondary Oxidizer to Absorber with Water,\n"
        "recycle from Absorber to Secondary Oxidizer,\n"
        "and final output Sulfuric Acid.\n\n"
        "This example text only shows how to add or divide more connections and recycle back stream."
    ),
    flagging_dir="/tmp/flagged_data"
)

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