import gradio as gr
import json
from collections import defaultdict

# -----------------------------
# Goal Setting Agent Algorithm
# -----------------------------
def goal_setting_agent(goals_text, priorities_text, dependencies_text, total_time):

    # Parse inputs
    G = [g.strip() for g in goals_text.split(",")]
    P = list(map(int, priorities_text.split(",")))

    # Dependencies format: g1:g2 means g1 depends on g2
    DG = defaultdict(list)
    if dependencies_text.strip():
        deps = dependencies_text.split(",")
        for d in deps:
            g1, g2 = d.split(":")
            DG[g1.strip()].append(g2.strip())

    # Step 1: Sort goals by priority
    goals_sorted = [g for _, g in sorted(zip(P, G), reverse=True)]

    # Step 2: Ensure dependencies satisfied (simple check)
    ordered_goals = []
    visited = set()

    def visit(g):
        if g in visited:
            return
        for dep in DG[g]:
            visit(dep)
        visited.add(g)
        ordered_goals.append(g)

    for g in goals_sorted:
        visit(g)

    # Step 3: Decompose into milestones
    action_plan = {}
    for g in ordered_goals:
        action_plan[g] = [
            f"{g} - Milestone 1",
            f"{g} - Milestone 2",
            f"{g} - Milestone 3"
        ]

    # Step 4: Assign duration
    time_per_goal = int(total_time) // len(ordered_goals)
    schedule = {}

    current_time = 0
    for g in ordered_goals:
        schedule[g] = {
            "start": current_time,
            "end": current_time + time_per_goal,
            "duration": time_per_goal
        }
        current_time += time_per_goal

    # Step 5: Adjust constraints (basic)
    # (can be extended for real-world constraints)
   
    return (
        json.dumps(schedule, indent=2),
        json.dumps(DG, indent=2),
        json.dumps(action_plan, indent=2)
    )


# -----------------------------
# Gradio UI
# -----------------------------
interface = gr.Interface(
    fn=goal_setting_agent,
    inputs=[
        gr.Textbox(
            label="Health Goals (comma-separated)",
            value="Reduce Vata imbalance, Improve digestion (Pitta), Enhance sleep quality, Balance Kapha"
        ),
        gr.Textbox(
            label="Priorities (comma-separated, higher = important)",
            value="3, 2, 1, 2"
        ),
        gr.Textbox(
            label="Dependencies (goal1:goal2 format, comma-separated)",
            value="Enhance sleep quality:Reduce Vata imbalance, Improve digestion (Pitta):Balance Kapha"
        ),
        gr.Number(
            label="Total Time Available (in days)",
            value=21
        )
    ],
    outputs=[
        gr.Code(label="Personalized Schedule (S)", language="json"),
        gr.Code(label="Dependency Graph (DG)", language="json"),
        gr.Code(label="Ayurvedic Action Plan", language="json")
    ],
    title="Ayurvedic Goal Setting Agent (Algorithm 4.2)",
    description="Generates a personalized healing plan based on dosha imbalance, priorities, and dependencies."
)

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