mg_datastreamer_v3.py

"""
High-Performance Cryptocurrency Data Streamer for RL Trading Simulations

This module provides a fast, memory-efficient data streaming solution for large
cryptocurrency datasets using Polars and producer/consumer patterns.
"""

import cProfile
import pstats
import time
from collections import deque
from dataclasses import dataclass
from pathlib import Path
from typing import Optional, Generator, Any
import threading
import queue
from datetime import datetime
import random
import numpy as np

import polars as pl


@dataclass
class StreamConfig:
    """Configuration for the data streamer."""
    dataset_path: str
    chunk_size: int = 10_000
    batch_size: int = 1
    max_timestamp: Optional[str] = None
    overlap_rows: int = 0
    queue_maxsize: int = 100
    verbose: bool = True
    shuffle: bool = False
    random_seed: Optional[int] = None
    use_streaming_engine: bool = True  # Use new streaming engine for memory efficiency
    smart_sampling: bool = False  # Enable frequency-based sampling for RL training
    sampling_bias_power: float = 2.0  # Higher = stronger bias toward unused chunks


class DataStreamer:
    """
    High-performance data streamer for cryptocurrency trading data.

    Uses Polars for fast I/O and implements producer/consumer pattern
    for optimal memory usage and throughput.
    """

    def __init__(self, config: StreamConfig):
        self.config = config
        self.data_queue = queue.Queue(maxsize=config.queue_maxsize)
        self._stop_event = threading.Event()
        self._producer_thread = None
        self._dataset_length = None  # Cache for shuffle mode
        self._chunk_usage_counts = None  # Array to track chunk usage for smart sampling
        self._max_possible_chunks = None  # Total number of possible chunk start positions
        self._stats = {
            'chunks_produced': 0,
            'chunks_consumed': 0,
            'total_rows_processed': 0,
            'start_time': None,
            'end_time': None,
            'unique_chunks_served': 0,
            'avg_chunk_reuse': 0.0
        }

        # Set random seed if provided
        if config.random_seed is not None:
            random.seed(config.random_seed)
            np.random.seed(config.random_seed)

        # Validate dataset exists
        if not Path(config.dataset_path).exists():
            raise FileNotFoundError(f"Dataset not found: {config.dataset_path}")

    def _load_and_filter_data(self) -> pl.LazyFrame:
        """Load data and apply timestamp filtering if specified."""
        if self.config.verbose:
            print(f"Loading dataset: {self.config.dataset_path}")

        # Use lazy loading for memory efficiency
        lf = pl.scan_ipc(self.config.dataset_path)

        # Apply timestamp filter if specified
        if self.config.max_timestamp:
            if self.config.verbose:
                print(f"Filtering data before timestamp: {self.config.max_timestamp}")
            lf = lf.filter(pl.col("timestamp") < self.config.max_timestamp)

        return lf

    def _get_dataset_length(self, lazy_frame: pl.LazyFrame) -> int:
        """Get the total number of rows in the dataset (cached)."""
        if self._dataset_length is None:
            if self.config.verbose:
                print("Calculating dataset length for shuffle mode...")
            # Use a fast count operation
            self._dataset_length = lazy_frame.select(pl.len()).collect().item()
            if self.config.verbose:
                print(f"Dataset contains {self._dataset_length:,} rows")
        return self._dataset_length

    def _initialize_smart_sampling(self, lazy_frame: pl.LazyFrame):
        """Initialize usage tracking for smart sampling."""
        if not self.config.smart_sampling or self._chunk_usage_counts is not None:
            return

        dataset_length = self._get_dataset_length(lazy_frame)
        # Calculate max possible chunk start positions
        self._max_possible_chunks = max(1, dataset_length - self.config.chunk_size + 1)

        # Initialize usage counts (start with zeros)
        self._chunk_usage_counts = np.zeros(self._max_possible_chunks, dtype=np.int32)

        if self.config.verbose:
            print(f"Smart sampling initialized: {self._max_possible_chunks:,} possible chunk positions")

    def _smart_sample_chunk_positions(self, batch_size: int) -> list[int]:
        """Sample chunk positions using frequency-based weighting."""
        if self._chunk_usage_counts is None:
            raise RuntimeError("Smart sampling not initialized")

        # Calculate inverse frequency weights (lower usage = higher weight)
        # Add 1 to avoid division by zero, then invert and apply power
        usage_counts = self._chunk_usage_counts.astype(np.float32)
        min_usage = usage_counts.min()

        # Bias toward less-used chunks: weight = (max_usage - usage + 1) ^ power
        max_usage = usage_counts.max()
        weights = np.power(max_usage - usage_counts + 1, self.config.sampling_bias_power)

        # Normalize weights for probability sampling
        weights = weights / weights.sum()

        # Sample chunk positions without replacement for this batch
        try:
            positions = np.random.choice(
                self._max_possible_chunks,
                size=min(batch_size, self._max_possible_chunks),
                replace=False,
                p=weights
            )
            return positions.tolist()
        except ValueError:
            # Fallback if sampling fails
            return np.random.choice(self._max_possible_chunks, size=batch_size, replace=True).tolist()

    def _smart_sampling_chunk_generator(self, lazy_frame: pl.LazyFrame) -> Generator[pl.DataFrame, None, None]:
        """Generate chunks using smart frequency-based sampling for RL training."""
        self._initialize_smart_sampling(lazy_frame)

        # Generate chunks indefinitely using smart sampling
        while not self._stop_event.is_set():
            # Sample positions for current batch
            positions = self._smart_sample_chunk_positions(self.config.batch_size)

            for start_offset in positions:
                # Fetch chunk at selected position
                if self.config.use_streaming_engine:
                    chunk = lazy_frame.slice(start_offset, self.config.chunk_size).collect(engine='streaming')
                else:
                    chunk = lazy_frame.slice(start_offset, self.config.chunk_size).collect()

                if chunk.height > 0:
                    # Update usage count for this position
                    self._update_chunk_usage([start_offset])
                    yield chunk

    def _update_chunk_usage(self, positions: list[int]):
        """Update usage counts for served chunk positions."""
        if self._chunk_usage_counts is not None:
            for pos in positions:
                if 0 <= pos < len(self._chunk_usage_counts):
                    self._chunk_usage_counts[pos] += 1

    def _chunk_generator(self, lazy_frame: pl.LazyFrame) -> Generator[pl.DataFrame, None, None]:
        """Generate chunks of data with optional overlap or smart sampling."""
        if self.config.smart_sampling:
            return self._smart_sampling_chunk_generator(lazy_frame)
        elif self.config.shuffle:
            return self._shuffled_chunk_generator(lazy_frame)
        else:
            return self._sequential_chunk_generator(lazy_frame)

    def _sequential_chunk_generator(self, lazy_frame: pl.LazyFrame) -> Generator[pl.DataFrame, None, None]:
        """Generate chunks sequentially (original behavior)."""
        offset = 0

        while True:
            # Fetch chunk with new streaming engine
            if self.config.use_streaming_engine:
                chunk = lazy_frame.slice(offset, self.config.chunk_size).collect(engine='streaming')
            else:
                chunk = lazy_frame.slice(offset, self.config.chunk_size).collect()

            if chunk.height == 0:
                break

            yield chunk

            # Calculate next offset with overlap
            if self.config.overlap_rows > 0:
                offset += max(1, self.config.chunk_size - self.config.overlap_rows)
            else:
                offset += self.config.chunk_size

            # Break if we got less than requested (end of data)
            if chunk.height < self.config.chunk_size:
                break

    def _shuffled_chunk_generator(self, lazy_frame: pl.LazyFrame) -> Generator[pl.DataFrame, None, None]:
        """Generate chunks from random positions in the dataset."""
        dataset_length = self._get_dataset_length(lazy_frame)
        max_start_offset = max(0, dataset_length - self.config.chunk_size)

        if max_start_offset <= 0:
            # Dataset is smaller than chunk size, just return the whole thing
            if self.config.use_streaming_engine:
                chunk = lazy_frame.collect(engine='streaming')
            else:
                chunk = lazy_frame.collect()
            if chunk.height > 0:
                yield chunk
            return

        # Pre-generate random offsets for better performance
        rng = random.Random(self.config.random_seed)

        # Generate chunks indefinitely from random positions
        while not self._stop_event.is_set():
            # Random start position
            start_offset = rng.randint(0, max_start_offset)

            # Fetch chunk at random position with new streaming engine
            if self.config.use_streaming_engine:
                chunk = lazy_frame.slice(start_offset, self.config.chunk_size).collect(engine='streaming')
            else:
                chunk = lazy_frame.slice(start_offset, self.config.chunk_size).collect()

            if chunk.height == 0:
                break

            yield chunk

            # In shuffle mode, we continue indefinitely until stopped
            # This allows for unlimited random sampling

    def _producer(self):
        """Producer thread function - generates and queues data chunks."""
        try:
            lazy_frame = self._load_and_filter_data()

            batch = []
            for chunk in self._chunk_generator(lazy_frame):
                if self._stop_event.is_set():
                    break

                batch.append(chunk)

                if len(batch) >= self.config.batch_size:
                    # Put batch in queue (blocks if queue is full)
                    self.data_queue.put(batch.copy())
                    self._stats['chunks_produced'] += len(batch)

                    if self.config.verbose and self._stats['chunks_produced'] % 10 == 0:
                        print(f"Produced {self._stats['chunks_produced']} chunks")

                    batch.clear()

            # Put remaining chunks if any
            if batch:
                self.data_queue.put(batch)
                self._stats['chunks_produced'] += len(batch)

            # Signal end of data
            self.data_queue.put(None)

        except Exception as e:
            print(f"Producer error: {e}")
            self.data_queue.put(None)

    def start_streaming(self):
        """Start the producer thread."""
        if self._producer_thread and self._producer_thread.is_alive():
            raise RuntimeError("Streaming already started")

        self._stop_event.clear()
        self._stats['start_time'] = time.time()
        self._producer_thread = threading.Thread(target=self._producer, daemon=True)
        self._producer_thread.start()

        if self.config.verbose:
            print("Data streaming started")

    def stop_streaming(self):
        """Stop the streaming process."""
        self._stop_event.set()
        if self._producer_thread:
            self._producer_thread.join(timeout=5.0)
        self._stats['end_time'] = time.time()

        if self.config.verbose:
            print("Data streaming stopped")

    def get_batch(self, timeout: float = 0.1) -> Optional[list[pl.DataFrame]]:
        """
        Get next batch of chunks from the queue.

        Returns:
            List of DataFrames or None if stream ended
        """
        try:
            batch = self.data_queue.get(block=True, timeout=timeout)
            if batch is None:  # End of stream signal
                return None

            self._stats['chunks_consumed'] += len(batch)
            total_rows = sum(chunk.height for chunk in batch)
            self._stats['total_rows_processed'] += total_rows

            # Mark task as done for queue management
            self.data_queue.task_done()

            return batch

        except queue.Empty:
            return []  # Timeout, but stream might continue

    def consume_all(self, process_fn=None):
        """
        Consume all data from the stream.

        Args:
            process_fn: Optional function to process each batch
        """
        if not self._producer_thread or not self._producer_thread.is_alive():
            self.start_streaming()

        while True:
            batch = self.get_batch()

            if batch is None:  # End of stream
                break
            elif batch:  # Non-empty batch
                if process_fn:
                    process_fn(batch)
                # Simulate processing by immediately discarding
                del batch

    def get_stats(self) -> dict:
        """Get streaming statistics."""
        stats = self._stats.copy()

        if stats['start_time'] and stats['end_time']:
            duration = stats['end_time'] - stats['start_time']
            stats['duration_seconds'] = duration
            stats['rows_per_second'] = stats['total_rows_processed'] / duration if duration > 0 else 0
            stats['chunks_per_second'] = stats['chunks_consumed'] / duration if duration > 0 else 0

        # Add smart sampling statistics
        if self.config.smart_sampling and self._chunk_usage_counts is not None:
            usage_counts = self._chunk_usage_counts
            stats['unique_chunks_served'] = np.count_nonzero(usage_counts)
            stats['total_possible_chunks'] = len(usage_counts)
            stats['chunk_coverage_percent'] = (stats['unique_chunks_served'] / len(usage_counts)) * 100
            stats['avg_chunk_reuse'] = usage_counts.mean()
            stats['max_chunk_reuse'] = usage_counts.max()
            stats['min_chunk_reuse'] = usage_counts.min()

            # Distribution of usage
            unused = np.count_nonzero(usage_counts == 0)
            used_once = np.count_nonzero(usage_counts == 1)
            used_multiple = np.count_nonzero(usage_counts > 1)

            stats['chunks_unused'] = unused
            stats['chunks_used_once'] = used_once
            stats['chunks_used_multiple'] = used_multiple

        return stats


