opt_train_v1.py

import os
import random
import numpy as np
import torch
import torch.nn as nn
import torch.optim as optim
import polars as pl
from rich import print as rprint
import wandb
from enum import Enum, auto
from dataclasses import dataclass
from typing import Tuple
import time

from opt_models import build_ensemble

# -------------------
# Global Config
# -------------------
SEED = 42
PROJECT_NAME = "options_trading_project"  # <-- configurable wandb project
DATASET_PATH = "datasets/Options_Final_Parquet/options_BTC_ETH_dataset_zscore.parquet"
class RunMode(Enum):
    OPTUNA = 1
    TRAIN = 2
    EVAL = 3

# -------------------
# Determinism Setup
# -------------------
os.environ["PYTHONHASHSEED"] = str(SEED)
os.environ["CUBLAS_WORKSPACE_CONFIG"] = ":4096:8"
os.environ["PYTORCH_CUDA_ALLOC_CONF"] = "max_split_size_mb:512"

random.seed(SEED)
np.random.seed(SEED)
torch.manual_seed(SEED)
torch.cuda.manual_seed_all(SEED)

# Reproducibility trade-offs
torch.backends.cudnn.deterministic = False
torch.backends.cudnn.benchmark = True
torch.use_deterministic_algorithms(False)


# -------------------
# Utilities
# -------------------
def print_run_header(mode: RunMode):
    """
    Print a nicely formatted run header with emoji + color based on mode.
    """
    configs = {
        RunMode.OPTUNA:   {"emoji": "🔬", "color": "magenta", "label": "Optuna Trials"},
        RunMode.TRAIN: {"emoji": "🏋️", "color": "green",   "label": "Training"},
        RunMode.EVAL:     {"emoji": "📊", "color": "cyan",    "label": "Evaluation"},
    }

    if mode not in configs:
        raise ValueError(f"Unsupported run mode: {mode}")

    cfg = configs[mode]
    bar = f"[{cfg['color']}] ######### [/{cfg['color']}]"
    msg = f"[{cfg['color']}]{cfg['emoji']} Start {cfg['label']} Pipeline: [/{cfg['color']}]"

    rprint(bar)
    rprint(msg)
    rprint(bar)

# -------------------
# Dataloader
# -------------------
def load_data_to_gpu(
    parquet_path: str,
    feature_window: int = 24,
    val_rows: int = 1000,
    val_asset: str = "BTCUSDT",
    device: str = "cuda"
):
    """
    Load full dataset into GPU-resident tensors.
    Returns train_X, train_y_dict, val_X, val_y_dict.

    train_y_dict and val_y_dict are dicts:
        {
          "ETargets": {2: tensor[B,2], 4: tensor[B,2], ..., 24: tensor[B,2]},
          "OTarget": tensor[B,1]
        }
    """
    df = pl.read_parquet(parquet_path)

    # Features
    feature_cols = [c for c in df.columns if c.startswith("feature_")]

    # Explicit required columns per horizon
    horizons = list(range(2, 25, 2))  # 2h, 4h, ..., 24h
    required_E = {
        h: [
            f"ETarget_{val_asset}_High_{h}h",
            f"ETarget_{val_asset}_Low_{h}h",
        ]
        for h in horizons
    }
    required_O = [f"OTarget_{val_asset}_Sigma_24h"]

    # Validate presence
    for h, cols in required_E.items():
        for col in cols:
            if col not in df.columns:
                raise ValueError(f"Missing expected ETarget column: {col}")
    for col in required_O:
        if col not in df.columns:
            raise ValueError(f"Missing expected OTarget column: {col}")

    n_rows = df.height
    train_end = n_rows - val_rows - feature_window
    val_start = train_end + feature_window

    df_train, df_val = df[:train_end], df[val_start:]

    def make_windows(split_df):
        fvals = split_df.select(feature_cols).to_numpy()
        Evals = {
            h: split_df.select(required_E[h]).to_numpy()
            for h in horizons
        }
        Ovals = split_df.select(required_O).to_numpy()

        X_list, E_dict, O_list = [], {h: [] for h in horizons}, []

        for i in range(len(split_df) - feature_window):
            x_window = fvals[i:i+feature_window]
            X_list.append(x_window)

            for h in horizons:
                E_dict[h].append(Evals[h][i+feature_window-1])  # row aligned
            O_list.append(Ovals[i+feature_window-1])

        X = torch.tensor(np.stack(X_list), dtype=torch.float32, device=device)
        E_targets = {
            h: torch.tensor(np.stack(E_dict[h]), dtype=torch.float32, device=device)
            for h in horizons
        }
        O_target = torch.tensor(np.stack(O_list), dtype=torch.float32, device=device)

        return X, {"ETargets": E_targets, "OTarget": O_target}

    train_X, train_y = make_windows(df_train)
    val_X, val_y = make_windows(df_val)

    return train_X, train_y, val_X, val_y


