# src/forecast_utils.py
import numpy as np
import pandas as pd

TRADING_DAYS = 252

def forecast_to_annual_return(forecast_series: pd.Series,
                              last_price: float,
                              method: str = "avg_daily",
                              cap: float = 2.0) -> float:
    """
    Convert a forecasted price series into an expected annual return.
    
    Args:
        forecast_series (pd.Series): Forecasted prices with DatetimeIndex.
        last_price (float): Last observed price before forecast starts.
        method (str): 'avg_daily' (default, safer) or 'horizon'.
        cap (float): Maximum allowed annualized return (as fraction, default 2.0 = 200%).
    
    Returns:
        float: Annualized expected return (capped).
    """
    forecast_series = forecast_series.dropna()
    if forecast_series.empty:
        raise ValueError("Forecast series is empty")

    if method == "horizon":
        horizon_price = forecast_series.iloc[-1]
        horizon_days = len(forecast_series)
        horizon_return = horizon_price / last_price - 1
        annual_return = (1 + horizon_return) ** (TRADING_DAYS / horizon_days) - 1
    elif method == "avg_daily":
        daily_returns = forecast_series.pct_change().dropna()
        if daily_returns.empty:
            raise ValueError("Not enough forecast points to compute daily returns")
        avg_daily = daily_returns.mean()
        annual_return = (1 + avg_daily) ** TRADING_DAYS - 1
    else:
        raise ValueError("method must be 'avg_daily' or 'horizon'")
    
    # Cap extreme values for stability
    annual_return = np.clip(annual_return, -cap, cap)
    return annual_return

import matplotlib.pyplot as plt

def plot_forecast_vs_actual(actual, forecast, title="TSLA Forecast vs Actual"):
    """
    Plot actual vs forecasted stock prices.
    """
    fig, ax = plt.subplots(figsize=(10, 5))
    ax.plot(actual.index, actual.values, label="Actual", color="blue")
    ax.plot(forecast.index, forecast.values, label="Forecast", color="red")
    ax.set_title(title)
    ax.set_xlabel("Date")
    ax.set_ylabel("Price ($)")
    ax.legend()
    ax.grid(True)
    return fig