reference_training_script.py

import os
import math
import random
import numpy as np
import pyarrow.parquet as pq

SEED = 42
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)

import torch
import torch.optim as optim
import polars as pl
from torch.utils.data import DataLoader, TensorDataset, IterableDataset
from torch.amp import autocast, GradScaler
from torch.optim.lr_scheduler import CosineAnnealingLR

# To this:
scaler = GradScaler(
    device='cuda',
    init_scale=2. ** 10,  # Start with higher scale
    growth_factor=2.0,  # Aggressive growth
    backoff_factor=0.5,
    growth_interval=2000,  # Grow every 2000 iterations
    enabled=True
)

torch.manual_seed(SEED)
torch.cuda.manual_seed_all(SEED)
torch.backends.cudnn.deterministic = False
torch.backends.cudnn.benchmark = True
torch.use_deterministic_algorithms(False)

# TF32 for Ampere+ GPUs
torch.backends.cuda.matmul.allow_tf32 = True
torch.backends.cudnn.allow_tf32 = True

# Memory optimization
if torch.cuda.is_available():
    torch.cuda.empty_cache()
    torch.cuda.set_per_process_memory_fraction(0.95)

import time
import wandb
import logging
import pandas as pd
from rich import print
from contextlib import contextmanager
from contexttimer import Timer
from datetime import datetime

# NEW IMPORTS FOR PRODUCER-CONSUMER SYSTEM
import multiprocessing as mp
from multiprocessing import Queue, Process, Event
import threading
from collections import deque

# from hilo_CTG_forecast_multiforecast_v2_larger import (
#     EnhancedConvTransformerGenerator as Generator,
#     create_generator,
#     estimate_model_size
# )

from hilo_CTG_forecast_multiforecast import (
    ConvTransformerGenerator as Generator,
    create_generator,
    estimate_model_size  # You'll need to add this function to the v1 file
)

from hilo_evals import (
    custom_loss_function,
    forecast_focused_loss,
    calculate_rmse, calculate_mae, calculate_mape,
    calculate_nrmse, calculate_weighted_mae,
    calculate_high_low_accuracy, calculate_predictions_within_accuracy_tolerance,
    calculate_undershoot_overshoot_ratio
)
from hilo_evals import forecast_focused_loss as custom_loss_function

# Debugs
FLAG_DEBUG_GPU_PROFILING = False
FLAG_DEBUG_PRINT_LOSS = True
FLAG_DEBUG_GRADIENT_FLOW = True
FLAG_DEBUG_TARGET_ALIGNMENT = False
FLAG_LOG_GRADIENTS_WANDB = False

# Control Flags
WANDB_PROJECT_NAME = "hilo-rev3-multitarget"
FLAG_SAVE_MODELS = False
MODEL_NAME = "Hilo_MEGA_CTG_Jun7_4h_20w_MultiTarget"
MODEL_SAVE_THRESHOLD = 0.9169
TOL_TEST_ACC_THRESHOLD = 0.002

# Dataset
BASE_DATA_DIR = 'datasets/unified_2_final'

# Viewing
FLAG_PRINT_TRAINING_PREDICTIONS = False
FLAG_PRINT_VALIDATION_PREDICTIONS = True
FLAG_PRINT_TIMER_BLOCKS = False
FLAG_PRINT_GPU_UTILIZATION = False
FLAG_PRINT_DATA_PRODUCER_DEBUG = False

# Noise Injection
FLAG_NOISE_INPUTS = False
FLAG_NOISE_LABELS = False
NOISE_LEVEL = 0.001

INTERVAL_COUNTS = {
    '5m': 571983,
    '15m': 190655,
    '30m': 95328,
    '1h': 47663,
    '2h': 23831,
    '4h': 11915,
    '6h': 7943,
    '8h': 5957,
    '12h': 3971,
}

# NEW: Multi-target configuration
# Define all available forecast distances (candles forward)
FORECAST_DISTANCES = ['1c', '4c', '6c', '8c', '12c', '24c', '72c']


# Legacy compatibility: if using single forecast_distance, convert to list
def ensure_forecast_distances_list(forecast_distances_or_single):
    """Convert single forecast distance to list for backward compatibility."""
    if isinstance(forecast_distances_or_single, (list, tuple)):
        return list(forecast_distances_or_single)
    else:
        return [forecast_distances_or_single]


# Logging
logging.basicConfig(
    filename='HILO_TRAINER.log',
    level=logging.DEBUG,
    format='%(asctime)s - %(levelname)s - %(message)s'
)


@contextmanager
def timed_block(name: str):
    with Timer() as t:
        yield
    if FLAG_PRINT_TIMER_BLOCKS:
        print(f"{name} took {t.elapsed:.4f} seconds")


'''
Learning Rate Setup:
'''
class GoldilocksZoneScheduler:
    """
    Custom scheduler that keeps learning rate in the 'goldilocks zone' for most of training.
    Based on your empirical observation that learning happens best between 0.0001-0.00015.
    """

    def __init__(self, optimizer, lr_g_init, total_steps, goldilocks_ratio=0.7):
        self.optimizer = optimizer
        self.total_steps = total_steps
        self.goldilocks_ratio = goldilocks_ratio  # 70% of training in goldilocks zone

        # Your empirically determined goldilocks zone
        self.lr_max = lr_g_init  # Starting LR
        self.goldilocks_high = 0.00022  # Upper goldilocks zone
        self.goldilocks_low = 0.00017  # Lower goldilocks zone
        self.lr_min = 1e-4  # Final minimum (still useful)

        # Calculate phase boundaries
        self.ramp_down_steps = int(total_steps * 0.1)  # 10% to reach goldilocks zone
        self.goldilocks_steps = int(total_steps * goldilocks_ratio)  # 70% in goldilocks
        self.final_steps = total_steps - self.ramp_down_steps - self.goldilocks_steps  # 20% final decay

        self.step_count = 0

        print(f"🎯 Goldilocks Zone Scheduler:")
        print(f"   Phase 1 (steps 0-{self.ramp_down_steps:,}): {self.lr_max} → {self.goldilocks_high} (ramp to goldilocks)")
        print(f"   Phase 2 (steps {self.ramp_down_steps:,}-{self.ramp_down_steps + self.goldilocks_steps:,}): {self.goldilocks_high} → {self.goldilocks_low} (goldilocks zone)")
        print(f"   Phase 3 (steps {self.ramp_down_steps + self.goldilocks_steps:,}-{total_steps:,}): {self.goldilocks_low} → {self.lr_min} (final polishing)")

    def step(self):
        if self.step_count < self.ramp_down_steps:
            # Phase 1: Quick ramp down to goldilocks zone (10% of training)
            progress = self.step_count / self.ramp_down_steps
            # Use cosine for smooth transition
            cosine_factor = 0.5 * (1 + math.cos(math.pi * progress))
            lr = self.goldilocks_high + (self.lr_max - self.goldilocks_high) * cosine_factor

        elif self.step_count < self.ramp_down_steps + self.goldilocks_steps:
            # Phase 2: Stay in goldilocks zone with gentle decay (70% of training)
            goldilocks_progress = (self.step_count - self.ramp_down_steps) / self.goldilocks_steps
            # Very gentle decay through goldilocks zone
            lr = self.goldilocks_low + (self.goldilocks_high - self.goldilocks_low) * (1 - goldilocks_progress) ** 0.3

        else:
            # Phase 3: Final decay to minimum (20% of training)
            final_progress = (self.step_count - self.ramp_down_steps - self.goldilocks_steps) / self.final_steps
            # Smooth final decay
            cosine_factor = 0.5 * (1 + math.cos(math.pi * final_progress))
            lr = self.lr_min + (self.goldilocks_low - self.lr_min) * cosine_factor

        # Update optimizer
        for param_group in self.optimizer.param_groups:
            param_group['lr'] = lr

        self.step_count += 1
        return lr

    def get_last_lr(self):
        """For compatibility with other schedulers"""
        return [param_group['lr'] for param_group in self.optimizer.param_groups]


'''
Data Preparations - Enhanced with Producer-Consumer System
'''

class TimescaleSampler:
    def __init__(self, interval_counts, weight_power=1.0, target_boost=None, target_interval=None):
        self.intervals = list(interval_counts.keys())
        raw_weights = np.array([interval_counts[k] for k in self.intervals], dtype=np.float64)

        # Normalize raw weights to form a distribution
        norm_weights = raw_weights / raw_weights.sum()

        # Apply shaping curve
        shaped_weights = norm_weights ** weight_power

        # Apply boost for target timescale
        if target_boost and target_interval in self.intervals:
            idx = self.intervals.index(target_interval)
            shaped_weights[idx] *= target_boost

        # Final normalization
        self.weights = shaped_weights / shaped_weights.sum()

    def sample(self, k=1):
        return random.choices(self.intervals, weights=self.weights, k=k)