def remove_data_from_gpu(*tensors): # train_X, train_y, val_X, val_y
    """
    Safely delete tensors from GPU and clear cache.
    """
    for t in tensors:
        if isinstance(t, torch.Tensor) and t.is_cuda:
            del t
    torch.cuda.empty_cache()

def verify_dataset_shapes(train_X, train_y, val_X, val_y, max_print: int = 3):
    """
    Print detailed dataset info to verify correctness.
    Shows shapes, column counts, target horizon checks, and example rows.
    """
    def summarize_split(name, X, y):
        print(f"\n=== {name} SPLIT ===")
        print(f"Features X: {X.shape}")  # [N, window, feature_dim]

        # Check E targets
        for h, t in y["ETargets"].items():
            print(f"  ETarget {h}h: {t.shape}")  # should be [N, 2]
            if t.shape[1] != 2:
                print(f"  ⚠️ WARNING: Expected 2 columns (High/Low), got {t.shape[1]}")
            # print sample rows
            print(f"    sample[0:{max_print}]: {t[:max_print].cpu().numpy()}")

        # Check O target
        O = y["OTarget"]
        print(f"  OTarget: {O.shape}")  # should be [N, 1]
        if O.shape[1] != 1:
            print(f"  ⚠️ WARNING: Expected 1 column (Sigma), got {O.shape[1]}")
        print(f"    sample[0:{max_print}]: {O[:max_print].cpu().numpy()}")

        # Quick sanity stats
        print(f"  NaNs in X: {torch.isnan(X).sum().item()}")
        print(f"  NaNs in ETargets: {sum(torch.isnan(t).sum().item() for t in y['ETargets'].values())}")
        print(f"  NaNs in OTarget: {torch.isnan(O).sum().item()}")

        print("  Min/Max check:")
        print(f"    X min {X.min().item():.4f}, max {X.max().item():.4f}")
        for h, t in y["ETargets"].items():
            print(f"    ETarget {h}h min {t.min().item():.4f}, max {t.max().item():.4f}")
        print(f"    OTarget min {O.min().item():.4f}, max {O.max().item():.4f}")

    # Summarize train + val
    summarize_split("TRAIN", train_X, train_y)
    summarize_split("VAL", val_X, val_y)
    print("\n✅ Dataset verification complete.")


def sanity_check_O_targets(train_y, val_y, max_print=5):
    print("\n--- Sanity Check OTarget ---")
    passed = True  # track if both splits pass

    for split, y in [("TRAIN", train_y), ("VAL", val_y)]:
        O = y["OTarget"]
        print(f"{split} OTarget shape: {O.shape}")
        print(f"First {max_print} rows: {O[:max_print].cpu().numpy().flatten()}")

        # Count zeros
        num_zeros = torch.sum(O == 0).item()
        if num_zeros > 0:
            print(f"  ⚠️ WARNING: Found {num_zeros} pure zero values in {split} OTarget")
            passed = False
        else:
            print(f"  ✅ No pure zero values detected in {split} OTarget")

    if passed:
        print("✅ OTarget verification complete — no zeros found in TRAIN or VAL.")
    else:
        print("❌ OTarget verification failed — zeros detected, check data!")



class EpochSampler:
    def __init__(self, X: torch.Tensor, y: dict, batch_size: int = 64, shuffle: bool = True):
        """
        General-purpose sampler for both training and validation.

        y is a dict with:
          - "ETargets": dict of horizon -> tensor [N, 2]
          - "OTarget": tensor [N, 1]
        """
        self.X = X
        self.y = y
        self.batch_size = batch_size
        self.shuffle_enabled = shuffle
        self.indices = np.arange(len(X))
        self.ptr = 0

        if self.shuffle_enabled:
            self.shuffle()

    def shuffle(self):
        if self.shuffle_enabled:
            np.random.shuffle(self.indices)
        self.ptr = 0

    def get_batch(self):
        if self.ptr >= len(self.indices):
            raise StopIteration
        idx = self.indices[self.ptr:self.ptr + self.batch_size]
        self.ptr += self.batch_size

        batch_y = {
            "E": {h: self.y["ETargets"][h][idx] for h in self.y["ETargets"]},
            "O": self.y["OTarget"][idx],
        }
        return self.X[idx], batch_y

    def __iter__(self):
        self.shuffle()
        return self

    def __next__(self):
        return self.get_batch()



