src/forecast.py

"""
TimesFM Forecasting Module

This module provides a simplified and robust interface for TimesFM forecasting,
handling both basic and covariates-enhanced forecasting with consistent quantile output.

Key Features:
- Single forecast method with optional covariates
- Always returns quantiles (never "maybe")
- Simplified logic: IF covariates -> use covariates, ELSE -> use basic
- Consistent return format: (point_forecast, quantile_forecast)
"""

import numpy as np
import pandas as pd
import logging
from typing import List, Dict, Optional, Tuple, Any, Union
import timesfm

logger = logging.getLogger(__name__)


class Forecaster:
    """
    Simplified TimesFM Forecaster with consistent quantile output.
    
    This class provides a single forecast method that handles both basic and
    covariates-enhanced forecasting, always returning quantiles.
    
    Example:
        >>> forecaster = Forecaster(model)
        >>> point_forecast, quantile_forecast = forecaster.forecast(
        ...     inputs=[1,2,3,4,5], 
        ...     use_covariates=True,
        ...     dynamic_numerical_covariates={'feature1': [[1,2,3,4,5]]}
        ... )
    """
    
    def __init__(self, model: timesfm.TimesFm):
        """
        Initialize the Forecaster with a loaded TimesFM model.
        
        Args:
            model: Initialized TimesFM model instance
        """
        self.model = model
        self.capabilities = self._detect_capabilities()
        logger.info(f"Forecaster initialized with capabilities: {list(self.capabilities.keys())}")
    
    def _detect_capabilities(self) -> Dict[str, bool]:
        """Detect available forecasting capabilities of the model."""
        return {
            'basic_forecasting': True,
            'quantile_forecasting': hasattr(self.model, 'experimental_quantile_forecast'),
            'covariates_support': hasattr(self.model, 'forecast_with_covariates')
        }
    
    def forecast(
        self,
        inputs: Union[List[float], List[List[float]]],
        freq: Union[int, List[int]] = 0,
        dynamic_numerical_covariates: Optional[Dict[str, List[List[float]]]] = None,
        dynamic_categorical_covariates: Optional[Dict[str, List[List[str]]]] = None,
        static_numerical_covariates: Optional[Dict[str, List[float]]] = None,
        static_categorical_covariates: Optional[Dict[str, List[str]]] = None,
        use_covariates: bool = False,
        xreg_mode: str = "xreg + timesfm",
        ridge: float = 0.0,
        normalize_xreg_target_per_input: bool = True
    ) -> Tuple[np.ndarray, np.ndarray]:
        """
        Perform TimesFM forecasting with optional covariates support.
        
        This is the main forecasting method that handles both basic and covariates-enhanced
        forecasting. Quantiles are always returned regardless of covariates usage.
        
        Args:
            inputs: Input time series data
            freq: Frequency indicator(s)
            dynamic_numerical_covariates: Dynamic numerical covariates (if use_covariates=True)
            dynamic_categorical_covariates: Dynamic categorical covariates (if use_covariates=True)
            static_numerical_covariates: Static numerical covariates (if use_covariates=True)
            static_categorical_covariates: Static categorical covariates (if use_covariates=True)
            use_covariates: Whether to use covariates-enhanced forecasting
            xreg_mode: Covariate integration mode ("xreg + timesfm" or "timesfm + xreg")
            ridge: Ridge regression parameter for covariates
            normalize_xreg_target_per_input: Whether to normalize covariates
            
        Returns:
            Tuple of (point_forecast, quantile_forecast) - both are always returned
            
        Raises:
            ValueError: If covariates are requested but not supported
            Exception: If forecasting fails
        """
        logger.info(f"Performing TimesFM forecasting (covariates={use_covariates})...")
        
        # Normalize inputs format
        if isinstance(inputs[0], (int, float)):
            # inputs is a single list of numbers
            inputs_norm = [inputs]
        else:
            # inputs is already a list of lists
            inputs_norm = inputs
            
        if isinstance(freq, int):
            freq_norm = [freq] * len(inputs_norm)
        else:
            freq_norm = freq

        try:
            if use_covariates and any([
                dynamic_numerical_covariates, dynamic_categorical_covariates,
                static_numerical_covariates, static_categorical_covariates
            ]):
                # Validate covariates support
                if not self.capabilities['covariates_support']:
                    raise ValueError("Model does not support covariates forecasting")
                
                # Validate covariates data structure
                self._validate_covariates(
                    inputs_norm, dynamic_numerical_covariates, dynamic_categorical_covariates,
                    static_numerical_covariates, static_categorical_covariates
                )
                
                logger.info(f"Using covariates-enhanced forecasting (mode: {xreg_mode})...")
                logger.info(f"Inputs shape: {[len(x) for x in inputs] if isinstance(inputs[0], list) else len(inputs)}")
                logger.info(f"Inputs type: {type(inputs)}")
                
                # Perform covariates forecasting with original mode
                covariates_result = self.model.forecast_with_covariates(
                    inputs=inputs_norm,
                    dynamic_numerical_covariates=dynamic_numerical_covariates or {},
                    dynamic_categorical_covariates=dynamic_categorical_covariates or {},
                    static_numerical_covariates=static_numerical_covariates or {},
                    static_categorical_covariates=static_categorical_covariates or {},
                    freq=freq_norm,
                    xreg_mode=xreg_mode,
                    ridge=ridge,
                    normalize_xreg_target_per_input=normalize_xreg_target_per_input
                )
                
                # Handle return format from forecast_with_covariates
                if isinstance(covariates_result, tuple) and len(covariates_result) == 2:
                    point_forecast, quantile_forecast = covariates_result
                    point_forecast = np.array(point_forecast)
                    quantile_forecast = np.array(quantile_forecast)
                    
                    logger.info(f"✅ Covariates forecasting completed.")
                    logger.info(f"  Point forecast shape: {point_forecast.shape}")
                    logger.info(f"  Quantile forecast shape: {quantile_forecast.shape}")
                    
                    # Check if we have proper quantiles (multiple quantiles, not just 1)
                    if quantile_forecast.ndim == 2 and (quantile_forecast.shape[0] == 1 or quantile_forecast.shape[1] == 1):
                        logger.warning("⚠️ Covariates forecasting returned insufficient quantiles, falling back to basic forecast for quantiles")
                        # Get quantiles from basic forecast method
                        _, quantile_forecast = self.model.forecast(inputs=inputs_norm, freq=freq_norm)
                        quantile_forecast = np.array(quantile_forecast)
                        logger.info(f"✅ Basic forecast quantiles obtained. Shape: {quantile_forecast.shape}")
                    else:
                        logger.info("✅ Using quantiles from covariates forecasting")
                else:
                    # Fallback: If forecast_with_covariates doesn't return quantiles, get them separately
                    logger.warning("⚠️ Covariates forecasting didn't return quantiles, getting them separately")
                    point_forecast = np.array(covariates_result)
                    _, quantile_forecast = self.model.forecast(inputs=inputs_norm, freq=freq_norm)
                    quantile_forecast = np.array(quantile_forecast)
            
            else:
                logger.info("Using basic forecasting...")
                
                # Perform basic forecasting - this should return (point, quantiles)
                point_forecast, quantile_forecast = self.model.forecast(inputs=inputs_norm, freq=freq_norm)
                point_forecast = np.array(point_forecast)
                quantile_forecast = np.array(quantile_forecast)
                
                logger.info(f"✅ Basic forecasting completed.")
            
            return point_forecast, quantile_forecast
            
        except Exception as e:
            logger.error(f"❌ Forecasting failed: {str(e)}")
            raise
    
    def _validate_covariates(
        self,
        inputs: List[List[float]],
        dynamic_numerical: Optional[Dict],
        dynamic_categorical: Optional[Dict],
        static_numerical: Optional[Dict],
        static_categorical: Optional[Dict]
    ) -> None:
        """Validate covariates data structure and compatibility."""
        logger.info("Validating covariates data structure...")
        
        # Check that all covariates have the same number of series as inputs
        num_series = len(inputs)
        
        for cov_type, cov_data in [
            ("dynamic_numerical", dynamic_numerical),
            ("dynamic_categorical", dynamic_categorical),
            ("static_numerical", static_numerical),
            ("static_categorical", static_categorical)
        ]:
            if cov_data:
                for name, data in cov_data.items():
                    if isinstance(data[0], (list, np.ndarray)):
                        # Dynamic covariates
                        if len(data) != num_series:
                            raise ValueError(f"Dynamic covariate '{name}' has {len(data)} series, expected {num_series}")
                    else:
                        # Static covariates
                        if len(data) != num_series:
                            raise ValueError(f"Static covariate '{name}' has {len(data)} values, expected {num_series}")
        
        logger.info("✅ Covariates validation passed")
    

