{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Environment Test - FBMC Chronos 2 Forecasting\n",
    "\n",
    "**Purpose**: Verify GPU, Python 3.11, Chronos 2, and dataset access\n",
    "\n",
    "**Date**: 2025-11-12  \n",
    "**Space**: evgueni-p/fbmc-chronos2-forecast"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 1. Check GPU Availability"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "\n",
    "print(\"PyTorch Version:\", torch.__version__)\n",
    "print(\"CUDA Available:\", torch.cuda.is_available())\n",
    "\n",
    "if torch.cuda.is_available():\n",
    "    print(f\"GPU: {torch.cuda.get_device_name(0)}\")\n",
    "    print(f\"VRAM: {torch.cuda.get_device_properties(0).total_memory / 1e9:.1f} GB\")\n",
    "    print(f\"CUDA Version: {torch.version.cuda}\")\n",
    "else:\n",
    "    print(\"[!] WARNING: No GPU detected!\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 2. Check Python Version"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import sys\n",
    "print(f\"Python: {sys.version}\")\n",
    "print(f\"Version Info: {sys.version_info}\")\n",
    "\n",
    "# Verify Python 3.11+\n",
    "assert sys.version_info >= (3, 11), \"Python 3.11+ required for Chronos 2\"\n",
    "print(\"[+] Python version check passed!\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 3. Test Chronos 2 Import"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "try:\n",
    "    from chronos import Chronos2Pipeline\n",
    "    import chronos\n",
    "    print(f\"Chronos version: {chronos.__version__}\")\n",
    "    print(\"[+] Chronos 2 imported successfully!\")\n",
    "except ImportError as e:\n",
    "    print(f\"[!] ERROR: Could not import Chronos 2: {e}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 4. Load HuggingFace Dataset"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from datasets import load_dataset\n",
    "import os\n",
    "\n",
    "# Check if HF_TOKEN is available\n",
    "hf_token = os.getenv('HF_TOKEN')\n",
    "print(f\"HF_TOKEN available: {hf_token is not None}\")\n",
    "\n",
    "# Load dataset\n",
    "print(\"\\nLoading dataset...\")\n",
    "dataset = load_dataset(\"evgueni-p/fbmc-features-24month\", split=\"train\")\n",
    "\n",
    "print(f\"[+] Dataset loaded: {len(dataset)} rows\")\n",
    "print(f\"    Columns: {len(dataset.column_names)} features\")\n",
    "print(f\"    First 5 columns: {dataset.column_names[:5]}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 5. Load Chronos 2 Model on GPU"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print(\"Loading Chronos 2 Large model...\")\n",
    "print(\"(This may take 2-3 minutes on first load)\\n\")\n",
    "\n",
    "pipeline = Chronos2Pipeline.from_pretrained(\n",
    "    \"amazon/chronos-2-large\",\n",
    "    device_map=\"cuda\",\n",
    "    torch_dtype=torch.float32\n",
    ")\n",
    "\n",
    "print(\"[+] Chronos 2 Large loaded successfully on GPU!\")\n",
    "print(f\"    Model device: {next(pipeline.model.parameters()).device}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 6. Quick Inference Test"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import pandas as pd\n",
    "import numpy as np\n",
    "from datetime import datetime, timedelta\n",
    "\n",
    "# Create minimal test data\n",
    "print(\"Creating test data...\")\n",
    "timestamps = pd.date_range(start='2024-01-01', periods=100, freq='H')\n",
    "test_context = pd.DataFrame({\n",
    "    'timestamp': timestamps,\n",
    "    'border': ['TEST'] * len(timestamps),\n",
    "    'target': np.random.randn(len(timestamps)) * 100 + 1000  # Random values around 1000 MW\n",
    "})\n",
    "\n",
    "future_timestamps = pd.date_range(start=timestamps[-1] + timedelta(hours=1), periods=24, freq='H')\n",
    "test_future = pd.DataFrame({\n",
    "    'timestamp': future_timestamps,\n",
    "    'border': ['TEST'] * len(future_timestamps)\n",
    "})\n",
    "\n",
    "print(f\"Context shape: {test_context.shape}\")\n",
    "print(f\"Future shape: {test_future.shape}\")\n",
    "\n",
    "# Run inference\n",
    "print(\"\\nRunning inference (24-hour forecast)...\")\n",
    "forecasts = pipeline.predict_df(\n",
    "    context_df=test_context,\n",
    "    future_df=test_future,\n",
    "    prediction_length=24,\n",
    "    id_column='border',\n",
    "    timestamp_column='timestamp',\n",
    "    num_samples=50\n",
    ")\n",
    "\n",
    "print(f\"[+] Inference complete! Forecast shape: {forecasts.shape}\")\n",
    "print(f\"    Forecast columns: {list(forecasts.columns)}\")\n",
    "\n",
    "if 'mean' in forecasts.columns:\n",
    "    print(f\"\\nForecast statistics:\")\n",
    "    print(f\"    Mean: {forecasts['mean'].mean():.2f}\")\n",
    "    print(f\"    Min: {forecasts['mean'].min():.2f}\")\n",
    "    print(f\"    Max: {forecasts['mean'].max():.2f}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## ✅ Summary"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print(\"=\"*60)\n",
    "print(\"ENVIRONMENT TEST SUMMARY\")\n",
    "print(\"=\"*60)\n",
    "print(\"[+] GPU: Available and working\")\n",
    "print(\"[+] Python: 3.11+\")\n",
    "print(\"[+] Chronos 2: Installed and imported\")\n",
    "print(\"[+] Dataset: Loaded successfully\")\n",
    "print(\"[+] Model: Chronos 2 Large on GPU\")\n",
    "print(\"[+] Inference: Working correctly\")\n",
    "print(\"=\"*60)\n",
    "print(\"\\n[+] Environment is READY for full inference run!\")\n",
    "print(\"=\"*60)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.11.0"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}