"""TimesFM (Time Series Foundation Model) integration for decline curve forecasting."""

import warnings

import numpy as np
import pandas as pd


def forecast_timesfm(series: pd.Series, horizon: int = 12) -> pd.Series:
    """
    Generate forecasts using Google's TimesFM model.

    Parameters:
    - series: Historical production data
    - horizon: Number of periods to forecast

    Returns:
    - Forecasted production series
    """
    try:
        # Try to import TimesFM
        import torch
        from transformers import AutoModelForSeq2SeqLM, AutoTokenizer  # noqa: F401

        # Check if CUDA is available
        device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

        # Load TimesFM model (this is a placeholder - actual implementation
        # would use Google's TimesFM)
        # For now, we'll use a simple transformer-based approach as a proxy
        model_name = "google/timesfm-1.0-200m"  # Placeholder model name

        try:
            model = AutoModelForSeq2SeqLM.from_pretrained(model_name)
            model.to(device)
            model.eval()
        except Exception:
            # Fallback to a simpler approach if TimesFM is not available
            warnings.warn("TimesFM model not available, using fallback method")
            return _fallback_timesfm_forecast(series, horizon)

        # Prepare input data
        values = series.values.astype(np.float32)

        # Normalize the data
        mean_val = np.mean(values)
        std_val = np.std(values) if np.std(values) > 0 else 1.0
        normalized_values = (values - mean_val) / std_val

        # Create input sequence (simplified approach)
        input_length = min(len(normalized_values), 512)  # Limit input length
        input_seq = normalized_values[-input_length:]

        # Generate forecast (this is a simplified implementation)
        with torch.no_grad():
            # For demonstration, we'll use a simple pattern-based forecast
            # In reality, this would use the actual TimesFM model architecture
            forecast_normalized = _generate_timesfm_forecast(input_seq, horizon)

        # Denormalize the forecast
        forecast_values = forecast_normalized * std_val + mean_val

        # Create forecast index
        last_date = series.index[-1]
        freq = series.index.freq or pd.infer_freq(series.index)
        _ = pd.date_range(start=last_date, periods=horizon + 1, freq=freq)[1:]

        # Combine historical and forecast data
        full_index = pd.date_range(
            start=series.index[0], periods=len(series) + horizon, freq=freq
        )
        full_forecast = np.concatenate([values, forecast_values])

        return pd.Series(full_forecast, index=full_index, name="timesfm_forecast")

    except ImportError:
        warnings.warn(
            "Transformers library not available for TimesFM. Using fallback method."
        )
        return _fallback_timesfm_forecast(series, horizon)
    except Exception as e:
        warnings.warn(f"TimesFM forecasting failed: {e}. Using fallback method.")
        return _fallback_timesfm_forecast(series, horizon)


def _generate_timesfm_forecast(input_seq: np.ndarray, horizon: int) -> np.ndarray:
    """Generate forecast using a simplified TimesFM-like approach.

    This is a placeholder implementation.
    """
    # Simple trend and seasonality extraction
    if len(input_seq) < 3:
        # Not enough data, use last value with slight decline
        return np.array([input_seq[-1] * (0.95**i) for i in range(1, horizon + 1)])

    # Calculate trend
    x = np.arange(len(input_seq))
    trend_coef = np.polyfit(x, input_seq, 1)[0]

    # Calculate recent volatility
    recent_changes = (
        np.diff(input_seq[-12:]) if len(input_seq) >= 12 else np.diff(input_seq)
    )
    volatility = np.std(recent_changes) if len(recent_changes) > 0 else 0.1

    # Generate forecast with trend and noise
    forecast = []
    last_value = input_seq[-1]

    for i in range(horizon):
        # Apply trend with some dampening
        trend_component = trend_coef * (i + 1) * 0.8  # Dampen trend over time

        # Add some realistic noise
        noise = np.random.normal(0, volatility * 0.5)

        # Ensure non-negative values for production data
        next_value = max(0, last_value + trend_component + noise)
        forecast.append(next_value)
        last_value = next_value

    return np.array(forecast)


def _fallback_timesfm_forecast(series: pd.Series, horizon: int) -> pd.Series:
    """Fallback forecasting method when TimesFM is not available.

    Uses exponential smoothing as a proxy.
    """
    from scipy.optimize import minimize_scalar

    values = series.values

    # Simple exponential smoothing
    def sse_alpha(alpha):
        forecast = [values[0]]
        for i in range(1, len(values)):
            forecast.append(alpha * values[i - 1] + (1 - alpha) * forecast[i - 1])
        return np.sum((np.array(forecast) - values) ** 2)

    # Optimize smoothing parameter
    result = minimize_scalar(sse_alpha, bounds=(0.01, 0.99), method="bounded")
    alpha = result.x

    # Generate forecast using optimized alpha
    # Apply exponential smoothing
    last_smoothed = values[-1]
    forecast_values = []

    for i in range(horizon):
        # Use optimized alpha for trend dampening
        damping = alpha * (0.98**i)  # Combine alpha with exponential decay
        next_value = last_smoothed * damping
        forecast_values.append(max(0, next_value))  # Ensure non-negative
        last_smoothed = next_value

    # Create full series (historical + forecast)
    freq = series.index.freq or pd.infer_freq(series.index)
    full_index = pd.date_range(
        start=series.index[0], periods=len(series) + horizon, freq=freq
    )
    full_values = np.concatenate([values, forecast_values])

    return pd.Series(full_values, index=full_index, name="timesfm_fallback")


def check_timesfm_availability() -> bool:
    """Check if TimesFM dependencies are available."""
    try:
        import torch  # noqa: F401
        import transformers  # noqa: F401

        return True
    except ImportError:
        return False