def forward_step(batch_X, ensemble_models, volatility_model):
    """
    Forward pass for all ensemble models and the O-model.
    Returns dictionary of predictions:
      preds = {
          "E": {2: out_2h, 4: out_4h, ..., 24: out_24h},
          "O": out_sigma
      }
    """
    preds = {"E": {}, "O": None}

    # Forward ensemble models (2h → 24h in 2h increments)
    horizons = list(range(2, 26, 2))
    for horizon, model in zip(horizons, ensemble_models):
        preds["E"][horizon] = model(batch_X)  # expect shape [B, 2] (high, low)

    # Forward volatility (OTarget)
    preds["O"] = volatility_model(batch_X)  # expect shape [B, 1]

    return preds

def backward_step(preds, batch_y, optimizers, loss_fn):
    """
    Backward step: compute loss for each model individually.
    Works with CombinedTradingLoss (returns combined + sub-losses).
    Returns:
        losses: dict {
            "E": {2: {"combined": float, "rmse": float, "huber": float, "direction": float}, ...},
            "O": {"combined": float, "rmse": float, "huber": float, "direction": float}
        }
    """
    # Zero grads first
    for opt in optimizers["E"].values():
        opt.zero_grad()
    optimizers["O"].zero_grad()

    # Track all losses separately
    losses = {"E": {}, "O": None}

    # Compute ensemble losses
    for horizon, pred in preds["E"].items():
        target = batch_y["E"][horizon]  # shape [B, 2]

        combined_loss, sub_losses = loss_fn(pred, target)
        combined_loss.backward(retain_graph=True)  # keep graph for other models
        optimizers["E"][horizon].step()

        # Store both combined and components
        losses["E"][horizon] = {
            "combined": combined_loss.item(),
            "rmse": sub_losses["rmse"].item(),
            "huber": sub_losses["huber"].item(),
            "direction": sub_losses["direction"].item(),
        }

    # Compute volatility loss
    combined_loss_O, sub_losses_O = loss_fn(preds["O"], batch_y["O"])
    combined_loss_O.backward()
    optimizers["O"].step()

    losses["O"] = {
        "combined": combined_loss_O.item(),
        "rmse": sub_losses_O["rmse"].item(),
        "huber": sub_losses_O["huber"].item(),
        "direction": sub_losses_O["direction"].item(),
    }

    return losses


@torch.no_grad()
def validation_measurement(preds, batch_y, val_loss_fn):
    """
    Compute validation losses + metrics via CombinedValLoss.

    preds: dict {"E": {h: [B,2]}, "O": [B,1]}
    batch_y: dict {"E": {h: [B,2]}, "O": [B,1]}
    val_loss_fn: CombinedValLoss instance

    Returns:
        dict {
            "E": {2: {...}, 4: {...}, ..., 24: {...}},
            "O": {...},
            "range_extremes": {...}
        }
    """
    return val_loss_fn(preds, batch_y)



'''
LOSS FUNCTIONS
'''

class RMSELoss(nn.Module):
    def __init__(self, eps=1e-8):
        super().__init__()
        self.eps = eps

    def forward(self, pred, target):
        return torch.sqrt(torch.mean((pred - target) ** 2) + self.eps)

class HuberLoss(nn.Module):
    def __init__(self, delta=1.0):
        super().__init__()
        self.delta = delta

    def forward(self, pred, target):
        return torch.nn.functional.huber_loss(pred, target, delta=self.delta)

class DirectionalLoss(nn.Module):
    def __init__(self):
        super().__init__()

    def forward(self, pred, target):
        # Flatten to [N]
        pred = pred.view(-1)
        target = target.view(-1)

        # Mask: only keep non-zero targets
        nonzero_mask = target != 0.0
        if nonzero_mask.sum() == 0:
            # If all are zero, return zero loss (no directionality info)
            return torch.tensor(0.0, device=pred.device)

        # Apply mask
        pred_masked = pred[nonzero_mask]
        target_masked = target[nonzero_mask]

        # Sign mismatch → 1, match → 0
        sign_mismatch = (torch.sign(pred_masked) != torch.sign(target_masked)).float()

        return sign_mismatch.mean()