# NEW: Producer function that runs in separate processes - MODIFIED for multi-target
def data_producer_worker(
        producer_id,
        data_queue,
        control_event,
        pause_event,  # NEW: For pausing during batch size changes
        base_dir,
        asset_name,
        forecast_distances,  # CHANGED: Now list of distances instead of single
        window_size,
        mode,
        val_cutoff_date,
        validation_timescale,
        validation_distance,  # NEW: For validation-specific distance
        intervals,
        weight_power,
        target_interval,
        target_boost,
        batch_size,
        seed_offset=0
):
    """
    Producer worker that generates batches of data and puts them in the queue.
    This runs in a separate process.
    MODIFIED: Now handles multiple forecast distances for multi-target learning.
    """
    # Set unique seed for this producer
    worker_seed = SEED + producer_id + seed_offset
    np.random.seed(worker_seed)
    random.seed(worker_seed)

    # Timing diagnostics
    batch_count = 0
    total_collect_time = 0
    total_stack_time = 0
    total_queue_time = 0

    try:
        # Create the dataset (same logic as before but with multi-target)
        dataset = MultiTimescaleParquetDataset(
            base_dir=base_dir,
            asset_name=asset_name,
            forecast_distances=forecast_distances,  # CHANGED
            window_size=window_size,
            mode=mode,
            val_cutoff_date=val_cutoff_date,
            validation_timescale=validation_timescale,
            validation_distance=validation_distance,  # NEW
            intervals=intervals,
            weight_power=weight_power,
            target_interval=target_interval,
            target_boost=target_boost
        )

        # Create iterator
        data_iter = iter(dataset)

        batch_x = []
        batch_y = []

        while not control_event.is_set():
            # Check if we should pause (for batch size changes)
            if pause_event.is_set():
                # Clear any partial batch when paused
                if batch_x or batch_y:
                    batch_x = []
                    batch_y = []
                time.sleep(0.1)
                continue

            try:
                # TIME: Sample collection
                collect_start = time.time()

                # Collect batch_size samples with pause checks
                batch_complete = True
                for _ in range(batch_size):
                    # Check pause before each sample
                    if pause_event.is_set():
                        batch_complete = False
                        break

                    x, y = next(data_iter)
                    batch_x.append(x)
                    batch_y.append(y)

                # Only process if we completed the batch and aren't paused
                if not batch_complete or pause_event.is_set():
                    # Clear partial batch
                    batch_x = []
                    batch_y = []
                    continue

                collect_time = time.time() - collect_start
                total_collect_time += collect_time

                # TIME: Tensor stacking
                stack_start = time.time()
                batch_x_tensor = torch.stack(batch_x)
                batch_y_tensor = torch.stack(batch_y)
                stack_time = time.time() - stack_start
                total_stack_time += stack_time

                # TIME: Queue insertion
                queue_start = time.time()
                try:
                    # Check pause one more time before putting in queue
                    if pause_event.is_set():
                        batch_x = []
                        batch_y = []
                        continue

                    # Use blocking put initially, then non-blocking
                    if batch_count < 10:  # First 10 batches use blocking to ensure startup
                        data_queue.put((batch_x_tensor, batch_y_tensor), timeout=5.0)
                    else:
                        data_queue.put((batch_x_tensor, batch_y_tensor), timeout=1.0)
                    queue_time = time.time() - queue_start
                    total_queue_time += queue_time
                except Exception as e:
                    # If queue is full or we're pausing, just skip this batch
                    if not pause_event.is_set() and batch_count < 10:
                        print(f"Producer {producer_id} queue error during startup: {e}")
                    pass

                # Reset batch lists
                batch_x = []
                batch_y = []

                batch_count += 1

                # Print timing diagnostics every 50 batches
                if batch_count % 50 == 0 and not pause_event.is_set():
                    avg_collect = total_collect_time / batch_count
                    avg_stack = total_stack_time / batch_count
                    avg_queue = total_queue_time / batch_count
                    if FLAG_PRINT_DATA_PRODUCER_DEBUG:
                        print(f"Producer {producer_id} (batch_size={batch_size}): "
                              f"collect={avg_collect:.3f}s, stack={avg_stack:.3f}s, queue={avg_queue:.3f}s per batch")

            except StopIteration:
                # For validation mode, restart iterator
                if mode == 'val':
                    data_iter = iter(dataset)
                else:
                    # For training, the dataset should be infinite, but just in case
                    continue

    except Exception as e:
        print(f"Producer {producer_id} error: {e}")
    finally:
        print(f"Producer {producer_id} shutting down")


