services/model_service.py

"""
Chronos 2 Model Service
Handles model loading, caching, and inference using Chronos2Pipeline
"""

import logging
import time
from typing import Dict, List, Optional, Tuple, Any
import numpy as np
import pandas as pd
import torch
from chronos import ChronosPipeline, Chronos2Pipeline

from config.constants import CHRONOS2_MODEL, CONFIDENCE_LEVELS
from config.settings import CONFIG, DEVICE, MODEL_CONFIG

logger = logging.getLogger(__name__)


class ChronosModelService:
    """
    Service for managing Chronos 2 model lifecycle and inference
    Uses Chronos2Pipeline with DataFrame-based API
    """

    def __init__(self):
        self.model = None
        self.device = None
        self.model_variant = None
        self.is_loaded = False
        self.load_time = None
        self.is_chronos2 = False  # Track which pipeline type is loaded

    def _get_device(self) -> str:
        """Determine the best available device"""
        if DEVICE == 'cuda':
            if not torch.cuda.is_available():
                logger.warning("CUDA requested but not available, falling back to CPU")
                return 'cpu'
            return 'cuda'
        elif DEVICE == 'cpu':
            return 'cpu'
        else:  # auto
            return 'cuda' if torch.cuda.is_available() else 'cpu'

    def load_model(self) -> Dict[str, Any]:
        """
        Load the Chronos 2 model at startup

        Returns:
            Dictionary with loading status and metadata
        """
        try:
            start_time = time.time()
            logger.info("Loading Chronos 2 model from HuggingFace paper 2510.15821")

            # Use the single Chronos-2 model
            model_path = CHRONOS2_MODEL
            self.model_variant = 'chronos-2'

            # Determine device
            self.device = self._get_device()
            logger.info(f"Using device: {self.device}")

            # Load model using Chronos2Pipeline
            self.model = Chronos2Pipeline.from_pretrained(
                model_path,
                device_map=self.device,
                torch_dtype=torch.bfloat16 if self.device == 'cuda' else torch.float32,
            )
            self.is_chronos2 = True

            self.load_time = time.time() - start_time
            self.is_loaded = True

            logger.info(f"Model loaded successfully in {self.load_time:.2f}s")

            # Warmup prediction
            if MODEL_CONFIG['warmup_enabled']:
                self._warmup()

            return {
                'status': 'success',
                'model': 'chronos-2',
                'device': self.device,
                'load_time': self.load_time,
                'model_name': model_path
            }

        except Exception as e:
            logger.error(f"Failed to load model: {str(e)}", exc_info=True)
            self.is_loaded = False
            return {
                'status': 'error',
                'error': str(e)
            }

    def _warmup(self):
        """Run a warmup prediction to initialize the model"""
        try:
            logger.info("Running warmup prediction")

            # Create warmup DataFrame in Chronos 2 format
            warmup_data = pd.DataFrame({
                'id': ['warmup'] * MODEL_CONFIG['warmup_length'],
                'timestamp': pd.date_range('2020-01-01', periods=MODEL_CONFIG['warmup_length'], freq='D'),
                'target': np.random.randn(MODEL_CONFIG['warmup_length'])
            })

            self.predict(
                warmup_data,
                horizon=MODEL_CONFIG['warmup_horizon'],
                confidence_levels=[80]
            )
            logger.info("Warmup completed successfully")

        except Exception as e:
            logger.warning(f"Warmup failed: {str(e)}")

    def predict(
        self,
        data: pd.DataFrame,
        horizon: int,
        confidence_levels: List[int] = None,
        future_df: Optional[pd.DataFrame] = None
    ) -> Dict[str, Any]:
        """
        Generate forecasts using Chronos 2 model with DataFrame API

        Args:
            data: DataFrame with columns ['id', 'timestamp', 'target']
                  Can also include covariates for multivariate forecasting
            horizon: Number of periods to forecast
            confidence_levels: List of confidence levels (e.g., [80, 90, 95])
            future_df: Optional DataFrame with future covariate values

        Returns:
            Dictionary with predictions and metadata
        """
        logger.info("=" * 80)
        logger.info("MODEL SERVICE: predict() - ENTRY")
        logger.info(f"Data shape: {data.shape}")
        logger.info(f"Data columns: {data.columns.tolist()}")
        logger.info(f"Horizon: {horizon}")
        logger.info(f"Confidence levels: {confidence_levels}")
        logger.info(f"Is loaded: {self.is_loaded}")
        logger.info("=" * 80)

        if not self.is_loaded:
            logger.error("✗ Model not loaded!")
            raise RuntimeError("Model not loaded. Call load_model() first.")

        try:
            start_time = time.time()
            logger.info("Starting prediction...")

            # Use default confidence levels if not provided
            if confidence_levels is None:
                confidence_levels = CONFIDENCE_LEVELS

            # Calculate quantile levels from confidence intervals
            quantile_levels = []
            for cl in sorted(confidence_levels):
                lower = (100 - cl) / 200  # e.g., 80% -> 0.10
                upper = 1 - lower  # e.g., 80% -> 0.90
                quantile_levels.extend([lower, upper])

            # Add median
            quantile_levels.append(0.5)
            quantile_levels = sorted(set(quantile_levels))

            logger.info(f"Generating forecast for horizon={horizon}, quantiles={quantile_levels}")

            # Ensure required columns exist
            required_cols = ['id', 'timestamp', 'target']
            logger.info(f"Checking for required columns: {required_cols}")
            if not all(col in data.columns for col in required_cols):
                error_msg = f"Data must contain columns: {required_cols}, but got: {data.columns.tolist()}"
                logger.error(f"✗ {error_msg}")
                raise ValueError(error_msg)
            logger.info("✓ All required columns present")

            # Generate forecast using appropriate API
            if self.is_chronos2:
                logger.info("Using Chronos2Pipeline.predict_df() method")
                logger.info(f"Calling predict_df with prediction_length={horizon}, quantile_levels={quantile_levels}")
                # Use Chronos 2 DataFrame API
                pred_df = self.model.predict_df(
                    df=data,
                    future_df=future_df,
                    prediction_length=horizon,
                    quantile_levels=quantile_levels,
                    id_column='id',
                    timestamp_column='timestamp',
                    target='target'
                )
                logger.info(f"✓ predict_df completed - result shape: {pred_df.shape}")
            else:
                # Use original Chronos tensor API
                # Convert DataFrame to tensor
                context_tensor = torch.tensor(data['target'].values, dtype=torch.float32).unsqueeze(0)

                # Generate forecast
                forecast_tensors = self.model.predict(
                    context=context_tensor,
                    prediction_length=horizon,
                    num_samples=20,  # Number of sample paths
                    limit_prediction_length=False
                )

                # Convert tensor output to DataFrame format
                # forecast_tensors shape: [batch, num_samples, prediction_length]
                quantiles_np = np.quantile(
                    forecast_tensors.squeeze(0).numpy(),
                    q=quantile_levels,
                    axis=0
                )

                # Create prediction DataFrame in Chronos 2 format
                last_timestamp = pd.to_datetime(data['timestamp'].iloc[-1])
                freq = pd.infer_freq(pd.to_datetime(data['timestamp']))
                if freq is None:
                    freq = 'D'  # Default to daily

                future_timestamps = pd.date_range(
                    start=last_timestamp,
                    periods=horizon + 1,
                    freq=freq
                )[1:]  # Exclude the last historical point

                pred_df = pd.DataFrame({
                    'id': [data['id'].iloc[0]] * horizon,
                    'timestamp': future_timestamps
                })

                # Add quantile columns
                for i, q in enumerate(quantile_levels):
                    pred_df[f'{q:.2f}'] = quantiles_np[i, :]

            # Process forecast results
            # pred_df contains columns: id, timestamp, and quantile columns

            # Extract forecast for the first series (if multiple)
            series_ids = pred_df['id'].unique()
            if len(series_ids) > 0:
                series_pred = pred_df[pred_df['id'] == series_ids[0]].copy()
            else:
                series_pred = pred_df.copy()

            # Create forecast dataframe with confidence intervals
            forecast_df = pd.DataFrame({
                'ds': series_pred['timestamp'],
                'forecast': series_pred['0.5']  # Median forecast
            })

            # Add confidence intervals
            for cl in confidence_levels:
                lower = (100 - cl) / 200
                upper = 1 - lower

                lower_col = f'{lower:.2f}'
                upper_col = f'{upper:.2f}'

                if lower_col in series_pred.columns:
                    forecast_df[f'lower_{cl}'] = series_pred[lower_col].values
                if upper_col in series_pred.columns:
                    forecast_df[f'upper_{cl}'] = series_pred[upper_col].values

            inference_time = time.time() - start_time

            logger.info(f"✓ Forecast generated successfully in {inference_time:.2f}s")
            logger.info(f"Returning forecast DataFrame with {len(forecast_df)} rows")
            logger.info("MODEL SERVICE: predict() - EXIT (success)")
            logger.info("=" * 80)

            return {
                'status': 'success',
                'forecast': forecast_df,
                'inference_time': inference_time,
                'horizon': horizon,
                'confidence_levels': confidence_levels,
                'full_prediction': pred_df  # Include full prediction for multivariate
            }

        except Exception as e:
            logger.error(f"✗ EXCEPTION in predict(): {str(e)}", exc_info=True)
            logger.info("MODEL SERVICE: predict() - EXIT (exception)")
            logger.info("=" * 80)
            return {
                'status': 'error',
                'error': str(e)
            }

    def backtest(
        self,
        data: pd.DataFrame,
        test_size: int,
        forecast_horizon: int,
        confidence_levels: List[int] = None
    ) -> Dict[str, Any]:
        """
        Perform backtesting on historical data to evaluate model performance

        Args:
            data: DataFrame with columns ['id', 'timestamp', 'target']
            test_size: Number of periods to use for testing
            forecast_horizon: Forecast horizon for each prediction
            confidence_levels: List of confidence levels

        Returns:
            Dictionary with backtest results including predictions vs actuals
        """
        logger.info("=" * 80)
        logger.info("MODEL SERVICE: backtest() - ENTRY")
        logger.info(f"Data shape: {data.shape}")
        logger.info(f"Test size: {test_size}")
        logger.info(f"Forecast horizon: {forecast_horizon}")
        logger.info("=" * 80)

        if not self.is_loaded:
            raise RuntimeError("Model not loaded. Call load_model() first.")

        try:
            start_time = time.time()

            # Split data into train and test
            train_size = len(data) - test_size
            if train_size < forecast_horizon * 2:
                raise ValueError(f"Insufficient training data. Need at least {forecast_horizon * 2} points.")

            # Use rolling window approach
            # We'll make predictions for the test period using the training data
            train_data = data.iloc[:train_size].copy()
            test_data = data.iloc[train_size:].copy()

            logger.info(f"Train size: {len(train_data)}, Test size: {len(test_data)}")

            # Make prediction on test period
            forecast_result = self.predict(
                data=train_data,
                horizon=test_size,
                confidence_levels=confidence_levels
            )

            if forecast_result['status'] == 'error':
                return forecast_result

            forecast_df = forecast_result['forecast']

            # Align forecast with actual values
            backtest_df = pd.DataFrame({
                'timestamp': test_data['timestamp'].values,
                'actual': test_data['target'].values,
                'predicted': forecast_df['forecast'].values[:len(test_data)]
            })

            # Add confidence intervals if available
            for cl in (confidence_levels or []):
                lower_col = f'lower_{cl}'
                upper_col = f'upper_{cl}'
                if lower_col in forecast_df.columns:
                    backtest_df[lower_col] = forecast_df[lower_col].values[:len(test_data)]
                if upper_col in forecast_df.columns:
                    backtest_df[upper_col] = forecast_df[upper_col].values[:len(test_data)]

            # Calculate metrics
            actual = backtest_df['actual'].values
            predicted = backtest_df['predicted'].values

            # Remove any NaN values
            mask = ~(np.isnan(actual) | np.isnan(predicted))
            actual = actual[mask]
            predicted = predicted[mask]

            if len(actual) == 0:
                raise ValueError("No valid data points for metric calculation")

            mae = np.mean(np.abs(actual - predicted))
            rmse = np.sqrt(np.mean((actual - predicted) ** 2))
            mape = np.mean(np.abs((actual - predicted) / (actual + 1e-10))) * 100

            # R-squared
            ss_res = np.sum((actual - predicted) ** 2)
            ss_tot = np.sum((actual - np.mean(actual)) ** 2)
            r2 = 1 - (ss_res / (ss_tot + 1e-10))

            metrics = {
                'MAE': float(mae),
                'RMSE': float(rmse),
                'MAPE': float(mape),
                'R2': float(r2)
            }

            inference_time = time.time() - start_time

            logger.info(f"✓ Backtest completed in {inference_time:.2f}s")
            logger.info(f"Metrics: MAE={mae:.2f}, RMSE={rmse:.2f}, MAPE={mape:.2f}%, R2={r2:.4f}")
            logger.info("MODEL SERVICE: backtest() - EXIT (success)")
            logger.info("=" * 80)

            return {
                'status': 'success',
                'backtest_data': backtest_df,
                'metrics': metrics,
                'inference_time': inference_time,
                'train_size': train_size,
                'test_size': test_size
            }

        except Exception as e:
            logger.error(f"✗ EXCEPTION in backtest(): {str(e)}", exc_info=True)
            logger.info("MODEL SERVICE: backtest() - EXIT (exception)")
            logger.info("=" * 80)
            return {
                'status': 'error',
                'error': str(e)
            }

    def get_status(self) -> Dict[str, Any]:
        """Get current model status"""
        return {
            'is_loaded': self.is_loaded,
            'variant': self.model_variant,
            'device': self.device,
            'load_time': self.load_time
        }


# Global model service instance
model_service = ChronosModelService()