class RangeExtremeValidator(nn.Module):
    """
    Validation metric: compares max high and min low between real and predicted ranges.
    Works across all 12 horizons. Operates fully on GPU.

    Returns dictionary of measurements:
      {
        "max_high_diff": float,
        "min_low_diff": float,
        "max_high_overshoot": int,  # 1 if pred > truth, else 0
        "min_low_overshoot": int,   # 1 if pred < truth, else 0
      }
    """
    def __init__(self):
        super().__init__()

    def forward(self, preds_E: dict, targets_E: dict):
        """
        preds_E: dict {h: tensor[B,2]} (col0=high, col1=low) for each horizon
        targets_E: same structure as preds_E
        """

        # stack all horizons: shape [num_horizons, B, 2]
        pred_tensor = torch.stack([preds_E[h] for h in sorted(preds_E.keys())], dim=0)
        target_tensor = torch.stack([targets_E[h] for h in sorted(targets_E.keys())], dim=0)

        # get max of highs and min of lows for each sample in batch
        pred_highs = pred_tensor[..., 0]  # [H, B]
        pred_lows  = pred_tensor[..., 1]  # [H, B]
        targ_highs = target_tensor[..., 0]
        targ_lows  = target_tensor[..., 1]

        # extremes across horizons
        pred_max_high, _ = torch.max(pred_highs, dim=0)  # [B]
        pred_min_low,  _ = torch.min(pred_lows,  dim=0)  # [B]
        true_max_high, _ = torch.max(targ_highs, dim=0)
        true_min_low,  _ = torch.min(targ_lows,  dim=0)

        # differences (signed, no abs — keep sign info)
        max_high_diff = pred_max_high - true_max_high   # [B]
        min_low_diff  = pred_min_low - true_min_low     # [B]

        # overshoot logic
        max_high_overshoot = (pred_max_high > true_max_high).float()  # 1 overshoot, 0 undershoot
        min_low_overshoot  = (pred_min_low < true_min_low).float()    # 1 overshoot (lower), 0 undershoot

        # aggregate to batch averages
        results = {
            "max_high_diff": max_high_diff.mean().item(),
            "min_low_diff": min_low_diff.mean().item(),
            "max_high_overshoot": max_high_overshoot.mean().item(),
            "min_low_overshoot": min_low_overshoot.mean().item(),
        }

        return results



class CombinedTrainingLoss(nn.Module):
    def __init__(self, weights=None, huber_delta=1.0, scale=1.0):
        super().__init__()
        self.rmse = RMSELoss()
        self.huber = HuberLoss(delta=huber_delta)
        self.direction = DirectionalLoss()
        self.weights = weights if weights else {"rmse": 0.5, "huber": 0.5, "direction": 0.0}
        self.scale = scale

    def forward(self, pred, target):
        # scale raw values before loss
        pred_scaled = pred * self.scale
        target_scaled = target * self.scale

        losses = {
            "rmse": self.rmse(pred_scaled, target_scaled),
            "huber": self.huber(pred_scaled, target_scaled),
            "direction": self.direction(pred_scaled, target_scaled),
        }
        combined = (
            self.weights["rmse"] * losses["rmse"] +
            self.weights["huber"] * losses["huber"] +
            self.weights["direction"] * losses["direction"]
        )
        return combined, losses

