from langchain.tools import Tool
from langchain.utilities import WikipediaAPIWrapper, ArxivAPIWrapper, DuckDuckGoSearchRun
import math

## ----- TOOLS DEFINITION ----- ##

# ** Math Tools ** #

def add_numbers(a: float, b: float) -> float:
    """
    Add two floating-point numbers.

    Args:
        a (float): The first number.
        b (float): The second number.

    Returns:
        float: The result of the addition.
    """
    return a + b

def subtract_numbers(a: float, b: float) -> float:
    """
    Subtract the second floating-point number from the first.

    Args:
        a (float): The first number.
        b (float): The second number.

    Returns:
        float: The result of the subtraction.
    """
    return a - b

def multiply_numbers(a: float, b: float) -> float:
    """
    Multiply two floating-point numbers.

    Args:
        a (float): The first number.
        b (float): The second number.

    Returns:
        float: The result of the multiplication.
    """
    return a * b

def divide_numbers(a: float, b: float) -> float:
    """
    Divide the first floating-point number by the second.

    Args:
        a (float): The numerator.
        b (float): The denominator.

    Returns:
        float: The result of the division.

    Raises:
        ValueError: If division by zero is attempted.
    """
    if b == 0:
        raise ValueError("Division by zero")
    return a / b

def power(a: float, b: float) -> float:
    """
    Raise the first number to the power of the second.

    Args:
        a (float): The base.
        b (float): The exponent.

    Returns:
        float: The result of the exponentiation.
    """
    return a ** b

def modulus(a: float, b: float) -> float:
    """
    Compute the modulus (remainder) of the division of a by b.

    Args:
        a (float): The dividend.
        b (float): The divisor.

    Returns:
        float: The remainder after division.
    """
    return a % b

def square_root(a: float) -> float:
    """
    Compute the square root of a number.

    Args:
        a (float): The number.

    Returns:
        float: The square root.

    Raises:
        ValueError: If a is negative.
    """
    if a < 0:
        raise ValueError("Cannot compute square root of a negative number")
    return math.sqrt(a)

def logarithm(a: float, base: float = math.e) -> float:
    """
    Compute the logarithm of a number with a specified base.

    Args:
        a (float): The number.
        base (float, optional): The logarithmic base (default is natural log).

    Returns:
        float: The logarithm.

    Raises:
        ValueError: If a or base is not positive.
    """
    if a <= 0 or base <= 0:
        raise ValueError("Logarithm arguments must be positive")
    return math.log(a, base)

# ** Search Tools ** #

# DuckDuckGo Web Search
duckduckgo_search = DuckDuckGoSearchRun()
web_search_tool = Tool.from_function(
    func=duckduckgo_search.run,
    name="Web Search",
    description="Use this tool to search the internet for general-purpose queries."
)

# Wikipedia Search
wikipedia_search = WikipediaAPIWrapper()
wikipedia_tool = Tool.from_function(
    func=wikipedia_search.run,
    name="Wikipedia Search",
    description="Use this tool to search Wikipedia for factual or encyclopedic information."
)

# ArXiv Search
arxiv_search = ArxivAPIWrapper()
arxiv_tool = Tool.from_function(
    func=arxiv_search.run,
    name="ArXiv Search",
    description="Use this tool to search ArXiv for scientific papers. Input should be a research topic or query."
)

## ----- TOOLS LIST ----- ##

tools = [
    # Math
    Tool.from_function(func=add_numbers, name="Add Numbers", description="Add two numbers."),
    Tool.from_function(func=subtract_numbers, name="Subtract Numbers", description="Subtract two numbers."),
    Tool.from_function(func=multiply_numbers, name="Multiply Numbers", description="Multiply two numbers."),
    Tool.from_function(func=divide_numbers, name="Divide Numbers", description="Divide two numbers."),
    Tool.from_function(func=power, name="Power", description="Raise one number to the power of another."),
    Tool.from_function(func=modulus, name="Modulus", description="Compute the modulus (remainder) of a division."),
    Tool.from_function(func=square_root, name="Square Root", description="Compute the square root of a number."),
    Tool.from_function(func=logarithm, name="Logarithm", description="Compute the logarithm of a number with a given base."),
    # Search
    web_search_tool,
    wikipedia_tool,
    arxiv_tool