# NEW: Consumer class that manages GPU transfer and provides batches - MODIFIED for multi-target
class AsyncDataLoader:
    def __init__(
            self,
            base_dir,
            asset_name,
            forecast_distances,  # CHANGED: Now list instead of single distance
            window_size,
            mode,
            batch_size,
            val_cutoff_date,
            validation_timescale=None,
            validation_distance=None,  # NEW: For validation-specific distance
            intervals=None,
            weight_power=1.0,
            target_interval=None,
            target_boost=None,
            num_producers=None,
            queue_size=None,
            gpu_buffer_size=None,
            device='cuda'
    ):
        self.device = device
        self.batch_size = batch_size
        self.mode = mode

        # HARD LIMITS - Never exceed these regardless of batch size
        MAX_PRODUCERS = 2  # HARD LIMIT: Never more than 2 processes
        MAX_QUEUE_SIZE = 20  # INCREASED: Was 10, now 20 for better buffering
        MAX_GPU_BUFFER = 8  # INCREASED: Was 3, now 8 for better GPU buffering

        # Conservative scaling - much more limited
        if num_producers is None:
            # CONSERVATIVE: Only 1-2 producers regardless of batch size
            self.num_producers = 1 if batch_size <= 16 else 2
        else:
            self.num_producers = min(num_producers, MAX_PRODUCERS)

        if queue_size is None:
            # CONSERVATIVE: Small queue regardless of batch size
            self.queue_size = min(5 if batch_size <= 16 else 8, MAX_QUEUE_SIZE)
        else:
            self.queue_size = min(queue_size, MAX_QUEUE_SIZE)

        if gpu_buffer_size is None:
            # INCREASED: More generous GPU buffer for better performance
            self.gpu_buffer_size = min(4 if batch_size <= 16 else 6, MAX_GPU_BUFFER)
        else:
            self.gpu_buffer_size = min(gpu_buffer_size, MAX_GPU_BUFFER)

        print(
            f"SAFE AsyncDataLoader: batch_size={batch_size}, producers={self.num_producers}, queue_size={self.queue_size}, gpu_buffer={self.gpu_buffer_size}")

        # Estimate memory usage and warn if too high
        estimated_ram_gb = self._estimate_memory_usage(forecast_distances)
        if estimated_ram_gb > 8.0:  # Warning if over 8GB
            print(f"⚠️  WARNING: Estimated RAM usage: {estimated_ram_gb:.1f}GB")
            print("⚠️  Consider reducing batch size if system becomes unstable")

        # Store parameters for batch size changes
        self.base_dir = base_dir
        self.asset_name = asset_name
        self.forecast_distances = forecast_distances  # CHANGED
        self.window_size = window_size
        self.val_cutoff_date = val_cutoff_date
        self.validation_timescale = validation_timescale
        self.validation_distance = validation_distance  # NEW
        self.intervals = intervals
        self.weight_power = weight_power
        self.target_interval = target_interval
        self.target_boost = target_boost

        # Multiprocessing setup
        self.data_queue = Queue(maxsize=self.queue_size)
        self.control_event = Event()
        self.pause_event = Event()
        self.producers = []

        # GPU buffer (deque for efficient pop/append)
        self.gpu_buffer = deque(maxlen=self.gpu_buffer_size)
        self.buffer_lock = threading.Lock()

        # Start producer processes
        self._start_producers()

        # Start GPU transfer thread
        self.gpu_thread = threading.Thread(target=self._gpu_transfer_worker, daemon=True)
        self.gpu_thread.start()

        print(f"AsyncDataLoader started with {self.num_producers} producers, batch_size={batch_size}")

        # ADDED: Wait for system to actually be ready before returning
        self._wait_for_system_ready()

    def _wait_for_system_ready(self):
        """Wait for producers and GPU transfer thread to actually start producing batches."""
        print("Waiting for AsyncDataLoader system to be ready...")
        start_time = time.time()

        # Wait for at least 1 batch to be available in GPU buffer
        target_batches = 1
        max_wait_time = 60.0  # 60 seconds max

        while time.time() - start_time < max_wait_time:
            # Check GPU buffer
            with self.buffer_lock:
                gpu_buffer_size = len(self.gpu_buffer)

            if gpu_buffer_size >= target_batches:
                elapsed = time.time() - start_time
                print(f"✅ AsyncDataLoader ready! {gpu_buffer_size} batches in GPU buffer after {elapsed:.1f}s")
                return

            # Show progress every 5 seconds
            elapsed = time.time() - start_time
            if int(elapsed) % 5 == 0 and elapsed >= 5:
                try:
                    cpu_queue_size = self.data_queue.qsize()
                except:
                    cpu_queue_size = "unknown"
                print(
                    f"Still warming up... {elapsed:.0f}s elapsed, GPU buffer: {gpu_buffer_size}, CPU queue: {cpu_queue_size}")
                time.sleep(1)  # Prevent spam

            time.sleep(0.1)  # Check every 100ms

        # Timeout - but continue anyway
        print(f"⚠️ Timeout waiting for system ready, continuing anyway...")
        with self.buffer_lock:
            final_gpu_size = len(self.gpu_buffer)
        try:
            final_cpu_size = self.data_queue.qsize()
        except:
            final_cpu_size = "unknown"
        print(f"Final state: GPU buffer: {final_gpu_size}, CPU queue: {final_cpu_size}")

    def _warmup_buffer(self):
        """Pre-fill the GPU buffer to avoid cold start delays."""
        print("Warming up AsyncDataLoader buffer...")
        warmup_start = time.time()

        # Wait for buffer to fill to at least 50% capacity
        target_fill = max(1, self.gpu_buffer_size // 2)

        # First wait for producers to start generating data
        print("Waiting for producers to start...")
        data_queue_had_items = False

        for attempt in range(30):  # 3 seconds max to see any queue activity
            try:
                queue_size = self.data_queue.qsize()
                if queue_size > 0:
                    data_queue_had_items = True
                    print(f"Producers active - queue has {queue_size} items")
                    break
            except:
                pass  # qsize() might not be available on all platforms
            time.sleep(0.1)

        if not data_queue_had_items:
            print("⚠️ Warning: No data in queue after 3s")
            # Try to get one item directly to test if producers are working
            try:
                print("Testing direct queue access...")
                test_batch = self.data_queue.get(timeout=5.0)  # Longer timeout
                print("✅ Direct queue access successful - producers are working")
                # Put the batch back (approximately)
                try:
                    self.data_queue.put(test_batch, timeout=1.0)
                except:
                    pass
                data_queue_had_items = True
            except Exception as e:
                print(f"❌ Direct queue access failed: {e}")
                print("Proceeding without warmup...")
                return

        # Now wait for GPU buffer to fill
        print(f"Waiting for GPU buffer to reach {target_fill} batches...")
        buffer_fill_start = time.time()

        while len(self.gpu_buffer) < target_fill:
            time.sleep(0.05)  # Check every 50ms

            # Show progress every 2 seconds
            elapsed = time.time() - buffer_fill_start
            if elapsed > 0 and int(elapsed) % 2 == 0:
                with self.buffer_lock:
                    current_fill = len(self.gpu_buffer)
                print(f"Buffer filling: {current_fill}/{target_fill} batches ({elapsed:.1f}s elapsed)")

            # Timeout after 15 seconds total
            if time.time() - warmup_start > 15.0:
                with self.buffer_lock:
                    current_fill = len(self.gpu_buffer)
                print(f"⚠️ Warmup timeout - buffer has {current_fill} batches (target: {target_fill})")
                if current_fill > 0:
                    print("Proceeding with partial buffer...")
                break

        warmup_time = time.time() - warmup_start
        with self.buffer_lock:
            final_fill = len(self.gpu_buffer)
        print(f"Buffer warmed up in {warmup_time:.2f}s with {final_fill} batches ready")

    def _estimate_memory_usage(self, forecast_distances):
        """Estimate memory usage to prevent system crashes. MODIFIED for multi-target."""
        # Rough estimates (very conservative)
        bytes_per_sample = 100 * 310 * 4  # window_size * features * float32
        # CHANGED: Now we have multiple targets (high/low for each distance)
        targets_per_sample = len(forecast_distances) * 2  # high + low for each distance
        bytes_per_target = targets_per_sample * 4  # float32
        total_bytes_per_sample = bytes_per_sample + bytes_per_target

        samples_per_batch = self.batch_size

        # Memory in queue + GPU buffer
        total_batches = self.queue_size + self.gpu_buffer_size
        total_samples = total_batches * samples_per_batch
        total_bytes = total_samples * total_bytes_per_sample

        # Add overhead for processes, etc.
        total_bytes_with_overhead = total_bytes * 3  # 3x overhead factor

        return total_bytes_with_overhead / (1024 ** 3)  # Convert to GB

    def _start_producers(self):
        """Start all producer processes."""
        for i in range(self.num_producers):
            producer = Process(
                target=data_producer_worker,
                args=(
                    i, self.data_queue, self.control_event, self.pause_event,
                    self.base_dir, self.asset_name, self.forecast_distances, self.window_size,  # CHANGED
                    self.mode, self.val_cutoff_date, self.validation_timescale, self.validation_distance,  # NEW
                    self.intervals, self.weight_power, self.target_interval, self.target_boost,
                    self.batch_size, i * 1000  # seed offset
                )
            )
            producer.start()
            self.producers.append(producer)

    def _gpu_transfer_worker(self):
        """
        Background thread that continuously moves batches from CPU queue to GPU buffer.
        Now properly respects pause events during batch size changes.
        """
        print(f"GPU transfer thread starting...")
        batch_count = 0
        consecutive_timeouts = 0

        while not self.control_event.is_set():
            # Check if we should pause
            if self.pause_event.is_set():
                time.sleep(0.1)
                consecutive_timeouts = 0  # Reset timeout counter during pause
                continue

            try:
                # Get batch from producer queue (with timeout)
                cpu_batch = self.data_queue.get(timeout=1.0)
                consecutive_timeouts = 0  # Reset on successful get

                # Unpack batch
                x_cpu, y_cpu = cpu_batch

                # Check if we should skip this batch (wrong size during transition)
                if x_cpu.shape[0] != self.batch_size:
                    # Skip batches with wrong size during batch size transitions
                    print(f"GPU thread: Skipping batch with wrong size {x_cpu.shape[0]} (expected {self.batch_size})")
                    continue

                # Transfer to GPU
                x_gpu = x_cpu.to(self.device, non_blocking=True)
                y_gpu = y_cpu.to(self.device, non_blocking=True)

                # Add to GPU buffer (thread-safe)
                with self.buffer_lock:
                    if len(self.gpu_buffer) < self.gpu_buffer_size:
                        self.gpu_buffer.append((x_gpu, y_gpu))
                        batch_count += 1

                        # Progress every 100 batches
                        if batch_count % 100 == 0:
                            print(f"GPU thread: Transferred {batch_count} batches, buffer size: {len(self.gpu_buffer)}")
                    else:
                        # Buffer full, wait a bit
                        time.sleep(0.01)

            except Exception as e:
                # Only count consecutive timeouts when not paused
                if not self.pause_event.is_set():
                    consecutive_timeouts += 1

                    # Only log if we're getting many timeouts in a row (not during normal operation)
                    if consecutive_timeouts > 10:
                        print(
                            f"GPU transfer thread: {consecutive_timeouts} consecutive timeouts (this is normal during low activity)")
                        consecutive_timeouts = 0  # Reset to avoid spam

                # Small sleep to prevent tight loop
                time.sleep(0.01)

        print("GPU transfer thread shutting down")

    def change_batch_size(self, new_batch_size):
        """
        Efficiently change batch size without recreating the entire loader.
        Now with smooth handling like startup - no aggressive timeouts!
        """
        if new_batch_size == self.batch_size:
            return  # No change needed

        print(f"\n{'=' * 60}")
        print(f"Initiating batch size change: {self.batch_size} → {new_batch_size}")
        change_start_time = time.time()

        # Step 1: Pause producers and GPU thread
        print("Step 1: Pausing producers and GPU thread...")
        self.pause_event.set()

        # Wait for producers to actually pause (check queue activity)
        pause_confirmed = False
        pause_wait_start = time.time()
        last_queue_size = -1
        stable_count = 0

        print("Waiting for producers to pause...")
        while not pause_confirmed and time.time() - pause_wait_start < 10.0:
            try:
                current_queue_size = self.data_queue.qsize()
                if current_queue_size == last_queue_size:
                    stable_count += 1
                    if stable_count >= 5:  # Queue size stable for 5 checks
                        pause_confirmed = True
                        print(f"✓ Producers paused (queue stable at {current_queue_size} items)")
                else:
                    stable_count = 0
                last_queue_size = current_queue_size
            except:
                # qsize not available, just wait a bit
                pass
            time.sleep(0.1)

        if not pause_confirmed:
            print("⚠️ Producers may not be fully paused, continuing anyway...")

        # Step 2: Drain CPU queue patiently
        print("\nStep 2: Draining CPU queue...")
        drained_count = 0
        drain_start = time.time()

        while True:
            try:
                # Try to get an item with a short timeout
                self.data_queue.get(timeout=0.1)
                drained_count += 1
                if drained_count % 10 == 0:
                    print(f"  Drained {drained_count} batches...")
            except:
                # Queue empty or timeout - check if really empty
                try:
                    if self.data_queue.qsize() == 0:
                        break
                except:
                    # If qsize not available, assume empty after timeout
                    break

            # Safety check - don't drain forever
            if time.time() - drain_start > 10.0:
                print(f"⚠️ Drain timeout after {drained_count} batches")
                break

        # Step 3: Clear GPU buffer
        print("\nStep 3: Clearing GPU buffer...")
        with self.buffer_lock:
            gpu_cleared = len(self.gpu_buffer)
            self.gpu_buffer.clear()

        print(f"✓ Cleared {drained_count} CPU batches and {gpu_cleared} GPU batches")

        # Step 4: Update batch size and limits
        print("\nStep 4: Updating configuration...")
        self.batch_size = new_batch_size

        # HARD LIMITS - Never exceed regardless of batch size
        MAX_PRODUCERS = 2
        MAX_QUEUE_SIZE = 20  # Keep higher limit from your config
        MAX_GPU_BUFFER = 8  # Keep higher limit from your config

        # Conservative recalculation with hard limits
        new_num_producers = min(1 if new_batch_size <= 16 else 2, MAX_PRODUCERS)
        new_queue_size = min(10 if new_batch_size <= 16 else 15, MAX_QUEUE_SIZE)
        new_gpu_buffer_size = min(4 if new_batch_size <= 16 else 6, MAX_GPU_BUFFER)

        self.num_producers = new_num_producers
        self.queue_size = new_queue_size
        self.gpu_buffer_size = new_gpu_buffer_size

        # Estimate memory and warn if too high
        estimated_ram_gb = self._estimate_memory_usage(self.forecast_distances)
        print(
            f"New configuration: producers={self.num_producers}, queue_size={self.queue_size}, gpu_buffer={self.gpu_buffer_size}")
        print(f"Estimated RAM usage: {estimated_ram_gb:.1f}GB")

        if estimated_ram_gb > 12.0:  # Hard limit at 12GB
            print(f"❌ ABORT: Estimated RAM usage exceeds 12GB limit!")
            print("❌ Keeping current batch size to prevent system crash")
            self.pause_event.clear()  # Resume with old settings
            return

        # Resize GPU buffer
        self.gpu_buffer = deque(maxlen=self.gpu_buffer_size)

        # Create new queue with new size
        old_queue = self.data_queue
        self.data_queue = Queue(maxsize=self.queue_size)

        # Step 5: Gracefully shutdown old producers
        print("\nStep 5: Shutting down old producers...")
        shutdown_start = time.time()

        # First try graceful shutdown
        for i, producer in enumerate(self.producers):
            if producer.is_alive():
                print(f"  Waiting for producer {i} (PID: {producer.pid}) to finish...")
                producer.join(timeout=5.0)  # Give each producer 5 seconds

                if producer.is_alive():
                    print(f"  Producer {i} still alive, terminating...")
                    producer.terminate()
                    producer.join(timeout=2.0)

                    if producer.is_alive():
                        print(f"  ⚠️ Producer {i} failed to terminate cleanly")
                else:
                    print(f"  ✓ Producer {i} shut down gracefully")

        shutdown_time = time.time() - shutdown_start
        print(f"Producer shutdown took {shutdown_time:.1f}s")

        # Clear the list
        self.producers = []

        # Step 6: Start new producers
        print(f"\nStep 6: Starting new producers with batch_size={new_batch_size}...")
        try:
            self._start_producers()
            print(f"✓ Started {self.num_producers} new producers")
        except Exception as e:
            print(f"❌ Error starting new producers: {e}")
            raise

        # Step 7: Clear pause and wait for system to be ready (like startup!)
        print("\nStep 7: Resuming system and waiting for readiness...")
        self.pause_event.clear()

        # Wait for producers to fill queue (similar to startup)
        ready_wait_start = time.time()
        target_queue_items = min(2, self.queue_size // 2)  # Wait for at least 2 items or half queue

        print(f"Waiting for at least {target_queue_items} items in queue...")
        while time.time() - ready_wait_start < 30.0:  # 30 second timeout
            try:
                queue_size = self.data_queue.qsize()
                if queue_size >= target_queue_items:
                    print(f"✓ Queue has {queue_size} items")
                    break
            except:
                pass

            # Show progress
            if int(time.time() - ready_wait_start) % 5 == 0:
                try:
                    queue_size = self.data_queue.qsize()
                    print(f"  Queue size: {queue_size}, waiting...")
                except:
                    print(f"  Waiting for queue to fill...")

            time.sleep(0.1)

        # Wait for GPU buffer to have at least 1 batch
        print("Waiting for GPU buffer to be ready...")
        gpu_ready_start = time.time()
        while time.time() - gpu_ready_start < 10.0:
            with self.buffer_lock:
                if len(self.gpu_buffer) > 0:
                    print(f"✓ GPU buffer has {len(self.gpu_buffer)} batches")
                    break
            time.sleep(0.1)

        # Final status
        total_time = time.time() - change_start_time
        try:
            final_queue_size = self.data_queue.qsize()
        except:
            final_queue_size = "unknown"

        with self.buffer_lock:
            final_gpu_size = len(self.gpu_buffer)

        print(f"\n✅ Batch size change completed in {total_time:.1f}s")
        print(f"   New batch size: {new_batch_size}")
        print(f"   CPU queue: {final_queue_size} items")
        print(f"   GPU buffer: {final_gpu_size} batches")
        print(f"{'=' * 60}\n")

    def get_batch(self):
        """
        Get a batch that's already on GPU. Should be fast since system is pre-warmed.
        """
        # Quick check for available batch
        with self.buffer_lock:
            if self.gpu_buffer:
                return self.gpu_buffer.popleft()

        # If buffer empty after warmup, something is wrong - but wait a bit
        print("⚠️ GPU buffer unexpectedly empty after warmup")

        # Wait briefly for GPU transfer thread to catch up
        max_wait_time = 2.0  # Short wait since system should be ready
        wait_start = time.time()

        while time.time() - wait_start < max_wait_time:
            with self.buffer_lock:
                if self.gpu_buffer:
                    elapsed = time.time() - wait_start
                    print(f"✅ Got batch after brief wait {elapsed:.3f}s")
                    return self.gpu_buffer.popleft()
            time.sleep(0.01)

        # Final fallback - should be rare after proper warmup
        print("Using fallback direct transfer")
        try:
            cpu_batch = self.data_queue.get(timeout=5.0)
            x_cpu, y_cpu = cpu_batch

            if x_cpu.shape[0] != self.batch_size:
                print(f"Warning: Wrong batch size {x_cpu.shape[0]}, expected {self.batch_size}")
                return self.get_batch()

            x_gpu = x_cpu.to(self.device, non_blocking=True)
            y_gpu = y_cpu.to(self.device, non_blocking=True)
            return x_gpu, y_gpu
        except Exception as e:
            raise RuntimeError(f"Failed to get batch: {e}")

    def __iter__(self):
        return self

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

    def shutdown(self):
        """Clean shutdown of all processes and threads."""
        print("Shutting down AsyncDataLoader...")
        self.control_event.set()
        self.pause_event.set()

        # Wait for producers to finish
        for producer in self.producers:
            producer.join(timeout=2.0)
            if producer.is_alive():
                producer.terminate()

        # Clear the queue
        while not self.data_queue.empty():
            try:
                self.data_queue.get_nowait()
            except Exception:  # multiprocessing Queue empty
                break

    def __del__(self):
        self.shutdown()


# MODIFIED: Dataset class for multi-target forecasting
class MultiTimescaleParquetDataset(IterableDataset):
    def __init__(self, base_dir, asset_name, forecast_distances, window_size,
                 mode, val_cutoff_date, validation_timescale=None, validation_distance=None,
                 intervals=None, weight_power=1.0, target_interval=None, target_boost=None):
        super().__init__()
        self.base_dir = base_dir
        self.asset_name = asset_name
        self.forecast_distances = forecast_distances  # CHANGED: Now list of distances
        self.window_size = window_size
        self.mode = mode
        self.val_cutoff_date = datetime.fromisoformat(val_cutoff_date)
        self.validation_timescale = validation_timescale
        self.validation_distance = validation_distance  # NEW: For validation-specific distance
        self.intervals = intervals or ['5m', '15m', '30m', '1h', '2h', '4h', '6h', '8h', '12h']
        self.weight_power = weight_power
        self.target_interval = target_interval
        self.target_boost = target_boost
        self.sampler = TimescaleSampler(INTERVAL_COUNTS, weight_power, target_boost, target_interval)

        # Preload data for all intervals
        self.interval_data = {}
        for interval in self.intervals:
            path = os.path.join(self.base_dir, interval, "unified_training_dataset_softclipped.parquet")
            if not os.path.exists(path):
                print(f"⚠️ Skipping missing file: {path}")
                continue

            df = pl.read_parquet(path).sort("timestamp").with_columns([
                pl.col("timestamp").cast(pl.Datetime)
            ])

            if self.mode == "train":
                df = df.filter(pl.col("timestamp") < self.val_cutoff_date)
            elif self.mode == "val":
                df = df.filter(pl.col("timestamp") >= self.val_cutoff_date)

            if df.height < self.window_size:
                print(f"⚠️ Not enough rows in {interval} after filtering; skipping...")
                continue

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

            # CHANGED: Build target columns for all forecast distances
            target_cols = []
            for distance in self.forecast_distances:
                # Handle both with and without 'c' suffix for backward compatibility
                if distance.endswith('c'):
                    distance_clean = distance  # Already has 'c' suffix
                else:
                    distance_clean = f"{distance}c"  # Add 'c' suffix

                high_col = f"target_{self.asset_name}_high_{distance_clean}"
                low_col = f"target_{self.asset_name}_low_{distance_clean}"
                target_cols.extend([high_col, low_col])

            # Check if all required target columns exist
            missing_cols = [col for col in target_cols if col not in df.columns]
            if missing_cols:
                print(f"⚠️ Missing target columns in {interval}: {missing_cols}; skipping...")
                continue

            features = df.select(feature_cols).to_numpy()
            targets = df.select(target_cols).to_numpy()

            self.interval_data[interval] = (features, targets)

        if not self.interval_data:
            raise RuntimeError("No valid datasets loaded.")

        # NEW: Create target index mapping for validation
        self._create_target_index_mapping()

        # Debug: Print loaded target structure
        if self.interval_data:
            sample_interval = next(iter(self.interval_data.keys()))
            sample_features, sample_targets = self.interval_data[sample_interval]
            print(f"📊 Loaded data structure:")
            print(f"   Features per sample: {sample_features.shape[1] if len(sample_features.shape) > 1 else 'N/A'}")
            print(f"   Targets per sample: {sample_targets.shape[1] if len(sample_targets.shape) > 1 else 'N/A'}")
            print(f"   Expected targets: {len(self.forecast_distances) * 2} (distances: {self.forecast_distances})")
            print(f"   Target indices mapping: {self.target_indices}")

    def _create_target_index_mapping(self):
        """
        Create mapping from forecast distances to target column indices.
        This helps us extract specific targets during validation.
        """
        self.target_indices = {}
        for i, distance in enumerate(self.forecast_distances):
            # Each distance has high and low targets
            high_idx = i * 2  # 0, 2, 4, 6, ...
            low_idx = i * 2 + 1  # 1, 3, 5, 7, ...
            self.target_indices[distance] = {'high': high_idx, 'low': low_idx}

        print(f"Target index mapping created: {self.target_indices}")

    def get_validation_target_indices(self, validation_distance):
        """
        Get the target column indices for a specific validation distance.
        Returns tuple of (high_idx, low_idx) for the specified distance.
        """
        if validation_distance not in self.target_indices:
            raise ValueError(
                f"Validation distance '{validation_distance}' not found in forecast distances: {self.forecast_distances}")

        high_idx = self.target_indices[validation_distance]['high']
        low_idx = self.target_indices[validation_distance]['low']
        return high_idx, low_idx

    def __iter__(self):
        if self.mode == "val":
            # Validation is fixed: yield from a single specified timescale
            if self.validation_timescale not in self.interval_data:
                raise RuntimeError(f"Validation timescale '{self.validation_timescale}' not found in loaded data.")
            features, targets = self.interval_data[self.validation_timescale]

            for i in range(self.window_size - 1, len(features)):
                x_window = features[i - self.window_size + 1: i + 1]
                y_target = targets[i]  # All targets for all distances
                yield torch.tensor(x_window, dtype=torch.float32).pin_memory(), torch.tensor(y_target,
                                                                                             dtype=torch.float32).pin_memory()

        else:
            # Training: infinite sampler across mixed intervals
            while True:
                interval = self.sampler.sample()[0]
                features, targets = self.interval_data[interval]

                if len(features) < self.window_size:
                    continue

                i = random.randint(self.window_size - 1, len(features) - 1)
                x_window = features[i - self.window_size + 1: i + 1]
                y_target = targets[i]  # All targets for all distances

                yield torch.tensor(x_window, dtype=torch.float32), torch.tensor(y_target, dtype=torch.float32)


def worker_init_fn(worker_id):
    worker_seed = SEED + worker_id
    np.random.seed(worker_seed)
    random.seed(worker_seed)


# MODIFIED: Update this function to use AsyncDataLoader for training, keep original for validation
def create_streaming_dataloader(
        base_dir,
        asset_name,
        forecast_distances=None,  # CHANGED: Now list instead of single distance
        forecast_distance=None,  # BACKWARD COMPATIBILITY: Keep old parameter
        window_size=20,
        mode='train',
        batch_size=16,
        val_cutoff_date='2024-11-10',
        validation_timescale=None,
        validation_distance=None,  # NEW: For validation-specific distance
        intervals=None,
        weight_power=1.0,
        target_interval=None,
        target_boost=None,
        device='cuda'
):
    # BACKWARD COMPATIBILITY: Handle both old and new parameter styles
    if forecast_distances is None and forecast_distance is not None:
        # Old style: single forecast_distance
        forecast_distances = [f"{forecast_distance}c"]
        print(
            f"⚠️ Using legacy forecast_distance={forecast_distance}, converted to forecast_distances={forecast_distances}")
    elif forecast_distances is None:
        # Default to original behavior
        forecast_distances = ['4c']
        print(f"⚠️ No forecast distances specified, using default: {forecast_distances}")

    # Ensure validation_distance is set
    if validation_distance is None:
        validation_distance = forecast_distances[0]  # Use first distance as default
        print(f"⚠️ No validation_distance specified, using first forecast distance: {validation_distance}")

    if mode == 'train':
        # Use new AsyncDataLoader for training
        return AsyncDataLoader(
            base_dir=base_dir,
            asset_name=asset_name,
            forecast_distances=forecast_distances,  # CHANGED
            window_size=window_size,
            mode=mode,
            batch_size=batch_size,
            val_cutoff_date=val_cutoff_date,
            validation_timescale=validation_timescale,
            validation_distance=validation_distance,  # NEW
            intervals=intervals,
            weight_power=weight_power,
            target_interval=target_interval,
            target_boost=target_boost,
            device=device
        )
    else:
        # Keep original DataLoader for validation (simpler, deterministic)
        dataset = MultiTimescaleParquetDataset(
            base_dir=base_dir,
            asset_name=asset_name,
            forecast_distances=forecast_distances,  # CHANGED
            window_size=window_size,
            mode=mode,
            val_cutoff_date=val_cutoff_date,
            validation_timescale=validation_timescale,
            validation_distance=validation_distance,  # NEW
            intervals=intervals,
            weight_power=weight_power,
            target_interval=target_interval,
            target_boost=target_boost
        )

        return DataLoader(
            dataset,
            batch_size=batch_size,
            shuffle=False,
            drop_last=True,
            num_workers=0,
            worker_init_fn=worker_init_fn
        )


# KEEP: Original create_dataloader function unchanged
def create_dataloader(X, y, batch_size):
    """
    Creates a DataLoader that handles X, y, and optional simulation data while maintaining alignment.

    Args:
        X: Feature data of shape (n_samples, window_size, n_features)
        y: Target data of shape (n_samples, n_targets)
        batch_size: Batch size for training
        device: Device to move tensors to ('cpu' or 'cuda')
    """
    dataloader_device = 'cpu'

    # Convert numpy arrays to tensors if needed
    if isinstance(X, np.ndarray):
        X = torch.tensor(X, dtype=torch.float32)
    if isinstance(y, np.ndarray):
        y = torch.tensor(y, dtype=torch.float32)

    # Move tensors to device
    tensor_X = X.to(dataloader_device)
    tensor_y = y.to(dataloader_device)

    # Create dataset without simulation data
    dataset = TensorDataset(tensor_X, tensor_y)

    # Create dataloader
    dataloader = DataLoader(
        dataset,
        batch_size=batch_size,
        shuffle=True,
        drop_last=True
    )

    # Add a quick verification step for the first batch
    verify_first_batch = False  # Set to True when you want to verify
    if verify_first_batch:
        for batch in dataloader:
            X_batch, y_batch = batch
            print("\nFirst batch verification:")
            print(f"X batch shape: {X_batch.shape}")
            print(f"y batch shape: {y_batch.shape}")
            break
    return dataloader


def get_num_features_from_dataloader(dataloader):
    if hasattr(dataloader, 'get_batch'):
        # AsyncDataLoader
        X, _ = dataloader.get_batch()
        return X.shape[-1]
    else:
        # Regular DataLoader
        first_batch = next(iter(dataloader))
        X, _ = first_batch
        return X.shape[-1]


def get_num_targets_from_dataloader(dataloader):
    """NEW: Get number of targets from dataloader to configure model output size."""
    if hasattr(dataloader, 'get_batch'):
        # AsyncDataLoader
        _, y = dataloader.get_batch()
        return y.shape[-1]
    else:
        # Regular DataLoader
        first_batch = next(iter(dataloader))
        _, y = first_batch
        return y.shape[-1]


def print_sample_alignment(train_loader, n=1):
    print("\n🔍 Sample Alignment Check:")
    if hasattr(train_loader, 'get_batch'):
        # AsyncDataLoader
        for i in range(n):
            x, y = train_loader.get_batch()
            print(f"Feature shape: {x.shape}, Target shape: {y.shape}")
            print(f"Feature (last window): {x[0][-1]}")
            print(f"Target: {y[0]}")
            print(f"Number of targets: {y.shape[-1]} (should be {len(FORECAST_DISTANCES) * 2})")
    else:
        # Regular DataLoader
        for i, (x, y) in enumerate(train_loader):
            print(f"Feature shape: {x.shape}, Target shape: {y.shape}")
            print(f"Feature (last window): {x[0][-1]}")
            print(f"Target: {y[0]}")
            print(f"Number of targets: {y.shape[-1]} (should be {len(FORECAST_DISTANCES) * 2})")
            if i + 1 >= n:
                break


'''
Training - MODIFIED to work with AsyncDataLoader and multi-target
'''


# update batch size based on learning rate descent
def lr_based_batch_size_schedule(current_lr, init_lr, min_lr, base_bs, max_multiplier, shaping_power=1.0):
    current_lr = max(current_lr, min_lr)
    decay_progress = (init_lr - current_lr) / (init_lr - min_lr)
    decay_progress = max(0.0, min(1.0, decay_progress))
    scaled_progress = decay_progress ** shaping_power
    max_bs = base_bs * max_multiplier
    interpolated_bs = base_bs + (max_bs - base_bs) * scaled_progress
    return int(round(interpolated_bs / 2) * 2)


def update_optimizer_and_scheduler(optimizer, scheduler, new_lr, current_metric):
    # Update optimizer's learning rate
    for param_group in optimizer.param_groups:
        param_group['lr'] = new_lr
    # Reset scheduler state
    scheduler.num_bad_epochs = 0
    scheduler.best = current_metric
    scheduler._last_lr = [new_lr]  # Update the last learning rate


def save_model(generator, model_name, nrmse_score, forecast_distances, model_dir="hilo_saved_models"):
    """
    Saves the generator model to a folder with the model name and MAE score in the filename.
    MODIFIED: Now includes forecast distances in filename.
    """
    # Ensure the directory exists
    os.makedirs(model_dir, exist_ok=True)

    # Create the filename with the NRMSE score and forecast distances
    distances_str = "_".join(forecast_distances)
    model_filename = f"{model_name}_F{distances_str}_NRMSE_{nrmse_score:.6f}.pt"

    # Save the model state_dict
    save_path = os.path.join(model_dir, model_filename)
    torch.save(generator.state_dict(), save_path)

    print(f"Model saved to: {save_path}")


def log_gradients_to_wandb(model, model_name):
    # Loop over each parameter in the model and log its gradient
    for name, param in model.named_parameters():
        if param.grad is not None:
            # Log the gradient norm (L2 norm) to WandB for each parameter
            wandb.log({f"{model_name}/{name}_grad_norm": param.grad.norm().item()})


'''
TRAINING PROCESS - MODIFIED for AsyncDataLoader and multi-target
'''


# MODIFIED: Model Training loop for multi-target
def train_model(
        generator,
        initial_train_loader,  # ← useful for debug/feature inference
        val_loader,
        validation_distance,  # NEW: For validation-specific distance
        init_train_batch_size,
        max_t_batch_multiplier,
        forecast_distances,  # NEW: For model saving
        lr_g_init,
        psi1, psi2, psi3, psi4,
        loss_scaler,
        total_steps,
        eval_frequency,
        device
):
    generator.train()
    optimizer_g = optim.Adam(generator.parameters(), lr=lr_g_init, betas=(0.5, 0.9))

    '''OLD LR SCHEDULER'''
    # scheduler_g = CosineAnnealingLR(
    #     optimizer_g,
    #     T_max=total_steps,
    #     eta_min=1e-6
    # )

    '''NEW LR SCHEDULER'''
    scheduler_g = GoldilocksZoneScheduler(
        optimizer_g,
        lr_g_init,
        total_steps=total_steps,
        goldilocks_ratio=0.7  # 70% of training in goldilocks zone
    )


    # For AsyncDataLoader, we need to handle batch size changes differently
    current_train_loader = initial_train_loader
    current_batch_size = init_train_batch_size

    step = 0

    # Metrics
    eval_list_rmse, eval_list_mae, eval_list_mape, eval_list_nrmse, eval_list_weighted_mae = [], [], [], [], []

    try:
        while step < total_steps:
            if step % 10 == 0:
                print(f"Step {step}")

            # Grab current LR
            current_lr = optimizer_g.param_groups[0]['lr']
            # Dynamic batch size
            bs = lr_based_batch_size_schedule(
                current_lr=current_lr,
                init_lr=lr_g_init,
                min_lr=1e-5,
                base_bs=init_train_batch_size,
                max_multiplier=max_t_batch_multiplier,
                shaping_power=1.3
            )

            # Handle batch size changes for AsyncDataLoader
            if bs != current_batch_size:
                print(f"Changing batch size from {current_batch_size} to {bs}")
                # Use the efficient batch size change method
                current_train_loader.change_batch_size(bs)
                current_batch_size = bs

            # Get batch (already on GPU!)
            with timed_block("Data Loading"):
                batch_start = time.time()
                X_batch, y_batch = current_train_loader.get_batch()
                data_load_time = time.time() - batch_start

                # Log data loading performance every 10 steps
                if FLAG_PRINT_TIMER_BLOCKS:
                    if step % 10 == 0:
                        print(f"Data loading took {data_load_time:.4f}s")

            # Apply noise if needed
            if FLAG_NOISE_INPUTS:
                noise = torch.randn_like(X_batch) * NOISE_LEVEL
                X_batch = torch.clamp(X_batch + noise, -1.0, 1.0)
            if FLAG_NOISE_LABELS:
                noise_std = NOISE_LEVEL * y_batch.std()
                y_batch = torch.clamp(y_batch + torch.randn_like(y_batch) * noise_std, -1.0, 1.0)

            optimizer_g.zero_grad()
            with timed_block("Forward & Loss"):
                with autocast('cuda'):
                    y_pred = generator(X_batch)

                    # CHANGED: Apply loss to all targets (multi-target learning)
                    # In your training loop:
                    custom_loss = forecast_focused_loss(y_batch, y_pred, '4c', forecast_distances)
                    mse = torch.mean((y_batch - y_pred) ** 2)
                    sign = torch.mean(torch.abs(torch.sign(y_batch) - torch.sign(y_pred)))
                    huber = torch.nn.functional.huber_loss(y_batch, y_pred, delta=1.0)
                    final_loss = psi1 * mse + psi2 * (sign / 100) + psi3 * huber + psi4 * custom_loss
                    final_loss = final_loss * loss_scaler

                    if FLAG_DEBUG_PRINT_LOSS:
                        if step % 5000 == 0:
                            print(f"Custom Loss: {custom_loss:.4f}")
                            print(f"MSE Loss: {mse:.4f}")
                            print(f"Sign Loss: {sign:.4f}")
                            print(f"Huber Loss: {huber:.4f}")
                            print(f"Final Loss: {final_loss:.4f}")

            with timed_block("Backward & Step"):
                scaler.scale(final_loss).backward()
                scaler.step(optimizer_g)
                scaler.update()
                scheduler_g.step()

            if FLAG_DEBUG_GRADIENT_FLOW:
                if step % 5000 == 0:
                    for name, param in generator.named_parameters():
                        if param.grad is None:
                            print(f"No gradient for {name}")
                        else:
                            print(f"{name} gradient norm: {param.grad.norm().item()}")
                    print('\n')
            if FLAG_LOG_GRADIENTS_WANDB:
                log_gradients_to_wandb(generator, "Generator")

            # Log training step with GPU monitoring
            if FLAG_DEBUG_GPU_PROFILING:
                if step % 10 == 0:
                    try:
                        # Simple GPU stats check
                        import subprocess
                        result = subprocess.run(
                            ['nvidia-smi', '--query-gpu=utilization.gpu,memory.used,memory.total,power.draw',
                             '--format=csv,noheader,nounits'],
                            capture_output=True, text=True, timeout=1)
                        if result.returncode == 0:
                            stats = result.stdout.strip().split(', ')
                            gpu_util = stats[0]
                            mem_used = stats[1]
                            mem_total = stats[2]
                            power = stats[3]

                            if FLAG_PRINT_GPU_UTILIZATION:
                                print(
                                    f"Step {step}: GPU={gpu_util}%, VRAM={mem_used}MB/{mem_total}MB, Power={power}W, Loss={final_loss.item():.4f}")
                        else:
                            if FLAG_PRINT_TIMER_BLOCKS:
                                print(f"Step {step}: Loss={final_loss.item():.4f} (GPU check failed)")

                    except Exception:
                        if FLAG_PRINT_TIMER_BLOCKS:
                            print(f"Step {step}: Loss={final_loss.item():.4f} (no GPU monitoring)")

            if step % 10 == 0:
                # Just log to wandb for non-10 steps
                wandb.log({
                    "Step": step,
                    "Batch Size": bs,
                    "Learning Rate": current_lr,
                    "Loss": final_loss.item()
                })

            # Periodic validation
            if step > 0 and step % eval_frequency == 0:
                # CHANGED: Pass validation_distance to validation function
                metrics = validate(generator, val_loader, validation_distance, device)
                (avg_rmse, avg_mae, avg_mape, avg_direction_accuracy,
                 avg_high_acc, avg_low_acc, avg_tol_test,
                 avg_uos_high, avg_uos_low,
                 avg_nrmse, avg_weighted_mae) = metrics

                # Save model
                if FLAG_SAVE_MODELS and avg_nrmse < MODEL_SAVE_THRESHOLD:
                    save_model(generator, MODEL_NAME, avg_nrmse, forecast_distances)

                # Log validation metrics
                wandb.log({
                    "Validation RMSE": avg_rmse,
                    "Validation MAE": avg_mae,
                    "Validation MAPE": avg_mape,
                    "Validation NRMSE": avg_nrmse,
                    "Validation Weight MAE": avg_weighted_mae,
                    "Validation DAP": avg_direction_accuracy,
                    "Val High Acc": avg_high_acc,
                    "Val Low Acc": avg_low_acc,
                    "Val Tol Test": avg_tol_test,
                    "Val UOS High": avg_uos_high,
                    "Val UOS Low": avg_uos_low,
                })

                # Accumulate
                eval_list_rmse.append(avg_rmse)
                eval_list_mae.append(avg_mae)
                eval_list_mape.append(avg_mape)
                eval_list_nrmse.append(avg_nrmse)
                eval_list_weighted_mae.append(avg_weighted_mae)

                generator.train()

            step += 1

    except KeyboardInterrupt:
        print("\n[INFO] Training interrupted by user (Ctrl+C). Cleaning up...")
    finally:
        # Clean shutdown
        if hasattr(current_train_loader, 'shutdown'):
            current_train_loader.shutdown()
        wandb.finish()
        torch.cuda.empty_cache()

    return eval_list_rmse, eval_list_mae, eval_list_mape, eval_list_nrmse, eval_list_weighted_mae


# MODIFIED: Validation Loop Function for multi-target with validation distance selection
def validate(generator, validation_loader, validation_distance, device):
    """
    MODIFIED: Now handles multi-target predictions but only evaluates on the specified validation distance.

    Args:
        generator: The model to validate
        validation_loader: DataLoader for validation data
        validation_distance: Specific forecast distance to evaluate (e.g. '4c')
        device: Computing device
    """
    generator.eval()  # Set to evaluation mode
    rmse_list, mae_list, mape_list, direction_accuracy_list = [], [], [], []
    high_acc_list, low_acc_list, acc_tol_test_list = [], [], []
    uos_high_list, uos_low_list, = [], []
    nrmse_list, weighted_mae_list = [], []

    # Get target indices for validation distance
    if hasattr(validation_loader.dataset, 'get_validation_target_indices'):
        high_idx, low_idx = validation_loader.dataset.get_validation_target_indices(validation_distance)
    else:
        # Fallback: calculate indices manually
        # For regular DataLoader, we need to access the dataset differently
        dataset = validation_loader.dataset
        if hasattr(dataset, 'forecast_distances'):
            forecast_distances = dataset.forecast_distances
        else:
            # If it's a single distance (backward compatibility), create indices for position 0
            forecast_distances = [validation_distance]

        if validation_distance not in forecast_distances:
            raise ValueError(
                f"Validation distance '{validation_distance}' not found in forecast distances: {forecast_distances}")

        distance_idx = forecast_distances.index(validation_distance)
        high_idx = distance_idx * 2
        low_idx = distance_idx * 2 + 1

    print(f"🎯 Validating on distance '{validation_distance}' using target indices: high={high_idx}, low={low_idx}")

    with torch.no_grad():  # Disable gradient calculation for validation
        for batch in validation_loader:
            X_val, y_val = batch[:2]  # Always expect X_val and y_val

            # Keep everything on GPU until needed on CPU
            X_val = X_val.to(device)
            y_val = y_val.to(device)

            # Generate predictions for ALL targets (multi-target model output)
            y_pred_all = generator(X_val)

            # CHANGED: Extract only the targets and predictions for validation distance
            # y_val shape: (batch_size, num_targets) where num_targets = len(forecast_distances) * 2
            # y_pred_all shape: (batch_size, num_targets)

            y_val_validation = y_val[:, [high_idx, low_idx]]  # Shape: (batch_size, 2)
            y_pred_validation = y_pred_all[:, [high_idx, low_idx]]  # Shape: (batch_size, 2)

            # Convert outputs to CPU for evaluation metrics
            y_pred_cpu = y_pred_validation.cpu().detach().numpy()
            y_val_cpu = y_val_validation.cpu().detach().numpy()

            if FLAG_PRINT_VALIDATION_PREDICTIONS:
                # Print only the first few samples for comparison
                print('\n')
                print(f"VAL: Real Future Data (First 5 samples) for distance {validation_distance}:")
                print(f"VAL: Fake Predictions (First 5 samples) for distance {validation_distance}:")

                num_samples_to_print = min(5, y_val_cpu.shape[0])

                # Iterate through the first 5 samples and print real and predicted side by side
                y_val_np = y_val_cpu[:num_samples_to_print]
                y_pred_np = y_pred_cpu[:num_samples_to_print]

                for i in range(num_samples_to_print):
                    print(f"  Real: {y_val_np[i]} (high={y_val_np[i][0]:.4f}, low={y_val_np[i][1]:.4f})")
                    print(f"  Pred: {y_pred_np[i]} (high={y_pred_np[i][0]:.4f}, low={y_pred_np[i][1]:.4f})")

                print('\n')

            # Calculate metrics (same as before, but now only on validation distance targets)
            rmse = calculate_rmse(y_val_cpu, y_pred_cpu)
            mae = calculate_mae(y_val_cpu, y_pred_cpu)
            mape = calculate_mape(y_val_cpu, y_pred_cpu)
            nrmse = calculate_nrmse(y_val_cpu, y_pred_cpu)
            weighted_mae = calculate_weighted_mae(y_val_cpu, y_pred_cpu)
            high_acc, low_acc = calculate_high_low_accuracy(y_val_cpu, y_pred_cpu)
            acc_tol_test_score = calculate_predictions_within_accuracy_tolerance(y_val_cpu, y_pred_cpu,
                                                                                 TOL_TEST_ACC_THRESHOLD)
            uos_high, uos_low = calculate_undershoot_overshoot_ratio(y_val_cpu, y_pred_cpu)

            # Calculate direction accuracy: comparing the sign of predictions with actual values
            correct_direction = np.sign(y_val_cpu) == np.sign(y_pred_cpu)
            direction_accuracy = np.mean(correct_direction)

            rmse_list.append(rmse)
            mae_list.append(mae)
            mape_list.append(mape)
            direction_accuracy_list.append(direction_accuracy)
            high_acc_list.append(high_acc)
            low_acc_list.append(low_acc)
            acc_tol_test_list.append(acc_tol_test_score)
            uos_high_list.append(uos_high)
            uos_low_list.append(uos_low)
            nrmse_list.append(nrmse)
            weighted_mae_list.append(weighted_mae)

    avg_rmse = np.mean(rmse_list)
    avg_mae = np.mean(mae_list)
    avg_mape = np.mean(mape_list)
    avg_direction_accuracy = np.mean(direction_accuracy_list) * 100  # Convert to percentage
    avg_high_acc = np.mean(high_acc_list)
    avg_low_acc = np.mean(low_acc_list)
    avg_tol_test = np.mean(acc_tol_test_list)
    avg_uos_high = np.mean(uos_high_list)
    avg_uos_low = np.mean(uos_low_list)
    avg_nrmse = np.mean(nrmse_list)
    avg_weighted_mae = np.mean(weighted_mae_list)

    print(
        f'Validation for {validation_distance}: RMSE: {avg_rmse:.4f}, MAE: {avg_mae:.4f}, MAPE: {avg_mape:.4f}, Sign Accuracy: {avg_direction_accuracy:.4f}')
    return (avg_rmse, avg_mae, avg_mape, avg_direction_accuracy,
            avg_high_acc, avg_low_acc, avg_tol_test, avg_uos_high, avg_uos_low,
            avg_nrmse, avg_weighted_mae)


'''
HYPERPARAMETERS - MODIFIED for multi-target
'''
def get_hyperparameters():
    hyperparams = {
        'asset_name': 'ADAUSDT',
        'total_steps': 4_000_000,
        'eval_frequency': 2_500,
        'forecast_distances': FORECAST_DISTANCES,  # NEW: All forecast distances
        'validation_distance': '4c',  # NEW: Specific distance for validation metrics
        'window_size': 100,  # 20
        'init_train_batch_size': 4,
        'max_train_batch_multiplier': 12,
        'val_batch_size': 16,
        'val_cutoff_date': '2024-11-10',  # or whatever cutoff you want
        'validation_timescale': '1h',
        # Maybe 1? more balance?
        'sampling_weight_curve_power': 0.5,  # INIT: 0.5 | <1 favors balance, >1 favors raw proportions
        # Maybe 1.3?
        'val_target_timescale_boost': 2.0,  # INIT: 2.0 | multiply its sampling weight

        # Learning Rate Config
        'lr_g_init': 0.0003,  # LAST STABLE: 0.0002

        # Hyperparameters for guiding terms in the generator loss
        'psi1': 0.7130,  # Weight for MSE term
        'psi2': 0.4350,  # Weight for sign consistency term 0.434724
        'psi3': 0.1870,  # Weight for huber loss term
        'psi4': 0.6850, # Weight for custom loss
        'loss_scaler': 10
    }
    return hyperparams


if __name__ == '__main__':
    # -------- Start Timer & Precision --------
    main_start_time = time.time()
    np.set_printoptions(suppress=True, precision=8)

    # -------- Device Setup --------
    print(torch.cuda.is_available())
    print(torch.cuda.device_count())
    print(torch.cuda.current_device())
    print(torch.cuda.get_device_name(torch.cuda.current_device()))
    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
    print(f"Using device: {device}")
    torch.cuda.empty_cache()

    # ADD THESE OPTIMIZATIONS:
    # Enable TF32 on Ampere+ GPUs (RTX 3000+)
    torch.backends.cuda.matmul.allow_tf32 = True
    torch.backends.cudnn.allow_tf32 = True

    # Set optimal number of threads
    torch.set_num_threads(torch.get_num_threads())
    torch.set_num_interop_threads(1)

    # -------- Hyperparameters --------
    hyperparams = get_hyperparameters()
    asset_name = hyperparams['asset_name']
    total_steps = hyperparams['total_steps']
    eval_frequency = hyperparams['eval_frequency']
    forecast_distances = hyperparams['forecast_distances']  # CHANGED: Now list
    validation_distance = hyperparams['validation_distance']  # NEW
    window_size = hyperparams['window_size']
    init_train_batch_size = hyperparams['init_train_batch_size']
    max_t_batch_multiplier = hyperparams['max_train_batch_multiplier']
    val_batch_size = hyperparams['val_batch_size']
    lr_g_init = hyperparams['lr_g_init']
    psi1, psi2, psi3, psi4 = (
        hyperparams['psi1'],
        hyperparams['psi2'],
        hyperparams['psi3'],
        hyperparams['psi4'],
    )
    loss_scaler = hyperparams['loss_scaler']
    val_cutoff_date = hyperparams['val_cutoff_date']
    sampling_weight_curve_power = hyperparams['sampling_weight_curve_power']
    validation_timescale = hyperparams['validation_timescale']
    val_target_timescale_boost = hyperparams['val_target_timescale_boost']

    # NEW: Print multi-target configuration
    print(f"\n🎯 Multi-Target Configuration:")
    print(f"   Training on forecast distances: {forecast_distances}")
    print(f"   Total targets: {len(forecast_distances) * 2} (high + low for each distance)")
    print(f"   Validation distance: {validation_distance}")
    print(f"   Validation timescale: {validation_timescale}")
    print()

    # -------- WandB Init --------
    wandb.init(project=WANDB_PROJECT_NAME, config=hyperparams)

    # MODIFIED: Load training using AsyncDataLoader with multiple forecast distances
    initial_train_loader = create_streaming_dataloader(
        base_dir=BASE_DATA_DIR,
        asset_name=asset_name,
        forecast_distances=forecast_distances,  # CHANGED: Now list
        window_size=window_size,
        mode='train',
        batch_size=init_train_batch_size,
        val_cutoff_date=val_cutoff_date,
        validation_distance=validation_distance,  # NEW
        weight_power=sampling_weight_curve_power,
        target_interval=validation_timescale,
        target_boost=val_target_timescale_boost,
        device=device
    )

    # MODIFIED: Load validation using original DataLoader with multiple forecast distances
    val_loader = create_streaming_dataloader(
        base_dir=BASE_DATA_DIR,
        asset_name=asset_name,
        forecast_distances=forecast_distances,  # CHANGED: Now list
        window_size=window_size,
        mode='val',
        batch_size=val_batch_size,
        val_cutoff_date=val_cutoff_date,
        validation_timescale=validation_timescale,
        validation_distance=validation_distance,  # NEW
        device=device
    )

    print_sample_alignment(initial_train_loader, n=1)
    input('Please ensure training and validation targets are aligned, then press Enter to continue...')

    # -------- Model Instantiation --------
    num_features = get_num_features_from_dataloader(initial_train_loader)
    num_targets = get_num_targets_from_dataloader(initial_train_loader)  # NEW
    print(f"Detected number of input features: {num_features}")
    print(f"Detected number of output targets: {num_targets}")
    print(f"Expected targets: {len(forecast_distances) * 2} (distances: {forecast_distances})")

    # Verify target count matches expectations
    expected_targets = len(forecast_distances) * 2
    if num_targets != expected_targets:
        print(f"⚠️ WARNING: Target count mismatch! Got {num_targets}, expected {expected_targets}")
        print("Please check your data or forecast_distances configuration.")
        input("Press Enter to continue anyway, or Ctrl+C to abort...")

    # MODIFIED: Create generator with multi-target support
    generator = create_generator(
    in_channels=num_features,
    device=device,
    forecast_distances=forecast_distances
)
    generator = generator.to(device)
    model_stats = estimate_model_size(generator)

    # Print generator info for verification
    target_info = generator.get_target_info()
    print(f"\n🏗️ Generator Configuration:")
    print(f"   Model outputs: {target_info['num_outputs']}")
    print(f"   Forecast distances: {target_info['forecast_distances']}")
    print(f"   Output structure: {target_info['output_structure'][:3]}...")  # Show first 3

    # -------- Step-Based Training Call --------
    try:
        eval_list_rmse, eval_list_mae, eval_list_mape, eval_list_nrmse, eval_list_weighted_mae = train_model(
            generator,
            initial_train_loader,
            val_loader,
            validation_distance,  # NEW: Pass validation distance
            init_train_batch_size,
            max_t_batch_multiplier,
            forecast_distances,  # NEW: Pass forecast distances for model saving
            lr_g_init,
            psi1, psi2, psi3, psi4,
            loss_scaler,
            total_steps,
            eval_frequency,
            device
        )

        # -------- Summarize --------
        print(f"Minimum RMSE:  {min(eval_list_rmse):.6f}")
        print(f"Minimum MAE:   {min(eval_list_mae):.6f}")
        print(f"Minimum MAPE:  {min(eval_list_mape):.6f}")
        print(f"Minimum NRMSE: {min(eval_list_nrmse):.6f}")
        print(f"(Validation metrics computed on distance: {validation_distance})")

    except KeyboardInterrupt:
        print("Training manually interrupted.")

    # -------- Wrap Up --------
    wandb.finish()
    total_time = time.time() - main_start_time
    print(f"Total time taken: {total_time:.2f} seconds")