def benchmark_consumer(batch: list[pl.DataFrame]):
    """Example consumer function for benchmarking."""
    # Simulate minimal processing
    total_rows = sum(chunk.height for chunk in batch)
    # Could add actual processing logic here
    pass


def profile_streaming(dataset_path: str,
                     chunk_size: int = 10_000,
                     batch_size: int = 1,
                     max_timestamp: Optional[str] = None,
                     max_chunks: Optional[int] = None,
                     shuffle: bool = False,
                     smart_sampling: bool = False,
                     sampling_bias_power: float = 2.0,
                     random_seed: Optional[int] = None):
    """Profile the streaming performance."""

    config = StreamConfig(
        dataset_path=dataset_path,
        chunk_size=chunk_size,
        batch_size=batch_size,
        max_timestamp=max_timestamp,
        shuffle=shuffle,
        smart_sampling=smart_sampling,
        sampling_bias_power=sampling_bias_power,
        random_seed=random_seed,
        verbose=True
    )

    streamer = DataStreamer(config)

    def consume_with_limit():
        streamer.start_streaming()
        chunks_processed = 0

        while True:
            batch = streamer.get_batch()

            if batch is None:  # End of stream
                break
            elif batch:  # Non-empty batch
                benchmark_consumer(batch)
                chunks_processed += len(batch)

                if max_chunks and chunks_processed >= max_chunks:
                    break

        streamer.stop_streaming()

    # Profile the consumption
    pr = cProfile.Profile()
    pr.enable()

    consume_with_limit()

    pr.disable()

    # Print stats
    stats = streamer.get_stats()
    print("\n" + "="*50)
    print("STREAMING PERFORMANCE STATS")
    print("="*50)
    for key, value in stats.items():
        print(f"{key}: {value}")

    print("\n" + "="*50)
    print("PROFILING RESULTS")
    print("="*50)

    # Sort by cumulative time
    ps = pstats.Stats(pr).sort_stats('cumulative')
    ps.print_stats(20)  # Top 20 functions

    return streamer, stats


