app.py

import gradio as gr
import json
import re
from collections import defaultdict
from ortools.sat.python import cp_model
from datetime import datetime, timedelta
from typing import Dict, List, Tuple, Optional, Union, Any

import pandas


def normalize_task_name(task: str) -> str:
    """Normalize task name to lowercase for consistent comparison."""
    return task.strip().lower().replace(" ", "_")


def display_task_name(task: str) -> str:
    return task


def parse_requirements(requirements: List[str]) -> Tuple[Dict[str, List[str]], List[str], Dict[str, str]]:
    """Parse requirements list into dependency graph."""
    dependencies = defaultdict(list)
    all_tasks = set()
    original_names = {}  # Store original casing

    for req in requirements:
        if not req.strip():
            continue
        # Match pattern like "TaskA requires TaskB" (now supports spaces)
        match = re.match(r"(.+?)\s+requires\s+(.+)", req.strip(), re.IGNORECASE)
        if match:
            task_orig, required_orig = match.groups()
            task = normalize_task_name(task_orig)
            required = normalize_task_name(required_orig)

            all_tasks.add(task)
            all_tasks.add(required)
            dependencies[task].append(required)

            # Store original names for display
            original_names[task] = task_orig
            original_names[required] = required_orig

    return dependencies, list(all_tasks), original_names


def solve_all_tasks(dependencies: Dict[str, List[str]], all_tasks: List[str]) -> Optional[List[str]]:
    """Try to schedule all tasks (works if DAG)."""
    model = cp_model.CpModel()
    n_tasks = len(all_tasks)

    position = {}
    for task in all_tasks:
        position[task] = model.NewIntVar(0, n_tasks - 1, f"pos_{task}")

    model.AddAllDifferent(list(position.values()))

    for task, required_list in dependencies.items():
        for required_task in required_list:
            if required_task in position:
                model.Add(position[required_task] < position[task])

    solver = cp_model.CpSolver()
    status = solver.Solve(model)

    if status in [cp_model.OPTIMAL, cp_model.FEASIBLE]:
        task_order = [(solver.Value(position[task]), task) for task in all_tasks]
        task_order.sort()
        return [task for _, task in task_order]

    return None


def solve_maximum_subset(dependencies: Dict[str, List[str]], all_tasks: List[str]) -> List[str]:
    """Find maximum subset of tasks that can be executed (handles cycles)."""
    model = cp_model.CpModel()
    n_tasks = len(all_tasks)

    include = {}
    position = {}

    for task in all_tasks:
        include[task] = model.NewBoolVar(f"include_{task}")
        position[task] = model.NewIntVar(0, n_tasks - 1, f"pos_{task}")

    for task, required_list in dependencies.items():
        for required_task in required_list:
            if required_task in include:
                model.AddImplication(include[task], include[required_task])
                model.Add(position[required_task] < position[task]).OnlyEnforceIf(
                    [include[task], include[required_task]]
                )

    for i, task1 in enumerate(all_tasks):
        for j, task2 in enumerate(all_tasks):
            if i < j:
                model.Add(position[task1] != position[task2]).OnlyEnforceIf(
                    [include[task1], include[task2]]
                )

    for task in all_tasks:
        model.Add(position[task] == n_tasks).OnlyEnforceIf(include[task].Not())

    model.Maximize(sum(include.values()))

    solver = cp_model.CpSolver()
    status = solver.Solve(model)

    if status in [cp_model.OPTIMAL, cp_model.FEASIBLE]:
        included_tasks = []
        for task in all_tasks:
            if solver.Value(include[task]) == 1:
                pos = solver.Value(position[task])
                included_tasks.append((pos, task))

        included_tasks.sort()
        return [task for _, task in included_tasks]

    return []


def parse_tasks(tasks_text: str) -> Tuple[List[str], Dict[str, str]]:
    """Parse tasks from text input."""
    if not tasks_text.strip():
        return [], {}

    # Split by lines or commas and clean up
    tasks = []
    original_names = {}

    for line in tasks_text.strip().split("\n"):
        for task_orig in line.split(","):
            task_orig = task_orig.strip()
            if task_orig:
                task = normalize_task_name(task_orig)
                if task not in [normalize_task_name(t) for t in tasks]:
                    tasks.append(task)
                    original_names[task] = task_orig

    return sorted(tasks), original_names


