import os
import torch
import numpy as np
from chronos import BaseChronosPipeline

# Fix: Point HF cache to a writable path for the non-root HF Spaces user
os.environ.setdefault("HF_HOME", "/home/user/.cache/huggingface")
os.environ.setdefault("TRANSFORMERS_CACHE", "/home/user/.cache/huggingface")


class SMCPredictionEngine:
    def __init__(self, model_name="amazon/chronos-bolt-base"):
        """
        Initializes the Hugging Face Time-Series Transformer.
        Runs on CPU for compatibility with standard HF Spaces containers.
        """
        print(f"Loading pre-trained AI weights: {model_name}...")
        self.pipeline = BaseChronosPipeline.from_pretrained(
            model_name,
            device_map="cpu",
            torch_dtype=torch.float32,
        )
        print("Model loaded successfully.")

    def calculate_forecast(self, history_array: list | np.ndarray, prediction_length: int = 12) -> dict:
        """
        Takes an array of historical prices and generates probabilistic forecasts.

        BaseChronosPipeline.predict() returns a tensor of shape:
            [batch_size, num_samples, prediction_length]

        We derive low / median / high by taking quantiles across the sample axis (axis=1).
        """
        context_tensor = torch.tensor(
            np.array(history_array, dtype=np.float32).reshape(1, -1),  # [1, seq_len]
            dtype=torch.float32,
        )

        # forecast shape: [1, num_samples, prediction_length]
        forecast = self.pipeline.predict(
            context=context_tensor,
            prediction_length=prediction_length,
        )

        # Convert to numpy: shape [num_samples, prediction_length]
        samples = forecast[0].numpy()  # drop batch dimension

        # Derive quantile bands across the sample axis
        low_bounds         = np.quantile(samples, 0.10, axis=0)  # 10th percentile
        median_predictions = np.quantile(samples, 0.50, axis=0)  # median
        high_bounds        = np.quantile(samples, 0.90, axis=0)  # 90th percentile

        # Projected momentum: % change from last known price to median forecast end
        last_price = float(history_array[-1])
        projected_change = ((median_predictions[-1] - last_price) / last_price) * 100.0

        return {
            "low": low_bounds,
            "median": median_predictions,
            "high": high_bounds,
            "projected_change_pct": float(projected_change),
        }