"""
⚡ ULTRON Systems // OHLCV Transformer Configuration Matrix
Central configuration for data streaming, feature extraction, model architecture, and Colab GPU training.
"""
from dataclasses import dataclass, field
from typing import List, Optional
import os
import torch

@dataclass
class DatasetConfig:
    repo_id: str = "mito0o852/OHLCV-1m"
    # Choose date ranges to process. Each file is data/ohlcv_YYYY-MM.parquet
    start_year: int = 2020
    end_year: int = 2025
    
    # Target top liquid tickers to maximize signal and limit noise (leave empty for all tickers)
    # Default liquid basket: Mega-cap tech, ETFs, high volume liquid equities
    liquid_tickers: List[str] = field(default_factory=lambda: [
        "AAPL", "MSFT", "NVDA", "AMZN", "GOOGL", "META", "TSLA",
        "SPY", "QQQ", "IWM", "AMD", "NFLX", "INTC", "JPM", "BAC"
    ])
    
    # Sequence length in minutes (e.g., 512 minutes ~ 8.5 trading hours)
    seq_len: int = 512
    
    # Overlapping window stride for dataset generation (e.g. step by 16 minutes)
    window_stride: int = 16
    
    # Future prediction horizons in minutes: [1m, 3m, 5m, 10m, 15m]
    pred_horizons: List[int] = field(default_factory=lambda: [1, 3, 5, 10, 15])
    
    # Local ephemeral storage path for downloading 1 parquet file at a time
    cache_dir: str = "/tmp/ohlcv_cache"
    
    # Hugging Face API key / authentication token for high-rate limit transfers
    hf_token: Optional[str] = field(default_factory=lambda: os.environ.get("HF_TOKEN"))

@dataclass
class ModelConfig:
    in_features: int = 8             # Number of stationary features generated per bar
    d_model: int = 128               # Transformer latent dimension
    n_heads: int = 4                 # Multi-head attention heads
    n_layers: int = 4                # Number of Conformer/Transformer layers
    d_ff: int = 256                  # Feed-forward expansion dimension
    conv_kernels: List[int] = field(default_factory=lambda: [3, 7, 15, 31]) # Multi-scale conv kernels
    patch_stride: int = 4            # Stride of overlapping 1D conv tokenizer
    conv_depthwise_kernel: int = 7   # Conformer depthwise conv kernel
    dropout: float = 0.1
    use_revin: bool = True           # Reversible Instance Normalization
    num_horizons: int = 5            # Matching len(DatasetConfig.pred_horizons)

@dataclass
class TrainingConfig:
    batch_size: int = 64             # Physical batch size on GPU
    grad_accum_steps: int = 4        # Effective batch size = 64 * 4 = 256
    learning_rate: float = 3e-4
    min_lr: float = 1e-6
    weight_decay: float = 1e-2
    warmup_steps: int = 300
    max_steps_per_shard: int = 250   # Steps per monthly shard to ensure rapid progression
    max_epochs: int = 5
    fp16: bool = True                # Mixed precision (FP16/BF16)
    
    # Checkpoint settings
    save_every_steps: int = 200
    checkpoint_dir: str = "/content/drive/MyDrive/ohlcv_checkpoints"
    local_checkpoint_dir: str = "./checkpoints"
    
    # Telemetry
    log_interval: int = 25
    device: str = "cuda" if torch.cuda.is_available() else "cpu"