class CombinedValLoss(nn.Module):
    """
    Composite validation loss/metric calculator.
    Combines scalar loss functions (RMSE, Huber, Directional) with
    higher-order metrics (like RangeExtremeValidator).

    Returns a dictionary with per-horizon losses, O-target losses,
    and any additional validators added in the future.
    """

    def __init__(self, weights=None, huber_delta=1.0, scale=1.0, device="cuda"):
        super().__init__()
        # Base scalar losses
        self.rmse = RMSELoss().to(device)
        self.huber = HuberLoss(delta=huber_delta).to(device)
        self.direction = DirectionalLoss().to(device)

        # Extended validators (easy to add more later)
        self.range_validator = RangeExtremeValidator().to(device)

        # Loss weighting for combined score
        self.weights = weights if weights else {
            "rmse": 0.5,
            "huber": 0.5,
            "direction": 0.0,
        }
        self.scale = scale

    def forward(self, preds, targets):
        """
        preds: dict {"E": {h: [B,2]}, "O": [B,1]}
        targets: dict {"E": {h: [B,2]}, "O": [B,1]}

        Returns:
            dict {
                "E": {h: {"combined", "rmse", "huber", "direction"}, ...},
                "O": {"combined", "rmse", "huber", "direction"},
                "range_extremes": {...}
            }
        """
        results = {"E": {}, "O": None, "range_extremes": None}

        # --- E-target losses (loop over horizons)
        for h, pred in preds["E"].items():
            target = targets["E"][h]
            pred_scaled = pred * self.scale
            target_scaled = target * self.scale

            losses = {
                "rmse": self.rmse(pred_scaled, target_scaled),
                "huber": self.huber(pred_scaled, target_scaled),
                "direction": self.direction(pred_scaled, target_scaled),
            }
            combined = (
                self.weights["rmse"] * losses["rmse"] +
                self.weights["huber"] * losses["huber"] +
                self.weights["direction"] * losses["direction"]
            )

            results["E"][h] = {
                "combined": combined.item(),
                "rmse": losses["rmse"].item(),
                "huber": losses["huber"].item(),
                "direction": losses["direction"].item(),
            }

        # --- O-target loss
        pred_scaled = preds["O"] * self.scale
        target_scaled = targets["O"] * self.scale
        losses = {
            "rmse": self.rmse(pred_scaled, target_scaled),
            "huber": self.huber(pred_scaled, target_scaled),
            "direction": self.direction(pred_scaled, target_scaled),
        }
        combined = (
            self.weights["rmse"] * losses["rmse"] +
            self.weights["huber"] * losses["huber"] +
            self.weights["direction"] * losses["direction"]
        )
        results["O"] = {
            "combined": combined.item(),
            "rmse": losses["rmse"].item(),
            "huber": losses["huber"].item(),
            "direction": losses["direction"].item(),
        }

        # --- Range extremes (new validator)
        results["range_extremes"] = self.range_validator(preds["E"], targets["E"])

        return results



'''
DEBUG HELPER
'''
def debug_print_training_batch(preds, batch_y, max_samples: int = 3):
    """
    Debug printer: show true vs predicted values for a few samples.

    preds: dict from forward_step
    batch_y: dict with true targets
    """
    print("\n===== DEBUG BATCH PRINT =====")
    for i in range(max_samples):
        print(f"\n--- Sample {i} ---")

        # Ensemble (high/low targets)
        for h in sorted(preds["E"].keys()):
            y_true = batch_y["E"][h][i].detach().cpu().numpy()
            y_pred = preds["E"][h][i].detach().cpu().numpy()
            print(f"Horizon {h}h:")
            print(f"  True: {y_true}")
            print(f"  Pred: {y_pred}")

        # OTarget (volatility)
        y_true_o = batch_y["O"][i].detach().cpu().numpy()
        y_pred_o = preds["O"][i].detach().cpu().numpy()
        print("OTarget:")
        print(f"  True: {y_true_o}")
        print(f"  Pred: {y_pred_o}")
    print("=============================\n")


def debug_print_training_losses(losses):
    for h, l in losses["E"].items():
        print(
            f"Loss {h}h -> combined: {l['combined']:.4f}, "
            f"rmse: {l['rmse']:.4f}, huber: {l['huber']:.4f}, dir: {l['direction']:.4f}"
        )

    O = losses['O']
    print(
        f"Loss OTarget -> combined: {O['combined']:.4f}, "
        f"rmse: {O['rmse']:.4f}, huber: {O['huber']:.4f}, dir: {O['direction']:.4f}"
    )


