"""Test LLM Math functionality."""

import pytest

from langchain.chains.llm_symbolic_math.base import LLMSymbolicMathChain
from langchain.chains.llm_symbolic_math.prompt import _PROMPT_TEMPLATE
from tests.unit_tests.llms.fake_llm import FakeLLM


@pytest.fixture
@pytest.mark.requires("sympy")
def fake_llm_symbolic_math_chain() -> LLMSymbolicMathChain:
    """Fake LLM Math chain for testing."""
    queries = {
        _PROMPT_TEMPLATE.format(question="What is 1 plus 1?"): "Answer: 2",
        _PROMPT_TEMPLATE.format(
            question="What is the square root of 2?"
        ): "```text\nsqrt(2)\n```",
        _PROMPT_TEMPLATE.format(
            question="What is the limit of sin(x) / x as x goes to 0?"
        ): "```text\nlimit(sin(x)/x,x,0)\n```",
        _PROMPT_TEMPLATE.format(
            question="What is the integral of e^-x from 0 to infinity?"
        ): "```text\nintegrate(exp(-x), (x, 0, oo))\n```",
        _PROMPT_TEMPLATE.format(
            question="What are the solutions to this equation x**2 - x?"
        ): "```text\nsolveset(x**2 - x, x)\n```",
        _PROMPT_TEMPLATE.format(question="foo"): "foo",
    }
    fake_llm = FakeLLM(queries=queries)
    return LLMSymbolicMathChain.from_llm(fake_llm, input_key="q", output_key="a")


@pytest.mark.requires("sympy")
def test_simple_question(fake_llm_symbolic_math_chain: LLMSymbolicMathChain) -> None:
    """Test simple question that should not need python."""
    question = "What is 1 plus 1?"
    output = fake_llm_symbolic_math_chain.run(question)
    assert output == "Answer: 2"


@pytest.mark.requires("sympy")
def test_root_question(fake_llm_symbolic_math_chain: LLMSymbolicMathChain) -> None:
    """Test irrational number that should need sympy."""
    import sympy

    question = "What is the square root of 2?"
    output = fake_llm_symbolic_math_chain.run(question)
    assert output == f"Answer: {sympy.sqrt(2)}"


@pytest.mark.requires("sympy")
def test_limit_question(fake_llm_symbolic_math_chain: LLMSymbolicMathChain) -> None:
    """Test question about limits that needs sympy"""
    question = "What is the limit of sin(x) / x as x goes to 0?"
    output = fake_llm_symbolic_math_chain.run(question)
    assert output == "Answer: 1"


@pytest.mark.requires("sympy")
def test_integration_question(
    fake_llm_symbolic_math_chain: LLMSymbolicMathChain,
) -> None:
    """Test question about integration that needs sympy"""
    question = "What is the integral of e^-x from 0 to infinity?"
    output = fake_llm_symbolic_math_chain.run(question)
    assert output == "Answer: 1"


@pytest.mark.requires("sympy")
def test_solver_question(fake_llm_symbolic_math_chain: LLMSymbolicMathChain) -> None:
    """Test question about solving algebraic equations that needs sympy"""
    question = "What are the solutions to this equation x**2 - x?"
    output = fake_llm_symbolic_math_chain.run(question)
    assert output == "Answer: {0, 1}"


@pytest.mark.requires("sympy")
def test_error(fake_llm_symbolic_math_chain: LLMSymbolicMathChain) -> None:
    """Test question that raises error."""
    with pytest.raises(ValueError):
        fake_llm_symbolic_math_chain.run("foo")