# ==========================================
# examples/langchain_example.py
# ==========================================
"""
LangChain integration examples for pyCorpFin
Demonstrates how to use pyCorpFin tools with LangChain agents

Requirements:
    pip install pycorpfin[langchain]
    pip install langchain-openai

Set environment variable:
    export OPENAI_API_KEY='your-api-key'
"""

import os
from pyCorpFin.tools import get_langchain_tools

try:
    from langchain_openai import ChatOpenAI
    from langchain.agents import AgentExecutor, create_tool_calling_agent
    from langchain_core.prompts import ChatPromptTemplate

    LANGCHAIN_AVAILABLE = True
except ImportError:
    LANGCHAIN_AVAILABLE = False
    print("LangChain not installed. Install with: pip install pycorpfin[langchain] langchain-openai")


def example_simple_query():
    """Simple example: Calculate IRR with natural language"""
    print("=" * 80)
    print("Example 1: Simple IRR Calculation")
    print("=" * 80)

    tools = get_langchain_tools()
    llm = ChatOpenAI(model="gpt-4", temperature=0)

    prompt = ChatPromptTemplate.from_messages([
        ("system", "You are a financial analyst assistant. Use the available tools to help with calculations."),
        ("human", "{input}"),
        ("placeholder", "{agent_scratchpad}"),
    ])

    agent = create_tool_calling_agent(llm, tools, prompt)
    agent_executor = AgentExecutor(agent=agent, tools=tools, verbose=True)

    query = "I invested $10,000 and received $3,000 per year for 4 years. What's the IRR?"

    print(f"\nQuery: {query}\n")
    result = agent_executor.invoke({"input": query})
    print(f"\nResult: {result['output']}\n")


def example_loan_calculation():
    """Example: Loan EMI calculation"""
    print("=" * 80)
    print("Example 2: Loan EMI Calculation")
    print("=" * 80)

    tools = get_langchain_tools()
    llm = ChatOpenAI(model="gpt-4", temperature=0)

    prompt = ChatPromptTemplate.from_messages([
        ("system", "You are a financial advisor. Help users understand loan payments and provide clear explanations."),
        ("human", "{input}"),
        ("placeholder", "{agent_scratchpad}"),
    ])

    agent = create_tool_calling_agent(llm, tools, prompt)
    agent_executor = AgentExecutor(agent=agent, tools=tools, verbose=True)

    query = "I want to buy a house for $500,000 with a 30-year mortgage at 7.5% interest. What will my monthly payment be?"

    print(f"\nQuery: {query}\n")
    result = agent_executor.invoke({"input": query})
    print(f"\nResult: {result['output']}\n")


def example_investment_comparison():
    """Example: Compare multiple investments"""
    print("=" * 80)
    print("Example 3: Investment Comparison")
    print("=" * 80)

    tools = get_langchain_tools()
    llm = ChatOpenAI(model="gpt-4", temperature=0)

    prompt = ChatPromptTemplate.from_messages([
        ("system", "You are an investment analyst. Compare investments and provide recommendations."),
        ("human", "{input}"),
        ("placeholder", "{agent_scratchpad}"),
    ])

    agent = create_tool_calling_agent(llm, tools, prompt)
    agent_executor = AgentExecutor(agent=agent, tools=tools, verbose=True)

    query = """I have two investment options:
    Option A: Invest $10,000 now and receive $3,000 per year for 5 years
    Option B: Invest $10,000 now and receive $2,500, $3,000, $3,500, $4,000, $4,500 over 5 years

    Calculate the IRR for both and tell me which is better."""

    print(f"\nQuery: {query}\n")
    result = agent_executor.invoke({"input": query})
    print(f"\nResult: {result['output']}\n")


def example_xirr_irregular():
    """Example: XIRR for irregular cash flows"""
    print("=" * 80)
    print("Example 4: XIRR for Irregular Cash Flows")
    print("=" * 80)

    tools = get_langchain_tools()
    llm = ChatOpenAI(model="gpt-4", temperature=0)

    prompt = ChatPromptTemplate.from_messages([
        ("system", "You are a financial analyst specializing in return calculations."),
        ("human", "{input}"),
        ("placeholder", "{agent_scratchpad}"),
    ])

    agent = create_tool_calling_agent(llm, tools, prompt)
    agent_executor = AgentExecutor(agent=agent, tools=tools, verbose=True)

    query = """Calculate the XIRR for this investment:
    - Jan 1, 2023: Invested $10,000
    - July 15, 2023: Received $2,000
    - March 20, 2024: Received $4,000
    - Dec 31, 2024: Received $6,000
    """

    print(f"\nQuery: {query}\n")
    result = agent_executor.invoke({"input": query})
    print(f"\nResult: {result['output']}\n")