def debug_print_val_measurements(losses, preds, batch_y, max_samples: int = 3):
    """
    Debug printer for validation measurements.
    Includes per-horizon E-target losses, O-target losses,
    extended validators (range_extremes),
    and per-sample batch comparison including min/max range.

    Args:
        losses: dict returned from CombinedValLoss
        preds: dict {"E": {h: [B,2]}, "O": [B,1]} from forward_step
        batch_y: dict {"E": {h: [B,2]}, "O": [B,1]} ground truth
        max_samples: how many samples from the batch to print
    """
    print("\n===== VALIDATION MEASUREMENTS =====")

    # --- E-target losses (all horizons)
    for h, l in losses["E"].items():
        print(
            f"[Val] Loss {h}h -> combined: {l['combined']:.4f}, "
            f"rmse: {l['rmse']:.4f}, huber: {l['huber']:.4f}, dir: {l['direction']:.4f}"
        )

    # --- O-target losses
    O = losses["O"]
    print(
        f"[Val] Loss OTarget -> combined: {O['combined']:.4f}, "
        f"rmse: {O['rmse']:.4f}, huber: {O['huber']:.4f}, dir: {O['direction']:.4f}"
    )

    # --- Extended validators (range extremes, etc.)
    if "range_extremes" in losses and losses["range_extremes"] is not None:
        r = losses["range_extremes"]
        print(
            f"[Val] RangeExtremes -> "
            f"max_high_diff: {r['max_high_diff']:.4f}, "
            f"min_low_diff: {r['min_low_diff']:.4f}, "
            f"max_high_overshoot: {r['max_high_overshoot']:.2f}, "
            f"min_low_overshoot: {r['min_low_overshoot']:.2f}"
        )

    # --- Per-sample batch debug
    print("\n===== VALIDATION BATCH SAMPLES =====")
    for i in range(min(max_samples, batch_y["O"].shape[0])):
        print(f"\n--- Sample {i} ---")

        # Ensemble horizons
        for h in sorted(preds["E"].keys()):
            y_true = batch_y["E"][h][i].detach().cpu().numpy()
            y_pred = preds["E"][h][i].detach().cpu().numpy()
            print(f"Horizon {h}h:")
            print(f"  True: {y_true}")
            print(f"  Pred: {y_pred}")

        # OTarget
        y_true_o = batch_y["O"][i].detach().cpu().numpy()
        y_pred_o = preds["O"][i].detach().cpu().numpy()
        print("OTarget:")
        print(f"  True: {y_true_o}")
        print(f"  Pred: {y_pred_o}")

        # Min/max range comparison
        true_highs = [batch_y["E"][h][i][0].item() for h in sorted(batch_y["E"].keys())]
        true_lows = [batch_y["E"][h][i][1].item() for h in sorted(batch_y["E"].keys())]
        pred_highs = [preds["E"][h][i][0].item() for h in sorted(preds["E"].keys())]
        pred_lows = [preds["E"][h][i][1].item() for h in sorted(preds["E"].keys())]

        true_max_high = max(true_highs)
        true_min_low = min(true_lows)
        pred_max_high = max(pred_highs)
        pred_min_low = min(pred_lows)

        print("Range extremes:")
        print(f"  True -> High: {true_max_high:.6f}, Low: {true_min_low:.6f}")
        print(f"  Pred -> High: {pred_max_high:.6f}, Low: {pred_min_low:.6f}")

    print("====================================\n")


'''
RUN LOOPS
'''
def set_models_train(ensemble_models, volatility_model):
    # Set all ensemble models and O-model into training mode
    for model in ensemble_models.values():
        model.train()
    volatility_model.train()
    print("Models set to training mode.")

def set_models_eval(ensemble_models, volatility_model):
    for model in ensemble_models.values():
        model.eval()
    volatility_model.eval()
    print("Models set to evaluation mode.")