def generate_mermaid_gantt(task_order: List[str], original_names: Dict[str, str], durations: Optional[Dict[str, int]] = None, start_time: Optional[datetime] = None) -> str:
    """Generate Mermaid Gantt chart syntax with minute granularity and selectable start time."""
    if not task_order:
        return "gantt\n    title Task Execution Timeline\n    dateFormat  YYYY-MM-DD HH:mm\n    section No Tasks\n    No tasks to display : 2024-01-01 08:00, 1m"

    if start_time is None:
        start_date = datetime.now().replace(second=0, microsecond=0)
    else:
        start_date = start_time.replace(second=0, microsecond=0)

    gantt = "```mermaid\ngantt\n"
    gantt += "    title Task Execution Timeline\n"
    gantt += "    dateFormat  YYYY-MM-DD HH:mm\n"
    gantt += "    axisFormat  %a %H:%M\n    tickInterval 30minute\n"
    gantt += "    section Tasks\n"

    current_dt = start_date
    tasks = list(enumerate(task_order))
    for i, task in tasks:
        # Get duration for this task in minutes (default 1m if not specified)
        minutes = 1
        if durations and task in durations:
            minutes = durations[task]
        elif durations is None:
            minutes = 1
        else:
            minutes = durations.get(task, 1)
        
        after = f"after {tasks[i-1][1]}" if i > 0 else f"{current_dt.strftime('%Y-%m-%d %H:%M')}"

        gantt += f"    {task.title().replace('_', ' ')} : {task}, {after}, {minutes}m\n"
        current_dt += timedelta(minutes=minutes)

    gantt += "```"
    return gantt


def generate_mermaid_flowchart(dependencies: Dict[str, List[str]], all_tasks: List[str], original_names: Dict[str, str]) -> str:
    """Generate Mermaid flowchart syntax."""
    if not all_tasks:
        return "flowchart TD\n    A[No tasks to display]"

    flowchart = "```mermaid\nflowchart TD\n"

    # Create nodes for all tasks
    for task in all_tasks:
        display_name = original_names.get(task, display_task_name(task))
        # Create a clean ID for the node
        node_id = task.replace(" ", "_").replace("-", "_")
        flowchart += f"    {node_id}[{display_name}]\n"

    # Add dependencies as edges
    if dependencies:
        for task, required_list in dependencies.items():
            task_id = task.replace(" ", "_").replace("-", "_")
            for required_task in required_list:
                req_id = required_task.replace(" ", "_").replace("-", "_")
                flowchart += f"    {req_id} --> {task_id}\n"
    else:
        # If no dependencies, just show all tasks
        flowchart += "    style A fill:#e1f5fe\n"

    flowchart += "```"
    return flowchart


def update_dropdowns(tasks_text: str) -> Tuple[gr.Dropdown, gr.CheckboxGroup]:
    """Update dropdown choices when tasks change."""
    tasks, original_names = parse_tasks(tasks_text)

    # Create display choices (original casing)
    task_display_choices = [""] + [
        original_names.get(task, display_task_name(task)) for task in tasks
    ]
    req_display_choices = [
        original_names.get(task, display_task_name(task)) for task in tasks
    ]

    return (
        gr.Dropdown(choices=task_display_choices, value=""),
        gr.CheckboxGroup(choices=req_display_choices, value=[]),
    )


