"""
CFA AI Agent - Finance Calculation Tools
This module contains various financial calculation functions for CFA analysis.
"""

import numpy as np
import pandas as pd
from typing import List, Dict, Union, Optional
import yfinance as yf
from langchain.tools import tool


@tool
def calculate_dcf(
    cash_flows: List[float],
    terminal_value: float,
    discount_rate: float
) -> Dict[str, float]:
    """
    Calculate Discounted Cash Flow (DCF) valuation.

    Args:
        cash_flows: List of projected free cash flows
        terminal_value: Terminal value at end of projection period
        discount_rate: Weighted average cost of capital (WACC) as decimal

    Returns:
        Dictionary with NPV, terminal value present value, and total enterprise value
    """
    try:
        if not cash_flows:
            raise ValueError("Cash flows list cannot be empty")
        if discount_rate <= 0:
            raise ValueError("Discount rate must be positive")

        # Calculate present value of cash flows
        pv_cash_flows = []
        for i, cf in enumerate(cash_flows, 1):
            pv = cf / ((1 + discount_rate) ** i)
            pv_cash_flows.append(pv)

        # Calculate present value of terminal value
        years = len(cash_flows)
        pv_terminal = terminal_value / ((1 + discount_rate) ** years)

        # Total enterprise value
        enterprise_value = sum(pv_cash_flows) + pv_terminal

        return {
            "pv_cash_flows": sum(pv_cash_flows),
            "pv_terminal_value": pv_terminal,
            "enterprise_value": enterprise_value,
            "cash_flow_details": pv_cash_flows
        }
    except Exception as e:
        return {"error": f"DCF calculation failed: {str(e)}"}


@tool
def calculate_sharpe_ratio(
    returns: List[float],
    risk_free_rate: float
) -> Dict[str, float]:
    """
    Calculate Sharpe Ratio for risk-adjusted returns.

    Args:
        returns: List of periodic returns (as decimals)
        risk_free_rate: Risk-free rate (as decimal)

    Returns:
        Dictionary with Sharpe ratio, average return, and standard deviation
    """
    try:
        if not returns:
            raise ValueError("Returns list cannot be empty")
        if len(returns) < 2:
            raise ValueError("Need at least 2 return observations")

        returns_array = np.array(returns)

        # Calculate metrics
        avg_return = np.mean(returns_array)
        std_dev = np.std(returns_array, ddof=1)  # Sample standard deviation
        excess_return = avg_return - risk_free_rate

        if std_dev == 0:
            raise ValueError("Standard deviation cannot be zero")

        sharpe_ratio = excess_return / std_dev

        return {
            "sharpe_ratio": sharpe_ratio,
            "average_return": avg_return,
            "standard_deviation": std_dev,
            "excess_return": excess_return,
            "risk_free_rate": risk_free_rate
        }
    except Exception as e:
        return {"error": f"Sharpe ratio calculation failed: {str(e)}"}


@tool
def compare_pe_ratios(ticker1: str, ticker2: str) -> Dict[str, Union[float, str]]:
    """
    Compare P/E ratios of two stocks using real-time data.

    Args:
        ticker1: First stock ticker symbol
        ticker2: Second stock ticker symbol

    Returns:
        Dictionary with P/E ratios and comparison analysis
    """
    try:
        # Fetch stock data
        stock1 = yf.Ticker(ticker1.upper())
        stock2 = yf.Ticker(ticker2.upper())

        # Get info
        info1 = stock1.info
        info2 = stock2.info

        # Extract P/E ratios
        pe1 = info1.get('trailingPE') or info1.get('forwardPE')
        pe2 = info2.get('trailingPE') or info2.get('forwardPE')

        if pe1 is None or pe2 is None:
            return {"error": f"Could not retrieve P/E ratios for {ticker1} or {ticker2}"}

        # Calculate comparison metrics
        pe_difference = pe1 - pe2
        pe_ratio = pe1 / pe2 if pe2 != 0 else None

        # Determine which is more expensive
        comparison = "equal"
        if pe1 > pe2:
            comparison = f"{ticker1} is more expensive"
        elif pe1 < pe2:
            comparison = f"{ticker2} is more expensive"

        return {
            f"{ticker1}_pe": pe1,
            f"{ticker2}_pe": pe2,
            "pe_difference": pe_difference,
            "pe_ratio": pe_ratio,
            "comparison": comparison,
            f"{ticker1}_name": info1.get('longName', ticker1),
            f"{ticker2}_name": info2.get('longName', ticker2)
        }
    except Exception as e:
        return {"error": f"P/E comparison failed: {str(e)}"}