if __name__ == "__main__":
    # Configuration
    DATASET_PATH = "datasets/arrow_1m/unified_1m_dataset.arrow"

    # Toggle these to test different modes
    USE_SHUFFLE = False         # Set to True for random sampling
    USE_SMART_SAMPLING = True   # Set to True for RL-optimized frequency-based sampling

    # Smart sampling configuration
    SAMPLING_BIAS_POWER = 2.0   # Higher = stronger bias toward unused chunks (1.0 = linear, 2.0 = quadratic)

    # Example usage and benchmarking
    if USE_SMART_SAMPLING:
        mode = "smart sampling (RL optimized)"
    elif USE_SHUFFLE:
        mode = "random shuffle"
    else:
        mode = "sequential"

    print(f"Starting high-performance crypto data streaming benchmark ({mode})...")

    try:
        streamer, stats = profile_streaming(
            dataset_path=DATASET_PATH,
            chunk_size=10_000,
            batch_size=2,
            max_chunks=100,  # Increased for better smart sampling demo
            shuffle=USE_SHUFFLE,
            smart_sampling=USE_SMART_SAMPLING,
            sampling_bias_power=SAMPLING_BIAS_POWER,
            random_seed=42
        )

        print(f"\nProcessed {stats['total_rows_processed']:,} rows")
        print(f"Throughput: {stats.get('rows_per_second', 0):,.0f} rows/second")

        # Print smart sampling statistics
        if USE_SMART_SAMPLING:
            print(f"\n📊 SMART SAMPLING STATS:")
            print(f"Unique chunks served: {stats.get('unique_chunks_served', 0):,}")
            print(f"Total possible chunks: {stats.get('total_possible_chunks', 0):,}")
            print(f"Coverage: {stats.get('chunk_coverage_percent', 0):.1f}%")
            print(f"Avg reuse per chunk: {stats.get('avg_chunk_reuse', 0):.2f}")
            print(f"Unused chunks: {stats.get('chunks_unused', 0):,}")
            print(f"Used once: {stats.get('chunks_used_once', 0):,}")
            print(f"Used multiple times: {stats.get('chunks_used_multiple', 0):,}")

    except FileNotFoundError as e:
        print(f"Error: {e}")
        print("Please ensure the dataset path is correct.")
    except Exception as e:
        print(f"Unexpected error: {e}")