def add_dependencies(task_display: str, required_display_list: List[str], current_deps: List[str]) -> Tuple[List[str], str]:
    """Add multiple dependencies for a single task."""
    if not task_display:
        return current_deps, "⚠️ Please select a task"

    if not required_display_list:
        return current_deps, "⚠️ Please select at least one requirement"

    # Normalize for comparison
    task = normalize_task_name(task_display)
    required_tasks = [normalize_task_name(req) for req in required_display_list]

    # Remove the selected task from requirements if somehow selected
    required_tasks = [req for req in required_tasks if req != task]

    if not required_tasks:
        return current_deps, "⚠️ A task cannot require itself"

    # Create new dependencies (using original display names)
    new_deps = []
    existing_deps = []

    for req_display in required_display_list:
        if normalize_task_name(req_display) in required_tasks:
            new_dep = f"{task_display} requires {req_display}"

            # Check if this dependency already exists (case-insensitive)
            exists = any(
                normalize_task_name(dep.split(" requires ")[0]) == task
                and normalize_task_name(dep.split(" requires ")[1])
                == normalize_task_name(req_display)
                for dep in current_deps
                if " requires " in dep
            )

            if exists:
                existing_deps.append(req_display)
            else:
                new_deps.append(new_dep)

    if not new_deps and existing_deps:
        return current_deps, "⚠️ All selected dependencies already exist"

    updated_deps = current_deps + new_deps

    # Create status message
    status_parts = []
    if new_deps:
        status_parts.append(f"✅ Added {len(new_deps)} dependencies")
    if existing_deps:
        status_parts.append(f"⚠️ {len(existing_deps)} already existed")

    status = " | ".join(status_parts)

    return updated_deps, status


def remove_dependency(deps_to_remove: List[str], current_deps: List[str]) -> Tuple[List[str], str]:
    """Remove selected dependencies."""
    if not deps_to_remove:
        return current_deps, "⚠️ Please select dependencies to remove"

    updated_deps = [dep for dep in current_deps if dep not in deps_to_remove]
    removed_count = len(current_deps) - len(updated_deps)
    return updated_deps, f"✅ Removed {removed_count} dependencies"


def format_dependencies_display(dependencies_list: List[str]) -> str:
    """Format dependencies for display."""
    if not dependencies_list:
        return "No dependencies added yet."

    # Group by task for better readability
    task_deps = defaultdict(list)
    for dep in dependencies_list:
        parts = dep.split(" requires ")
        if len(parts) == 2:
            task, required = parts
            task_deps[task].append(required)

    display = "**Current Dependencies:**\n"
    for task in sorted(task_deps.keys(), key=str.lower):
        requirements = ", ".join(sorted(task_deps[task], key=str.lower))
        display += f"• **{task}** requires: {requirements}\n"

    return display