@tool
def calculate_beta(ticker: str, market_ticker: str = "^GSPC", period: str = "2y") -> Dict[str, float]:
    """
    Calculate beta coefficient for a stock relative to market.

    Args:
        ticker: Stock ticker symbol
        market_ticker: Market index ticker (default S&P 500)
        period: Time period for calculation

    Returns:
        Dictionary with beta, correlation, and other metrics
    """
    try:
        # Fetch data
        stock = yf.Ticker(ticker.upper())
        market = yf.Ticker(market_ticker)

        # Get historical data
        stock_data = stock.history(period=period)
        market_data = market.history(period=period)

        if stock_data.empty or market_data.empty:
            raise ValueError("Could not fetch historical data")

        # Calculate returns
        stock_returns = stock_data['Close'].pct_change().dropna()
        market_returns = market_data['Close'].pct_change().dropna()

        # Align data
        aligned_data = pd.concat([stock_returns, market_returns], axis=1, join='inner')
        aligned_data.columns = ['stock', 'market']
        aligned_data = aligned_data.dropna()

        if len(aligned_data) < 20:
            raise ValueError("Insufficient data points for beta calculation")

        # Calculate beta
        covariance = np.cov(aligned_data['stock'], aligned_data['market'])[0][1]
        market_variance = np.var(aligned_data['market'], ddof=1)
        beta = covariance / market_variance

        # Calculate correlation
        correlation = np.corrcoef(aligned_data['stock'], aligned_data['market'])[0][1]

        return {
            "beta": beta,
            "correlation": correlation,
            "stock_volatility": np.std(aligned_data['stock'], ddof=1),
            "market_volatility": np.std(aligned_data['market'], ddof=1),
            "data_points": len(aligned_data),
            "period": period
        }
    except Exception as e:
        return {"error": f"Beta calculation failed: {str(e)}"}


@tool
def calculate_wacc(
    cost_of_equity: float,
    cost_of_debt: float,
    tax_rate: float,
    market_value_equity: float,
    market_value_debt: float
) -> Dict[str, float]:
    """
    Calculate Weighted Average Cost of Capital (WACC).

    Args:
        cost_of_equity: Cost of equity as decimal
        cost_of_debt: Cost of debt as decimal
        tax_rate: Corporate tax rate as decimal
        market_value_equity: Market value of equity
        market_value_debt: Market value of debt

    Returns:
        Dictionary with WACC and component calculations
    """
    try:
        total_value = market_value_equity + market_value_debt

        if total_value <= 0:
            raise ValueError("Total market value must be positive")

        # Calculate weights
        weight_equity = market_value_equity / total_value
        weight_debt = market_value_debt / total_value

        # Calculate after-tax cost of debt
        after_tax_cost_debt = cost_of_debt * (1 - tax_rate)

        # Calculate WACC
        wacc = (weight_equity * cost_of_equity) + (weight_debt * after_tax_cost_debt)

        return {
            "wacc": wacc,
            "weight_equity": weight_equity,
            "weight_debt": weight_debt,
            "after_tax_cost_debt": after_tax_cost_debt,
            "cost_of_equity": cost_of_equity,
            "cost_of_debt": cost_of_debt,
            "tax_rate": tax_rate
        }
    except Exception as e:
        return {"error": f"WACC calculation failed: {str(e)}"}


@tool
def financial_ratios_analysis(ticker: str) -> Dict[str, Union[float, str]]:
    """
    Perform comprehensive financial ratios analysis for a stock.

    Args:
        ticker: Stock ticker symbol

    Returns:
        Dictionary with various financial ratios and metrics
    """
    try:
        stock = yf.Ticker(ticker.upper())
        info = stock.info

        # Extract key metrics
        ratios = {
            "ticker": ticker.upper(),
            "company_name": info.get('longName', 'N/A'),

            # Valuation ratios
            "pe_ratio": info.get('trailingPE'),
            "forward_pe": info.get('forwardPE'),
            "price_to_book": info.get('priceToBook'),
            "price_to_sales": info.get('priceToSalesTrailing12Months'),
            "peg_ratio": info.get('pegRatio'),

            # Profitability ratios
            "profit_margin": info.get('profitMargins'),
            "operating_margin": info.get('operatingMargins'),
            "roe": info.get('returnOnEquity'),
            "roa": info.get('returnOnAssets'),

            # Financial health
            "current_ratio": info.get('currentRatio'),
            "quick_ratio": info.get('quickRatio'),
            "debt_to_equity": info.get('debtToEquity'),
            "total_debt": info.get('totalDebt'),
            "total_cash": info.get('totalCash'),

            # Market data
            "market_cap": info.get('marketCap'),
            "enterprise_value": info.get('enterpriseValue'),
            "beta": info.get('beta'),
            "52_week_high": info.get('fiftyTwoWeekHigh'),
            "52_week_low": info.get('fiftyTwoWeekLow'),

            # Dividend info
            "dividend_yield": info.get('dividendYield'),
            "payout_ratio": info.get('payoutRatio'),

            # Growth metrics
            "earnings_growth": info.get('earningsGrowth'),
            "revenue_growth": info.get('revenueGrowth')
        }

        # Remove None values
        ratios = {k: v for k, v in ratios.items() if v is not None}

        return ratios
    except Exception as e:
        return {"error": f"Financial ratios analysis failed: {str(e)}"}