def example_complex_scenario():
    """Example: Complex multi-step financial analysis"""
    print("=" * 80)
    print("Example 5: Complex Financial Analysis")
    print("=" * 80)

    tools = get_langchain_tools()
    llm = ChatOpenAI(model="gpt-4", temperature=0)

    prompt = ChatPromptTemplate.from_messages([
        ("system", """You are a senior financial advisor. Provide comprehensive analysis with:
        1. Clear calculations using the tools
        2. Interpretation of results
        3. Practical recommendations"""),
        ("human", "{input}"),
        ("placeholder", "{agent_scratchpad}"),
    ])

    agent = create_tool_calling_agent(llm, tools, prompt)
    agent_executor = AgentExecutor(agent=agent, tools=tools, verbose=True)

    query = """I'm considering starting a business with these projections:
    - Initial investment: $50,000
    - Year 1 return: $10,000
    - Year 2 return: $15,000
    - Year 3 return: $20,000
    - Year 4 return: $25,000
    - Year 5 return: $30,000

    Also, I need to take out a loan of $30,000 at 8% for 5 years to supplement the investment.

    Please:
    1. Calculate the IRR of the business
    2. Calculate the monthly loan payment
    3. Tell me if this is a good investment considering a 12% discount rate"""

    print(f"\nQuery: {query}\n")
    result = agent_executor.invoke({"input": query})
    print(f"\nResult: {result['output']}\n")


def example_conversational():
    """Example: Multi-turn conversation"""
    print("=" * 80)
    print("Example 6: Conversational Agent")
    print("=" * 80)

    tools = get_langchain_tools()
    llm = ChatOpenAI(model="gpt-4", temperature=0)

    prompt = ChatPromptTemplate.from_messages([
        ("system", "You are a helpful financial assistant. Answer questions and perform calculations as needed."),
        ("placeholder", "{chat_history}"),
        ("human", "{input}"),
        ("placeholder", "{agent_scratchpad}"),
    ])

    agent = create_tool_calling_agent(llm, tools, prompt)
    agent_executor = AgentExecutor(agent=agent, tools=tools, verbose=True)

    # Simulate a conversation
    queries = [
        "What's the IRR if I invest $5,000 and get back $1,500 per year for 4 years?",
        "Now calculate the MIRR assuming a financing rate of 10% and reinvestment rate of 8%",
        "Is this a good investment if my required return is 15%?"
    ]

    chat_history = []
    for query in queries:
        print(f"\nUser: {query}\n")
        result = agent_executor.invoke({
            "input": query,
            "chat_history": chat_history
        })
        print(f"\nAssistant: {result['output']}\n")
        print("-" * 80)

        chat_history.append(("human", query))
        chat_history.append(("assistant", result['output']))


def main():
    """Run all examples"""
    if not LANGCHAIN_AVAILABLE:
        print("Please install required packages:")
        print("pip install pycorpfin[langchain] langchain-openai")
        return

    if not os.getenv("OPENAI_API_KEY"):
        print("Please set OPENAI_API_KEY environment variable")
        print("export OPENAI_API_KEY='your-api-key'")
        return

    print("\n" + "=" * 80)
    print("PYCORPFIN - LangChain Integration Examples")
    print("=" * 80 + "\n")

    try:
        example_simple_query()
        example_loan_calculation()
        example_investment_comparison()
        example_xirr_irregular()
        example_complex_scenario()
        example_conversational()

        print("\n" + "=" * 80)
        print("All examples completed!")
        print("=" * 80)

    except Exception as e:
        print(f"\nError running examples: {e}")
        print("\nMake sure you have:")
        print("1. Set OPENAI_API_KEY environment variable")
        print("2. Installed required packages: pip install pycorpfin[langchain] langchain-openai")


if __name__ == "__main__":
    main()