# thermo_assistant.py
#
# A chemical engineering assistant that uses Ollama tool calling
# to compute vapor pressures with the Antoine equation.
#
# CHEG 667-013
# E. M. Furst

import math
from ollama import chat

# Antoine equation: log10(P) = A - B / (C + T)
# Returns vapor pressure in mmHg given temperature in degrees Celsius.

ANTOINE_CONSTANTS = {
    'water':    {'A': 8.07131, 'B': 1730.63,  'C': 233.426},
    'ethanol':  {'A': 8.20417, 'B': 1642.89,  'C': 230.300},
    'benzene':  {'A': 6.90565, 'B': 1211.033, 'C': 220.790},
    'toluene':  {'A': 6.95464, 'B': 1344.800, 'C': 219.482},
    'acetone':  {'A': 7.02447, 'B': 1161.0,   'C': 224.0},
}

def vapor_pressure(compound: str, temperature_C: float) -> str:
    """
    Calculate the vapor pressure of a compound using the Antoine equation.

    Args:
        compound: Name of the compound (water, ethanol, benzene, toluene, or acetone)
        temperature_C: Temperature in degrees Celsius

    Returns:
        The vapor pressure in mmHg, or an error message if the compound is unknown
    """
    compound = compound.lower().strip()
    if compound not in ANTOINE_CONSTANTS:
        return f"Unknown compound: {compound}. Available: {', '.join(ANTOINE_CONSTANTS.keys())}"

    c = ANTOINE_CONSTANTS[compound]
    log_p = c['A'] - c['B'] / (c['C'] + temperature_C)
    p_mmhg = 10 ** log_p
    return f"{p_mmhg:.2f} mmHg"


def available_compounds() -> str:
    """
    List the compounds available for vapor pressure calculations.

    Returns:
        A comma-separated list of compound names
    """
    return ', '.join(ANTOINE_CONSTANTS.keys())


# Tool dispatch table
tools = {
    'vapor_pressure': vapor_pressure,
    'available_compounds': available_compounds,
}

# System prompt: tell the model what it is
system_message = {
    'role': 'system',
    'content': (
        'You are a chemical engineering assistant. You have access to tools '
        'for computing vapor pressures using the Antoine equation. Use them '
        'when asked about vapor pressures or boiling behavior. Report results '
        'with units. If asked about a compound you do not have data for, say so.'
    ),
}


def ask(question: str) -> str:
    """Send a question to the assistant and return the response."""
    messages = [system_message, {'role': 'user', 'content': question}]

    response = chat(
        'llama3.1:8b',
        messages=messages,
        tools=[vapor_pressure, available_compounds],
    )

    # Handle tool calls (there may be more than one)
    while response.message.tool_calls:
        for tool_call in response.message.tool_calls:
            name = tool_call.function.name
            args = tool_call.function.arguments
            print(f'  [tool] {name}({args})')
            result = tools[name](**args)
            print(f'  [result] {result}')
            messages.append(response.message)
            messages.append({
                'role': 'tool',
                'content': str(result),
                'tool_name': name,
            })

        response = chat(
            'llama3.1:8b',
            messages=messages,
            tools=[vapor_pressure, available_compounds],
        )

    return response.message.content


# Interactive loop
if __name__ == '__main__':
    print('Thermo Assistant (type "quit" to exit)\n')
    while True:
        question = input('You: ')
        if question.strip().lower() in ('quit', 'exit', 'q'):
            break
        answer = ask(question)
        print(f'\nAssistant: {answer}\n')