def run_forecast(
    forecaster: 'Forecaster',
    target_inputs: List[List[float]],
    covariates: Optional[Dict[str, Any]] = None,
    use_covariates: bool = False,
    freq: Union[int, List[int]] = 0
) -> Dict[str, Any]:
    """
    Centralized forecasting function that handles both basic and covariates-enhanced forecasting.
    
    This function implements the logic to decide whether to run forecast_with_covariates
    or the basic forecast, including fallback mechanisms and proper error handling.
    
    Args:
        forecaster: Initialized Forecaster instance
        target_inputs: Input time series data
        covariates: Dictionary containing covariate data (if use_covariates=True)
        use_covariates: Whether to use covariates-enhanced forecasting
        freq: Frequency indicator(s)
        
    Returns:
        Dictionary containing forecast results with keys:
        - 'enhanced_forecast' or 'point_forecast': Main forecast array
        - 'quantile_forecast': Quantile forecast array (always present)
        - 'method': String indicating the forecasting method used
        - 'metadata': Additional forecast metadata
        
    Raises:
        Exception: If forecasting fails
    """
    logger.info(f"🚀 Running centralized forecast (covariates={use_covariates})...")
    
    try:
        results = {}
        
        if use_covariates and covariates:
            logger.info("Using covariates-enhanced forecasting...")
            
            # Extract covariate data
            dynamic_numerical = covariates.get('dynamic_numerical_covariates')
            dynamic_categorical = covariates.get('dynamic_categorical_covariates')
            static_numerical = covariates.get('static_numerical_covariates')
            static_categorical = covariates.get('static_categorical_covariates')
            
            # Perform covariates forecasting
            point_forecast, quantile_forecast = forecaster.forecast(
                inputs=target_inputs,
                freq=freq,
                dynamic_numerical_covariates=dynamic_numerical,
                dynamic_categorical_covariates=dynamic_categorical,
                static_numerical_covariates=static_numerical,
                static_categorical_covariates=static_categorical,
                use_covariates=True
            )
            
            results['point_forecast'] = point_forecast
            results['method'] = 'covariates_enhanced'
            
        else:
            logger.info("Using basic forecasting...")
            
            # Perform basic forecasting
            point_forecast, quantile_forecast = forecaster.forecast(
                inputs=target_inputs,
                freq=freq,
                use_covariates=False
            )
            
            results['point_forecast'] = point_forecast
            results['method'] = 'basic_timesfm'
        
        # Check for NaN values before returning
        if np.any(np.isnan(point_forecast)):
            logger.error(f"❌ NaN values detected in point_forecast: {np.isnan(point_forecast).sum()} out of {point_forecast.size}")
            logger.error(f"Point forecast values: {point_forecast}")
            raise ValueError(f"Forecasting produced NaN values in point forecast. This may be due to insufficient data or model issues.")
        
        if np.any(np.isnan(quantile_forecast)):
            logger.error(f"❌ NaN values detected in quantile_forecast: {np.isnan(quantile_forecast).sum()} out of {quantile_forecast.size}")
            logger.error(f"Quantile forecast shape: {quantile_forecast.shape}")
            raise ValueError(f"Forecasting produced NaN values in quantile forecast. This may be due to insufficient data or model issues.")
        
        # Quantiles are always available
        results['quantile_forecast'] = quantile_forecast
        logger.info(f"✅ Quantile forecast obtained. Shape: {quantile_forecast.shape}")
        
        # Add metadata
        results['metadata'] = {
            'input_series_count': len(target_inputs),
            'forecast_length': results.get('point_forecast').shape[-1],
            'covariates_used': use_covariates and covariates is not None,
            'quantiles_available': True  # Always true now
        }
        
        logger.info(f"✅ Centralized forecast completed successfully!")
        logger.info(f"   Method: {results['method']}")
        logger.info(f"   Forecast shape: {results['metadata']['forecast_length']}")
        logger.info(f"   Quantiles: Yes (shape: {quantile_forecast.shape})")
        logger.info(f"   Point forecast range: {np.min(point_forecast):.2f} to {np.max(point_forecast):.2f}")
        
        return results
        
    except Exception as e:
        logger.error(f"❌ Centralized forecasting failed: {str(e)}")
        raise


