import gradio as gr
import os

os.environ["OPENAI_API_KEY"] = os.getenv('api_key')

import math
import types
import uuid

from langchain.chat_models import init_chat_model
from langchain.embeddings import init_embeddings
from langgraph.store.memory import InMemoryStore

from langgraph_bigtool import create_agent
from langgraph_bigtool.utils import (
    convert_positional_only_function_to_tool
)

# Collect functions from `math` built-in
all_tools = []
for function_name in dir(math):
    function = getattr(math, function_name)
    if not isinstance(
        function, types.BuiltinFunctionType
    ):
        continue
    # This is an idiosyncrasy of the `math` library
    if tool := convert_positional_only_function_to_tool(
        function
    ):
        all_tools.append(tool)

# Create registry of tools. This is a dict mapping
# identifiers to tool instances.
tool_registry = {
    str(uuid.uuid4()): tool
    for tool in all_tools
}

# Index tool names and descriptions in the LangGraph
# Store. Here we use a simple in-memory store.
embeddings = init_embeddings("openai:text-embedding-3-small")

store = InMemoryStore(
    index={
        "embed": embeddings,
        "dims": 1536,
        "fields": ["description"],
    }
)
for tool_id, tool in tool_registry.items():
    store.put(
        ("tools",),
        tool_id,
        {
            "description": f"{tool.name}: {tool.description}",
        },
    )

# Initialize agent
llm = init_chat_model("openai:gpt-4o-mini")

builder = create_agent(llm, tool_registry)
agent = builder.compile(store=store)

from langchain_core.tools import Tool
import sympy
from sympy import symbols

def make_sympy_tool(func, name, description):
    def _tool(expr: str) -> str:
        local_symbols = symbols("x y z a b c n")
        parsed_expr = sympy.sympify(expr, locals={s.name: s for s in local_symbols})
        result = func(parsed_expr)
        return str(result)

    return Tool.from_function(
        name=name,
        description=description,
        func=_tool
    )

from sympy import simplify, expand, factor

sympy_tools = [
    make_sympy_tool(simplify, "simplify", "Simplifies a symbolic expression"),
    make_sympy_tool(expand, "expand", "Expands a symbolic expression"),
    make_sympy_tool(factor, "factor", "Factors a symbolic expression"),
]

for tool in sympy_tools:
    tool_id = str(uuid.uuid4())
    tool_registry[tool_id] = tool
    store.put(
        ("tools",),
        tool_id,
        {"description": f"{tool.name}: {tool.description}"},
    )

builder = create_agent(llm, tool_registry)
agent = builder.compile(store=store)

def pvsnp(problem):
    output = []
    for step in agent.stream(
    {"messages": "Use tools to answer:"+problem},
    stream_mode="updates",
):
        for _, update in step.items():
           for message in update.get("messages", []):
              message.pretty_print()
              output.append(message.pretty_print())
    return output
                                      

iface = gr.Interface(
    fn=pvsnp, 
    inputs=gr.Textbox(label="What problem would you like to classify as P or NP?"), 
    outputs=gr.Markdown(label="Agent's response"),  # Output as HTML
    title="PolyMath",
    description="PolyMath is an advanced AI agent that guides users through the intricate maze of computational complexity. This agent scrutinizes problem descriptions with sophisticated LLM prompts and symbolic reasoning. It classifies problems into categories such as P, NP, NP-complete, NP-hard, or beyond (e.g., PSPACE, EXPTIME), while providing clear, concise explanations of its reasoning. As part of AI Quotient’s Millennium Math Challenge, it is the first step towards solving the P vs NP problem.",
    theme = gr.themes.Ocean(),
    examples = ["Simplify x*2+2x+1"]
)

# Launch the app
iface.launch()