dataset.py

"""
⚡ ULTRON Systems // Streaming & Sharded Parquet Dataset Loader
Safely streams the 90GB HuggingFace OHLCV dataset month-by-month.
Guarantees O(1) constant disk and RAM usage to prevent Colab quota violations.
"""
import os
import gc
import glob
import shutil
import urllib.request
import numpy as np
import pandas as pd
import pyarrow.parquet as pq
import torch
from torch.utils.data import Dataset, DataLoader
from typing import Generator, List, Tuple, Optional
from huggingface_hub import HfFileSystem, hf_hub_download

from config import DatasetConfig
from features import compute_stationary_features

class ShardedOHLCVDataset(Dataset):
    """
    In-memory dataset representing windows from a single monthly shard.
    Efficiently created from vectorized sliding window views.
    """
    def __init__(self, windows: np.ndarray, targets: np.ndarray):
        # windows: [N_samples, seq_len, in_features]
        # targets: [N_samples, num_horizons]
        self.windows = torch.from_numpy(windows).float()
        self.targets = torch.from_numpy(targets).float()
        
    def __len__(self):
        return len(self.windows)
        
    def __getitem__(self, idx: int) -> Tuple[torch.Tensor, torch.Tensor]:
        return self.windows[idx], self.targets[idx]

class StreamingShardManager:
    """
    Orchestrates downloading, windowing, and memory pruning for monthly Parquet files.
    """
    def __init__(self, cfg: DatasetConfig):
        self.cfg = cfg
        os.makedirs(self.cfg.cache_dir, exist_ok=True)
        self.parquet_files = self._discover_remote_shards()
        print(f"⚡ [ULTRON] Discovered {len(self.parquet_files)} target monthly shards between {cfg.start_year} and {cfg.end_year}.")
        
    def _discover_remote_shards(self) -> List[str]:
        """Discovers Parquet file paths in HF repo matching year filters."""
        try:
            fs = HfFileSystem(token=self.cfg.hf_token)
            all_files = fs.ls(f"datasets/{self.cfg.repo_id}/data", detail=False)
        except Exception as e:
            print(f"⚠️ [ULTRON] Direct HfFileSystem query failed ({e}). Constructing standard filenames...")
            all_files = []
            for y in range(self.cfg.start_year, self.cfg.end_year + 1):
                for m in range(1, 13):
                    all_files.append(f"data/ohlcv_{y}-{m:02d}.parquet")
                    
        # Filter matching year range and format
        filtered = []
        for f in all_files:
            fname = os.path.basename(f)
            if fname.startswith("ohlcv_") and fname.endswith(".parquet"):
                # Extract year
                try:
                    parts = fname.replace("ohlcv_", "").replace(".parquet", "").split("-")
                    year = int(parts[0])
                    if self.cfg.start_year <= year <= self.cfg.end_year:
                        filtered.append(f"data/{fname}")
                except Exception:
                    continue
        return sorted(list(set(filtered)))

    def stream_monthly_shards(self) -> Generator[ShardedOHLCVDataset, None, None]:
        """
        Yields a ShardedOHLCVDataset for each month, cleaning up disk and RAM immediately after.
        """
        for rel_path in self.parquet_files:
            filename = os.path.basename(rel_path)
            local_path = os.path.join(self.cfg.cache_dir, filename)
            
            print(f"\n⚡ [ULTRON] Fetching shard: {filename} ...")
            try:
                # Download single shard (~300MB) directly from HuggingFace with authentication
                downloaded_file = hf_hub_download(
                    repo_id=self.cfg.repo_id,
                    repo_type="dataset",
                    filename=rel_path,
                    local_dir=self.cfg.cache_dir,
                    token=self.cfg.hf_token,
                    force_download=False
                )
            except Exception as e:
                print(f"⚠️ [ULTRON] Failed to download {rel_path} via HF Hub: {e}. Trying raw URL fallback...")
                url = f"https://huggingface.co/datasets/{self.cfg.repo_id}/resolve/main/{rel_path}"
                try:
                    urllib.request.urlretrieve(url, local_path)
                    downloaded_file = local_path
                except Exception as ex2:
                    print(f"❌ [ULTRON] Shard {filename} download failed: {ex2}. Skipping.")
                    continue

            # Read and process shard
            dataset = self._process_parquet_to_dataset(downloaded_file)
            
            # Remove file from disk immediately to preserve Colab storage quota!
            try:
                if os.path.exists(downloaded_file):
                    os.remove(downloaded_file)
            except Exception:
                pass
                
            if dataset is not None and len(dataset) > 0:
                print(f"✅ [ULTRON] Shard {filename} ingested: Generated {len(dataset)} overlapping windows.")
                yield dataset
            else:
                print(f"⚠️ [ULTRON] Shard {filename} yielded no valid samples. Proceeding to next shard.")
                
            # Trigger garbage collection
            gc.collect()

    def _process_parquet_to_dataset(self, file_path: str) -> Optional[ShardedOHLCVDataset]:
        """Reads parquet, computes stationary features, and extracts sliding windows."""
        try:
            # Read only required columns
            table = pq.read_table(
                file_path, 
                columns=["timestamp", "open", "high", "low", "close", "volume", "ticker"]
            )
            df = table.to_pandas()
            del table
            gc.collect()
        except Exception as e:
            print(f"❌ Error reading parquet file {file_path}: {e}")
            return None
            
        # Filter tickers if specified
        if self.cfg.liquid_tickers:
            df = df[df["ticker"].isin(self.cfg.liquid_tickers)]
            
        if len(df) == 0:
            del df
            return None
            
        all_windows = []
        all_targets = []
        
        # Group by ticker to prevent cross-ticker leakage
        for ticker, ticker_df in df.groupby("ticker"):
            feats, targets = compute_stationary_features(ticker_df, self.cfg.pred_horizons)
            n_rows = len(feats)
            
            if n_rows < self.cfg.seq_len:
                continue
                
            # Vectorized overlapping sliding window slicing:
            # Stride over sequence with window_stride (e.g. 16 minutes)
            num_windows = (n_rows - self.cfg.seq_len) // self.cfg.window_stride + 1
            if num_windows <= 0:
                continue
                
            for i in range(0, n_rows - self.cfg.seq_len + 1, self.cfg.window_stride):
                w_feats = feats[i : i + self.cfg.seq_len]
                # Target corresponds to the future return from the final bar of this window
                w_target = targets[i + self.cfg.seq_len - 1]
                
                # Check for NaNs or Inf
                if not np.isnan(w_target).any() and not np.isnan(w_feats).any():
                    all_windows.append(w_feats)
                    all_targets.append(w_target)
                    
        del df
        gc.collect()
        
        if not all_windows:
            return None
            
        windows_arr = np.array(all_windows, dtype=np.float32)
        targets_arr = np.array(all_targets, dtype=np.float32)
        
        return ShardedOHLCVDataset(windows_arr, targets_arr)