import numpy as np
import torch
import warnings

# Make chronos import optional
try:
    from chronos import BaseChronosPipeline
    CHRONOS_AVAILABLE = True
except ImportError:
    warnings.warn("Chronos-forecasting not available. Using fallback predictions.")
    CHRONOS_AVAILABLE = False
    BaseChronosPipeline = None

class ModelHandler:
    def __init__(self):
        self.model_name = "amazon/chronos-2" 
        self.pipeline = None
        self.device = "cuda" if torch.cuda.is_available() else "cpu"
        self.load_model()
    
    def load_model(self):
        """Load Chronos-2 model using the official BaseChronosPipeline"""
        if not CHRONOS_AVAILABLE:
            print("Chronos-forecasting not installed. Using fallback prediction method.")
            return
            
        try:
            print(f"Loading {self.model_name} on {self.device}...")
            
            self.pipeline = BaseChronosPipeline.from_pretrained(
                self.model_name,
                device_map=self.device,
            )
            print("Chronos-2 pipeline loaded successfully.")
            
        except Exception as e:
            print(f"Error loading Chronos-2 model: {e}")
            print("Using fallback prediction method")
            self.pipeline = None
    
    def predict(self, data, horizon=10):
        """Generate predictions using Chronos-2 or fallback."""
        if not CHRONOS_AVAILABLE or self.pipeline is None:
            # Fallback to simple trend-based prediction
            return self._fallback_predict(data, horizon)
            
        try:
            if data is None or not isinstance(data, dict) or 'original' not in data or len(data['original']) < 20:
                return self._fallback_predict(data, horizon)
            
            # --- Chronos-2 Inference ---
            predictions_samples = self.pipeline.predict(
                data['original'],
                prediction_length=horizon,
            )
            
            # Mengambil nilai rata-rata (mean) dari semua sampel atau single trajectory
            if predictions_samples.ndim > 1 and predictions_samples.shape[0] > 1:
                mean_predictions = np.mean(predictions_samples, axis=0)
            elif predictions_samples.ndim > 1 and predictions_samples.shape[0] == 1:
                mean_predictions = predictions_samples[0]
            else:
                mean_predictions = predictions_samples
            
            return mean_predictions
            
        except Exception as e:
            print(f"Prediction error with Chronos: {e}. Using fallback.")
            return self._fallback_predict(data, horizon)
    
    def _fallback_predict(self, data, horizon=10):
        """Fallback prediction method when Chronos is unavailable"""
        try:
            if data is None or not isinstance(data, dict) or 'original' not in data:
                # Return zero predictions if no data
                return np.zeros(horizon)
            
            values = data['original']
            if len(values) < 5:
                return np.zeros(horizon)
            
            # Simple trend extrapolation
            recent_trend = np.polyfit(range(len(values[-20:])), values[-20:], 1)[0]
            predictions = []
            last_value = values[-1]
            
            for i in range(horizon):
                next_value = last_value + recent_trend * (i + 1)
                noise = np.random.normal(0, data.get('std', 1.0) * 0.1)
                predictions.append(next_value + noise)
            
            return np.array(predictions)
            
        except Exception as e:
            print(f"Fallback prediction error: {e}")
            return np.zeros(horizon)