def solve_dependencies(
    tasks_text: str, dependencies_list: List[str], durations_state: Optional[Dict[str, Union[str, int]]] = None
) -> Tuple[str, str, str, str, str, str, str]:
    """Solve the task ordering problem."""
    tasks, task_original_names = parse_tasks(tasks_text)

    durations = durations_state or {}
    # durations is a dict: {normalized_task_name: int_minutes}, may be string-int if loaded from input
    # normalize keys and intify
    durations = {
        normalize_task_name(k): int(v)
        for k, v in durations.items()
        if str(v).isdigit() and int(v) > 0
    }

    if not tasks:
        return "❌ Please enter some tasks!", "", "", "", "", "", ""

    if not dependencies_list:
        # No dependencies, just return tasks in alphabetical order
        output = "✅ **No dependencies - alphabetical order:**\n\n"
        display_tasks = [
            task_original_names.get(task, display_task_name(task)) for task in tasks
        ]
        for i, task_display in enumerate(display_tasks, 1):
            output += f"{i}. {task_display}\n"

        json_output = json.dumps(display_tasks, indent=2)
        gantt = generate_mermaid_gantt(tasks, task_original_names, durations=durations)
        flowchart = generate_mermaid_flowchart({}, tasks, task_original_names)
        # Provide raw for copy-paste, strip code block
        gantt_raw = gantt.replace("```mermaid\n", "").rstrip("`\n")
        flow_raw = flowchart.replace("```mermaid\n", "").rstrip("`\n")
        return output, "", json_output, gantt, flowchart, gantt_raw, flow_raw

    try:
        dependencies, all_tasks, dep_original_names = parse_requirements(
            dependencies_list
        )

        # Merge original names
        all_original_names = {**task_original_names, **dep_original_names}

        # Add any tasks that aren't in dependencies
        for task in tasks:
            if task not in all_tasks:
                all_tasks.append(task)

        # Create dependency summary
        dep_summary = "**Parsed Dependencies:**\n"
        if dependencies:
            for task, deps in dependencies.items():
                task_display = all_original_names.get(task, display_task_name(task))
                deps_display = [
                    all_original_names.get(dep, display_task_name(dep)) for dep in deps
                ]
                dep_summary += f"• {task_display} requires: {', '.join(deps_display)}\n"

        dep_summary += f"\n**Total tasks:** {len(all_tasks)}\n"
        task_displays = [
            all_original_names.get(task, display_task_name(task))
            for task in sorted(all_tasks)
        ]
        dep_summary += f"**Tasks:** {', '.join(task_displays)}"

        # Try to solve
        result = solve_all_tasks(dependencies, all_tasks)

        if result:
            # All tasks can be executed
            output = "✅ **All tasks can be executed!**\n\n"
            output += "**Optimal execution order:**\n"

            result_display = []
            for i, task in enumerate(result, 1):
                task_display = all_original_names.get(task, display_task_name(task))
                output += f"{i}. {task_display}\n"
                result_display.append(task_display)

            json_output = json.dumps(result_display, indent=2)
            gantt = generate_mermaid_gantt(
                result, all_original_names, durations=durations
            )
            flowchart = generate_mermaid_flowchart(
                dependencies, all_tasks, all_original_names
            )
            gantt_raw = gantt.replace("```mermaid\n", "").rstrip("`\n")
            flow_raw = flowchart.replace("```mermaid\n", "").rstrip("`\n")

        else:
            # Try maximum subset
            result = solve_maximum_subset(dependencies, all_tasks)

            if result:
                excluded_tasks = set(all_tasks) - set(result)
                output = "⚠️ **Circular dependencies detected!**\n\n"
                output += (
                    f"**Maximum executable tasks ({len(result)}/{len(all_tasks)}):**\n"
                )

                result_display = []
                for i, task in enumerate(result, 1):
                    task_display = all_original_names.get(task, display_task_name(task))
                    output += f"{i}. {task_display}\n"
                    result_display.append(task_display)

                if excluded_tasks:
                    output += (
                        "\n**❌ Excluded tasks (due to circular dependencies):**\n"
                    )
                    for task in sorted(excluded_tasks):
                        task_display = all_original_names.get(
                            task, display_task_name(task)
                        )
                        output += f"• {task_display}\n"

                json_output = json.dumps(result_display, indent=2)
                gantt = generate_mermaid_gantt(
                    result, all_original_names, durations=durations
                )
                flowchart = generate_mermaid_flowchart(
                    dependencies, all_tasks, all_original_names
                )
                gantt_raw = gantt.replace("```mermaid\n", "").rstrip("`\n")
                flow_raw = flowchart.replace("```mermaid\n", "").rstrip("`\n")
            else:
                output = "❌ **No solution found!** There might be complex circular dependencies."
                json_output = "[]"
                gantt = generate_mermaid_gantt(
                    [], all_original_names, durations=durations
                )
                flowchart = generate_mermaid_flowchart(
                    dependencies, all_tasks, all_original_names
                )
                gantt_raw = gantt.replace("```mermaid\n", "").rstrip("`\n")
                flow_raw = flowchart.replace("```mermaid\n", "").rstrip("`\n")

        return output, dep_summary, json_output, gantt, flowchart, gantt_raw, flow_raw

    except Exception as e:
        return f"❌ **Error:** {str(e)}", "", "", "", "", "", ""


# Example tasks
example_tasks = """Sleep
Dinner
Toothbrushing
Prep
Shopping
Money
Clean dining room
Cleaning time"""

