from smolagents import CodeAgent,DuckDuckGoSearchTool, HfApiModel,load_tool,tool
import datetime
import requests
import pytz
import os
import yaml
import pandas as pd
from typing import Dict, Union, List
import json
from tools.final_answer import FinalAnswerTool

from Gradio_UI import GradioUI

# Define a tool to retrieve the Discounted Cash Flow (DCF) of a company

@tool
def get_dcf_of_company(symbol: str) -> str:
    """Retrieve the Discounted Cash Flow (DCF) of a company based on its stock symbol.
    
    Args:
        symbol: The stock symbol of the company (e.g., 'AAPL' for Apple, 'IBM' for IBM).
    
    Returns:
        str: A message with the DCF value or an error message if the calculation fails.
    """
    try:
        api_key = 'NR9AFISYRYH2B5U3'
        # Retrieve the company's financial statements
        financials = {}
        for statement in ['INCOME_STATEMENT', 'BALANCE_SHEET', 'CASH_FLOW']:
            url = f'https://www.alphavantage.co/query?function={statement}&symbol={symbol}&apikey={api_key}'
            response = requests.get(url)
            data = response.json()
            if 'annualReports' in data:
                financials[statement] = data['annualReports'][0]  # Use the most recent annual report
            else:
                return f"Error: Unable to retrieve {statement} for symbol '{symbol}'."

        # Extract necessary financial data
        income_statement = financials['INCOME_STATEMENT']
        cash_flow_statement = financials['CASH_FLOW']
        balance_sheet = financials['BALANCE_SHEET']

        net_income = float(income_statement.get('netIncome', 0))
        depreciation = float(cash_flow_statement.get('depreciationDepletionAndAmortization', 0))  # Corrected key
        capex = float(cash_flow_statement.get('capitalExpenditures', 0))
        change_in_working_capital = (
            float(balance_sheet.get('totalCurrentAssets', 0)) - float(balance_sheet.get('totalCurrentLiabilities', 0))
        )

        # Calculate Free Cash Flow (FCF)
        fcf = net_income + depreciation - capex - change_in_working_capital

        # Assume a growth rate and discount rate (WACC)
        growth_rate = 0.03  # 3% growth rate
        discount_rate = 0.10  # 10% discount rate
        years = 5  # Number of years to project

        # Calculate DCF
        dcf_value = 0
        for year in range(1, years + 1):
            future_cash_flow = fcf * ((1 + growth_rate) ** year)
            discounted_cash_flow = future_cash_flow / ((1 + discount_rate) ** year)
            dcf_value += discounted_cash_flow

        # Calculate Terminal Value
        terminal_value = (fcf * (1 + growth_rate)) / (discount_rate - growth_rate)
        discounted_terminal_value = terminal_value / ((1 + discount_rate) ** years)

        # Total DCF value
        total_value = dcf_value + discounted_terminal_value

        return f"The Discounted Cash Flow (DCF) for company '{symbol}' is: ${total_value:,.2f}"

    except Exception as e:
        return f"Error calculating DCF for company '{symbol}': {str(e)}"



final_answer = FinalAnswerTool()

# If the agent does not answer, the model is overloaded, please use another model or the following Hugging Face Endpoint that also contains qwen2.5 coder:
# model_id='https://pflgm2locj2t89co.us-east-1.aws.endpoints.huggingface.cloud' 

model = HfApiModel(
max_tokens=2096,
temperature=0.5,
model_id='https://pflgm2locj2t89co.us-east-1.aws.endpoints.huggingface.cloud',# it is possible that this model may be overloaded
custom_role_conversions=None,
)


# Import tool from Hub
image_generation_tool = load_tool("agents-course/text-to-image", trust_remote_code=True)

with open("prompts.yaml", 'r') as stream:
    prompt_templates = yaml.safe_load(stream)
    
agent = CodeAgent(
    model=model,
    tools=[final_answer,
          get_dcf_of_company], ## add your tools here (don't remove final answer)
    max_steps=6,
    verbosity_level=1,
    grammar=None,
    planning_interval=None,
    name=None,
    description=None,
    prompt_templates=prompt_templates
)


GradioUI(agent).launch()