def run_training(num_epochs, batch_size,
                 train_X, train_y, val_X, val_y,
                 ensemble_models, volatility_model,
                 E_lr_init, O_lr_init,
                 training_debug_flag, validation_debug_flag,
                 device):

    # Set model group to Training Mode
    # for model in ensemble_models:
    #     model.train()
    # volatility_model.train()

    # Create optimizers for each model
    optimizers = {
        "E": {h: optim.Adam(model.parameters(), lr=E_lr_init) for h, model in ensemble_models.items()},
        "O": optim.Adam(volatility_model.parameters(), lr=O_lr_init),
    }

    # Loss function
    RangeValidator_fn = RangeExtremeValidator().to(device)
    RMSELoss_fn = RMSELoss().to(device)
    HuberLoss_fn = HuberLoss(delta=1.0).to(device)  # delta can be tuned
    DirLoss_fn = DirectionalLoss().to(device)
    CombinedTrainingLoss_fn = CombinedTrainingLoss(
        weights={"rmse": 0.5, "huber": 0.3, "direction": 0.2},
        huber_delta=1.0,
        scale=10.0  # <--- try 10, 100, etc.
    ).to(device)
    CombinedValLoss_fn = CombinedValLoss(
        weights={"rmse": 0.5, "huber": 0.3, "direction": 0.2},
        huber_delta=1.0,
        scale=10.0
    ).to(device)

    # Initialize Data Samplers
    train_sampler = EpochSampler(train_X, train_y, batch_size, shuffle=True)
    val_sampler = EpochSampler(val_X, val_y, batch_size, shuffle=False)


    # Training Metrics
    try:
        # Set all models to train at start of session
        set_models_train(ensemble_models, volatility_model)

        global_step = 0
        global_samples = 0
        for epoch in range(num_epochs):

            print("Shuffling data sheets.")
            train_sampler.shuffle()

            step = 0

            print(f"\n--- Epoch {epoch} ---")
            # Log steps
            for batch_X, batch_y in train_sampler:

                # Forward step
                preds = forward_step(batch_X, ensemble_models.values(), volatility_model)

                # Backward step
                losses = backward_step(preds, batch_y, optimizers, CombinedTrainingLoss_fn)

                if step % 10 == 0:
                    print("Training Update:")
                    print(f"GStep: {global_step} - Epoch {epoch} - Step {step} - Total Samples: {global_samples}")

                    # PRINT DEBUG
                    if step % 100 == 0 and training_debug_flag:
                        debug_print_training_batch(preds, batch_y, max_samples=3)
                        debug_print_training_losses(losses)

                step += 1
                global_step += 1
                global_samples += batch_size

                # VALIDATION
                if global_step % 250 == 0:
                    val_stats = run_validation(val_sampler, CombinedValLoss_fn,
                                               validation_debug_flag,
                                               ensemble_models, volatility_model)

                    '''
                    Wandb Logging
                    '''
                    # Flatten validation stats into wandb format
                    wandb_log = {}
                    for h, stats in val_stats["E"].items():
                        for k, v in stats.items():
                            wandb_log[f"val/E/{h}h/{k}"] = v

                    for k, v in val_stats["O"].items():
                        wandb_log[f"val/O/{k}"] = v

                    for k, v in val_stats["range_extremes"].items():
                        wandb_log[f"val/range_extremes/{k}"] = v

                    # Single wandb.log call with correct step
                    wandb.log(wandb_log, step=global_step)



    except KeyboardInterrupt:
        print("Training interrupted by user (Ctrl+C). Cleaning up...")
    finally:
        # Clean shutdown
        wandb.finish()
        torch.cuda.empty_cache()

    return


def run_validation(val_sampler: EpochSampler, CombinedValLoss_fn,
                   validation_debug_flag,
                   ensemble_models, volatility_model):
    """
    Run validation through all val samples once, computing metrics.
    """
    print("Running validation...")
    val_sampler.ptr = 0  # reset to beginning

    # Switch to evaluation mode (no dropout, batchnorm uses running stats)
    set_models_eval(ensemble_models, volatility_model)

    # --- Collectors for metrics across batches ---
    metrics_collector = {
        "E": {h: {"combined": [], "rmse": [], "huber": [], "direction": []}
              for h in range(2, 25, 2)},
        "O": {"combined": [], "rmse": [], "huber": [], "direction": []},
        "range_extremes": {
            "max_high_diff": [], "min_low_diff": [],
            "max_high_overshoot": [], "min_low_overshoot": []
        }
    }

    # Iterate over all val batches
    val_step = 0
    for batch_X, batch_y in val_sampler:
        print(f"Val Step: {val_step}")
        # === Forward pass through models goes here ===
        val_preds = forward_step(batch_X, ensemble_models.values(), volatility_model)

        # compute metrics
        val_measurements = validation_measurement(val_preds, batch_y, CombinedValLoss_fn)

        # --- Add per-horizon E metrics
        for h, m in val_measurements["E"].items():
            metrics_collector["E"][h]["combined"].append(m["combined"])
            metrics_collector["E"][h]["rmse"].append(m["rmse"])
            metrics_collector["E"][h]["huber"].append(m["huber"])
            metrics_collector["E"][h]["direction"].append(m["direction"])

        # --- Add O metrics
        for k, v in val_measurements["O"].items():
            metrics_collector["O"][k].append(v)

        # --- Add range extremes
        for k, v in val_measurements["range_extremes"].items():
            metrics_collector["range_extremes"][k].append(v)

        if val_step == 10 and validation_debug_flag:
            debug_print_val_measurements(val_measurements, val_preds, batch_y, max_samples=1)

        val_step += 1

    # --- Aggregate into validation_stats (averages over batches) ---
    validation_stats = {"E": {}, "O": {}, "range_extremes": {}}

    # Ensemble horizons
    for h in metrics_collector["E"]:
        validation_stats["E"][h] = {
            "combined": np.mean(metrics_collector["E"][h]["combined"]),
            "rmse": np.mean(metrics_collector["E"][h]["rmse"]),
            "huber": np.mean(metrics_collector["E"][h]["huber"]),
            "direction": np.mean(metrics_collector["E"][h]["direction"]),
        }

    # O target
    validation_stats["O"] = {
        "combined": np.mean(metrics_collector["O"]["combined"]),
        "rmse": np.mean(metrics_collector["O"]["rmse"]),
        "huber": np.mean(metrics_collector["O"]["huber"]),
        "direction": np.mean(metrics_collector["O"]["direction"]),
    }

    # Range extremes
    for k, v in metrics_collector["range_extremes"].items():
        validation_stats["range_extremes"][k] = np.mean(v)

    print("Validation Complete.")
    print("\n...")
    set_models_train(ensemble_models, volatility_model)
    return validation_stats