# Create Gradio interface
with gr.Blocks(title="Task Dependency Solver", theme=gr.themes.Soft()) as demo:
    gr.Markdown("""
    # 🔄 Task Dependency Solver
    
    Build task dependencies efficiently with multi-select requirements. Now with **case-insensitive** matching and **Mermaid.js visualizations**!
    
    **How to use:**
    1. Enter your tasks (one per line or comma-separated)
    2. (Optional) Enter the expected duration (in minutes) for each task.
    3. Select a task, then select all its requirements at once
    4. Click "Add Dependencies" to add them all
    5. Click "Solve Dependencies" to get the optimal execution order with visualizations
    6. You can copy and paste the generated Mermaid.js text below each diagram for use elsewhere!
    """)

    # State to store current dependencies
    dependencies_state = gr.State([])
    durations_state = gr.State(
        {}
    )  # dict: normalized_task -> int (user input durations)

    with gr.Row():
        with gr.Column(scale=1):
            gr.Markdown("### 📝 Step 1: Define Your Tasks")
            tasks_input = gr.Textbox(
                label="Tasks (one per line or comma-separated)",
                placeholder="Enter your tasks like:\nSleep\nDinner\nPrep\nShopping",
                lines=6,
                value=example_tasks,
                info="Case-insensitive: 'Sleep' and 'sleep' are treated the same",
            )

            gr.Markdown(
                "### ⏱️ (Optional) Step 1.5: Set Task Durations (minutes)", visible=False
            )
            durations_table = gr.Dataframe(
                headers=["Task", "Minutes"],
                datatype=["str", "number"],
                value=[],
                col_count=(2, "fixed"),
                row_count=(0, "dynamic"),
                label="Set duration (minutes) for each listed task.",
                interactive=True,
                wrap=True,
                visible=False,
            )

            gr.Markdown("### 🔗 Step 2: Build Dependencies")
            task_dropdown = gr.Dropdown(
                label="Select Task",
                choices=[],
                value="",
                info="Choose the task that has requirements",
            )

            requirements_checkbox = gr.CheckboxGroup(
                label="Select Requirements (can select multiple)",
                choices=[],
                value=[],
                info="Choose what the task requires",
            )

            add_btn = gr.Button("➕ Add Dependencies", variant="primary", scale=2)
            clear_all_btn = gr.Button("🗑️ Clear All", variant="secondary", scale=1)

    with gr.Row(scale=1):
        gr.Markdown("### 📋 Current Dependencies")
        dependencies_display = gr.Markdown("No dependencies added yet.")

        with gr.Accordion("🗑️ Remove Dependencies", open=False):
            remove_deps = gr.CheckboxGroup(
                label="Select dependencies to remove:", choices=[], value=[]
            )

            remove_btn = gr.Button("Remove Selected", variant="secondary")
            remove_status = gr.Markdown("")

    gr.Markdown("### 🚀 Step 3: Solve")
    solve_btn = gr.Button("🎯 Solve Dependencies", variant="primary", size="lg")

    with gr.Row():
        with gr.Column(scale=1):
            result_output = gr.Markdown()

        with gr.Accordion("📊 Analysis", open=False):
            dep_analysis = gr.Markdown()

        with gr.Accordion("💾 JSON Output", open=False):
            json_output = gr.Code(language="json")

    # Mermaid visualizations
    gr.Markdown("### 📊 Visualizations")

    gr.Markdown("#### 📅 Gantt Chart (Timeline)")
    gantt_mermaid = gr.Code(
        label="Copy-paste Gantt Mermaid.js", value="", language="markdown", interactive=True
    )
    gantt_chart = gr.Markdown()
    gantt_mermaid.change(
        fn=lambda x: x,  # Just to trigger the change event
        inputs=gantt_mermaid,
        outputs=gantt_chart,
    )
    gr.Markdown(
        "[🔗 Edit/view on Mermaid Live](https://mermaid.live/edit)", visible=True
    )

    gr.Markdown("#### 🔀 Dependency Flowchart")
    flow_mermaid = gr.Code(
        label="Copy-paste Flowchart Mermaid.js", value="", language="markdown", interactive=True
    )
    dependency_flowchart = gr.Markdown()
    flow_mermaid.change(
        fn=lambda x: x,  # Just to trigger the change event
        inputs=flow_mermaid,
        outputs=dependency_flowchart,
    )
    gr.Markdown(
        "[🔗 Edit/view on Mermaid Live](https://mermaid.live/edit)", visible=True
    )

    # Examples section
    with gr.Accordion("📚 Quick Examples", open=False):
        with gr.Row():
            morning_btn = gr.Button("🌅 Morning Routine")
            cooking_btn = gr.Button("🍳 Cooking Tasks")
            project_btn = gr.Button("💼 Project Tasks")

        gr.Markdown("""
        **Click the buttons above to load example task sets!**
        
        Each example includes pre-defined tasks. After loading, you can quickly build dependencies using the multi-select interface.
        
        **Note:** All matching is case-insensitive. You can mix cases freely!
        """)

    # Event handlers
    def update_ui_after_tasks_change(tasks_text: str, current_deps: List[str]) -> Tuple[gr.Dropdown, gr.CheckboxGroup, gr.CheckboxGroup]:
        """Update dropdowns and dependency checkboxes when tasks change."""
        tasks, original_names = parse_tasks(tasks_text)

        # Create display choices (original casing)
        task_display_choices = [""] + [
            original_names.get(task, display_task_name(task)) for task in tasks
        ]
        req_display_choices = [
            original_names.get(task, display_task_name(task)) for task in tasks
        ]

        # Update dependency removal checkboxes
        checkbox_choices = current_deps if current_deps else []

        return (
            gr.Dropdown(choices=task_display_choices, value=""),  # task_dropdown
            gr.CheckboxGroup(
                choices=req_display_choices, value=[]
            ),  # requirements_checkbox
            gr.CheckboxGroup(choices=checkbox_choices, value=[]),  # remove_deps
        )

    def handle_add_dependencies(task: str, required_tasks: List[str], current_deps: List[str]) -> Tuple[List[str], str, gr.CheckboxGroup, str, List[str]]:
        """Handle adding dependencies and update UI."""
        updated_deps, status = add_dependencies(task, required_tasks, current_deps)
        display = format_dependencies_display(updated_deps)
        checkbox_choices = updated_deps if updated_deps else []

        return (
            updated_deps,  # dependencies_state
            display,  # dependencies_display
            gr.CheckboxGroup(choices=checkbox_choices, value=[]),  # remove_deps
            "",  # task_dropdown (clear)
            [],  # requirements_checkbox (clear)
        )

    def handle_remove_dependencies(deps_to_remove: List[str], current_deps: List[str]) -> Tuple[List[str], str, str, gr.CheckboxGroup]:
        """Handle removing dependencies and update UI."""
        updated_deps, status = remove_dependency(deps_to_remove, current_deps)
        display = format_dependencies_display(updated_deps)
        checkbox_choices = updated_deps if updated_deps else []

        return (
            updated_deps,  # dependencies_state
            display,  # dependencies_display
            status,  # remove_status
            gr.CheckboxGroup(choices=checkbox_choices, value=[]),  # remove_deps
        )

    def clear_all_dependencies() -> Tuple[List[str], str, gr.CheckboxGroup]:
        """Clear all dependencies."""
        return (
            [],  # dependencies_state
            "No dependencies added yet.",  # dependencies_display
            gr.CheckboxGroup(choices=[], value=[]),  # remove_deps
        )

    # Example loaders
    def load_morning_example() -> Tuple[str, List[str], gr.CheckboxGroup]:
        tasks = "Wake_up\nShower\nBrush_teeth\nGet_dressed\nEat_breakfast\nMake_coffee\nLeave_house"
        return (
            tasks,
            [],
            gr.CheckboxGroup(choices=[], value=[]),
        )

    def load_cooking_example() -> Tuple[str, List[str], gr.CheckboxGroup]:
        tasks = "Shop_ingredients\nPrep_vegetables\nPreheat_oven\nCook_main_dish\nMake_sauce\nPlate_food\nServe_dinner"
        return (
            tasks,
            [],
            gr.CheckboxGroup(choices=[], value=[]),
        )

    def load_project_example() -> Tuple[str, List[str], gr.CheckboxGroup]:
        tasks = "Research\nPlan\nDesign\nDevelop\nTest\nDeploy\nDocument\nReview"
        return (
            tasks,
            [],
            gr.CheckboxGroup(choices=[], value=[]),
        )

    # Utility: update durations table to match tasks
    def update_durations_table(tasks_text: str, old_duration_rows: List[List[Any]]) -> List[List[str]]:
        tasks, original_names = parse_tasks(tasks_text)
        shown_task_set = set()
        updated_rows = []
        prev_map = {
            normalize_task_name(row[0]): row[1]
            for row in old_duration_rows
            if len(row) >= 2 and str(row[0]).strip()
        }
        for task in tasks:
            disp = original_names.get(task, display_task_name(task))
            shown_task_set.add(task)
            duration = prev_map.get(task, "")
            updated_rows.append([disp, duration])
        return updated_rows

    # Converts Dataframe rows to durations dict (normalized)
    def table_rows_to_durations(df: pandas.DataFrame) -> List[Dict[str, Any]]:
        return df.to_dict(orient="records")

    def update_durations_state(rows: pandas.DataFrame) -> List[Dict[str, Any]]:
        return table_rows_to_durations(rows)

    # Wire up events
    tasks_input.change(
        fn=update_ui_after_tasks_change,
        inputs=[tasks_input, dependencies_state],
        outputs=[task_dropdown, requirements_checkbox, remove_deps],
    )
    # Keep durations table in sync
    tasks_input.change(
        fn=update_durations_table,
        inputs=[tasks_input, durations_table],
        outputs=[durations_table],
    )
    # Whenever durations table is changed, update state and solve-display outputs
    # durations_table.change(
    #     fn=lambda rows: solve_dependencies(
    #         tasks_input.value, dependencies_state.value, table_rows_to_durations(rows)
    #     ),
    #     inputs=[durations_table],
    #     outputs=[
    #         result_output,
    #         dep_analysis,
    #         json_output,
    #         gantt_chart,
    #         dependency_flowchart,
    #         gantt_mermaid,
    #         flow_mermaid,
    #     ],
    # )

    add_btn.click(
        fn=handle_add_dependencies,
        inputs=[task_dropdown, requirements_checkbox, dependencies_state],
        outputs=[
            dependencies_state,
            dependencies_display,
            remove_deps,
            task_dropdown,
            requirements_checkbox,
        ],
    )

    remove_btn.click(
        fn=handle_remove_dependencies,
        inputs=[remove_deps, dependencies_state],
        outputs=[dependencies_state, dependencies_display, remove_status, remove_deps],
    )

    clear_all_btn.click(
        fn=clear_all_dependencies,
        outputs=[dependencies_state, dependencies_display, remove_deps],
    )

    solve_btn.click(
        fn=solve_dependencies,
        inputs=[tasks_input, dependencies_state, durations_state],
        outputs=[
            result_output,
            dep_analysis,
            json_output,
            gantt_mermaid,
            flow_mermaid,
        ],
    )

    # Example buttons
    morning_btn.click(
        fn=load_morning_example,
        outputs=[tasks_input, dependencies_state, dependencies_display, remove_deps],
    )

    cooking_btn.click(
        fn=load_cooking_example,
        outputs=[tasks_input, dependencies_state, dependencies_display, remove_deps],
    )

    project_btn.click(
        fn=load_project_example,
        outputs=[tasks_input, dependencies_state, dependencies_display, remove_deps],
    )

    # Initialize on load
    demo.load(
        fn=update_ui_after_tasks_change,
        inputs=[tasks_input, dependencies_state],
        outputs=[task_dropdown, requirements_checkbox, remove_deps],
    )
    # Also set durations table on load
    demo.load(
        fn=update_durations_table,
        inputs=[tasks_input, durations_table],
        outputs=[durations_table],
    )

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