def process_quantile_bands(
    quantile_forecast: np.ndarray,
    selected_indices: List[int] = None
) -> Dict[str, Any]:
    """
    Centralized function to process quantile forecasts into quantile bands for visualization.
    
    This function contains the logic for sorting quantiles and creating the quantile band
    dictionary, as used in both the webapp and notebook.
    
    Args:
        quantile_forecast: Array of quantile forecasts with shape (horizon, num_quantiles) or (num_quantiles, horizon)
        selected_indices: List of quantile indices to use for bands (default: [1, 3, 5, 7, 9])
        
    Returns:
        Dictionary of quantile bands ready for visualization with keys:
        - 'quantile_band_0_lower', 'quantile_band_0_upper', 'quantile_band_0_label'
        - 'quantile_band_1_lower', 'quantile_band_1_upper', 'quantile_band_1_label'
        - etc.
    """
    logger.info("🔄 Processing quantile bands...")
    logger.info(f"Input quantile_forecast type: {type(quantile_forecast)}")
    logger.info(f"Input quantile_forecast shape: {quantile_forecast.shape if hasattr(quantile_forecast, 'shape') else 'N/A'}")

    # logger.info(f"!!!!!!!!!!!!! selected_indices: {selected_indices}")
    # logger.info(f"!!!!!!!!!!!!! quantile_forecast.shape: {quantile_forecast.shape}")
    
    if quantile_forecast is None:
        logger.warning("No quantile forecast provided")
        return {}
    
    try:
        # logger.info(f"!!!!!!!!!!!!! selected_indices: {selected_indices}")
        # logger.info(f"!!!!!!!!!!!!! quantile_forecast.shape: {quantile_forecast.shape}")

        # Handle quantile indices - only use default if explicitly None (not empty list)
        if selected_indices is None:
            # This means no quantile selection was made, use default
            selected_indices = [1, 3, 5, 7, 9]  # Q10, Q30, Q50, Q70, Q90
        elif selected_indices == []:
            # This means user explicitly selected no quantiles, return empty
            logger.info("No quantiles selected by user - returning empty quantile bands")
            return {}
        
        # Handle different array dimensions
        if quantile_forecast.ndim == 3:
            # Shape is (1, horizon, num_quantiles) - squeeze out first dimension
            q_mat = quantile_forecast.squeeze(0)
            logger.info(f"3D array detected, squeezed to shape: {q_mat.shape}")
        elif quantile_forecast.ndim == 1:
            # Shape is (horizon,) - reshape to (1, horizon)
            q_mat = quantile_forecast.reshape(1, -1)
            logger.info(f"1D array detected, reshaped to: {q_mat.shape}")
        else:
            # Shape is 2D - determine if we need to transpose
            # For quantiles, we expect (horizon, num_quantiles) format
            # If we have more horizon than quantiles, it's likely (horizon, num_quantiles) and should be kept as-is
            if quantile_forecast.shape[0] > quantile_forecast.shape[1]:
                # Shape is (horizon, num_quantiles) - keep as is
                q_mat = quantile_forecast
                logger.info(f"2D array kept as is (horizon, quantiles): {q_mat.shape}")
            else:
                # Shape is (num_quantiles, horizon) - transpose to (horizon, num_quantiles)
                q_mat = quantile_forecast.T
                logger.info(f"2D array transposed from {quantile_forecast.shape} to {q_mat.shape}")
        
        horizon_len, num_quantiles = q_mat.shape
        logger.info(f"📊 Available quantiles: {num_quantiles} (indices 0-{num_quantiles-1})")
        logger.info(f"📊 Note: Index 0 is legacy mean forecast, using indices 1-{num_quantiles-1} for actual quantiles")
        
        # Check if we have enough quantiles for band creation (need at least 3 total: 0=legacy, 1=Q10, 2=Q20)
        if num_quantiles < 3:
            logger.warning(f"Not enough quantiles for band creation. Have {num_quantiles}, need at least 3")
            return {}
        
        # Filter selected indices to valid range (skip index 0)
        valid_indices = [idx for idx in selected_indices if 1 <= idx < num_quantiles]  # Skip index 0
        if not valid_indices:
            logger.warning("No valid quantile indices selected (after skipping legacy index 0)")
            return {}
        
        # logger.info(f"!!!!!!!!!!!!! valid_indices: {valid_indices}")
        
        # Sort quantiles by their median magnitude to ensure proper ordering
        quantile_medians = np.median(q_mat, axis=0)
        sorted_indices = np.argsort(quantile_medians)
        
        # Create quantile bands from selected indices
        quantile_bands = {}
        band_count = 0
        
        for i in range(len(valid_indices) - 1):
            lower_idx = valid_indices[i]
            upper_idx = valid_indices[i + 1]
            
            # Get the sorted indices for these quantiles
            lower_sorted_idx = sorted_indices[lower_idx]
            upper_sorted_idx = sorted_indices[upper_idx]
            
            # Extract quantile values
            lower_quantile = q_mat[:, lower_sorted_idx]
            upper_quantile = q_mat[:, upper_sorted_idx]
            
            # Create band labels
            lower_pct = idx_to_percent(lower_idx, num_quantiles)
            upper_pct = idx_to_percent(upper_idx, num_quantiles)
            band_label = f"Q{lower_pct:02d}–Q{upper_pct:02d}"
            
            # Store band data
            quantile_bands[f'quantile_band_{band_count}_lower'] = lower_quantile.tolist()
            quantile_bands[f'quantile_band_{band_count}_upper'] = upper_quantile.tolist()
            quantile_bands[f'quantile_band_{band_count}_label'] = band_label
            
            logger.info(f"   Band {band_count}: {band_label} - Lower: {len(lower_quantile)}, Upper: {len(upper_quantile)}")
            band_count += 1
        
        logger.info(f"✅ Created {band_count} quantile bands from indices: {valid_indices}")
        for i in range(band_count):
            label = quantile_bands[f'quantile_band_{i}_label']
            logger.info(f"   Band {i}: {label}")
        
        return quantile_bands
        
    except Exception as e:
        logger.error(f"❌ Quantile band processing failed: {str(e)}")
        raise


def idx_to_percent(idx: int, num_quantiles: int) -> int:
    """
    Convert quantile index to percentage for labeling.
    
    Note: Index 0 is legacy mean forecast and should be skipped.
    Actual quantiles start at index 1: 1->Q10, 2->Q20, ..., 9->Q90
    
    Args:
        idx: Quantile index (1-based for actual quantiles, 0 is legacy)
        num_quantiles: Total number of quantiles (including legacy index 0)
        
    Returns:
        Percentage value (e.g., 10 for Q10, 90 for Q90)
    """
    if num_quantiles == 10:
        # Special case for 10 quantiles: 1->Q10, 2->Q20, ..., 9->Q90
        # Index 0 is legacy mean, so actual quantiles start at index 1
        return idx * 10
    else:
        # General case: distribute evenly, accounting for skipped index 0
        # If we have 10 total quantiles (0-9), actual quantiles are 1-9
        actual_quantiles = num_quantiles - 1  # Subtract 1 for legacy index 0
        return int(100 * idx / actual_quantiles)