'''
HYPERPARAMETERS
''' # yeah
@dataclass
class HyperParameters:
    run_mode: RunMode = RunMode.TRAIN
    asset_name: str = "BTCUSDT"
    feature_window: int = 24
    batch_size: int = 32
    val_rows: int = 1000

    # Training Runs
    num_epochs: int = 500

    # Evaluation Runs:
    val_max_steps: int = 10

    # Training Config
    init_train_batch_size: int = 4
    max_train_batch_multiplier: int = 12

    # Validation Config
    val_frequency: int = 2_500
    val_batch_size: int = 16

    # learning rate config
    E_lr_init: float = 1e-5
    O_lr_init: float = 1e-5

    # generator loss weights
    psi1: float = 0.7130
    psi2: float = 0.4350
    psi3: float = 0.1870
    psi4: float = 0.6850
    loss_scaler: float = 10.0


def get_hyperparameters() -> HyperParameters:
    return HyperParameters()

# -------------------
# Main
# -------------------
def main():
    print("Market Volatility Training System Loading...")

    '''
    System Setup
    '''
    wandb.init(project=PROJECT_NAME, config={"seed": SEED})
    hyperparams = get_hyperparameters()


    '''
    Run Mode Selection
    '''
    if hyperparams.run_mode == RunMode.OPTUNA:
        print_run_header(RunMode.OPTUNA)
        pass

    elif hyperparams.run_mode == RunMode.TRAIN:
        print_run_header(RunMode.TRAIN)

        train_X, train_y, val_X, val_y = load_data_to_gpu(DATASET_PATH,
                                                          feature_window=hyperparams.feature_window,
                                                          val_rows=hyperparams.val_rows,
                                                          val_asset=hyperparams.asset_name
                                                          )
        # Gather input dim from data
        feature_dim = train_X.shape[-1]
        print("Total Features: {}".format(feature_dim))
        # Create ensemble hilo models and volatility model
        device = "cuda" if torch.cuda.is_available() else "cpu"
        print("Device: {}".format(device))
        ensemble_models, volatility_model = build_ensemble(feature_dim, hyperparams.feature_window, device=device)
        # Ensure models are on GPU
        for h, model in ensemble_models.items():
            print(h, next(model.parameters()).device)
        print("O model:", next(volatility_model.parameters()).device)

        verify_dataset_shapes(train_X, train_y, val_X, val_y)
        sanity_check_O_targets(train_y, val_y, max_print=5)

        input("\nPress Enter to begin training...")

        # Run training loop
        run_training(hyperparams.num_epochs, hyperparams.batch_size,
                     train_X, train_y, val_X, val_y,
                     ensemble_models, volatility_model,
                     hyperparams.E_lr_init, hyperparams.O_lr_init,
                     training_debug_flag=False, validation_debug_flag=True,
                     device=device)


    elif hyperparams.run_mode == RunMode.EVAL:
        print_run_header(RunMode.EVAL)
    else:
        raise ValueError(f"Invalid run mode: {hyperparams.run_mode}")


if __name__ == "__main__":
    main()