game.py

"""
Stage 3: Multi-Stream Token Architecture with Hybrid Tokenization

Event-driven architecture demonstrating:
- Configuration-based tokenization (no hardcoded values)
- Service-oriented design (TokenizationService)
- Efficient token storage (integer IDs, not strings)
- Hybrid tokenization (symbolic + continuous)
- Episode management for training
- Token lineage tracking

The game loop is now a simple coordinator - it doesn't know about
tokenization details, just routes events to services.
"""

import pygame
import sys
import time
from typing import Set, Optional
from core import GameEngine
from core.config import TokenizerConfig
from core.tokenizer import GameTokenizer
from core.token_store import MultiStreamStore, StreamType
from core.tokenization_service import TokenizationService
from core.events import InputChangeEvent

# Hybrid tokenization components
from core.storage import EpisodeManager, ContinuousStore
from core.lineage import ProvenanceTracker
from core.hybrid_tokenizer import HybridTokenizer, HybridTokenizerConfig


class MultiStreamPanel:
    """Visualization showing all three token streams."""

    def __init__(self, x: int, y: int, width: int, height: int):
        self.x, self.y = x, y
        self.width, self.height = width, height
        self.bg_color = (20, 20, 20)
        self.border_color = (100, 100, 100)

        # Fonts
        self.title_font = pygame.font.Font(None, 28)
        self.stream_title_font = pygame.font.Font(None, 22)
        self.token_font = pygame.font.Font(None, 16)
        self.stat_font = pygame.font.Font(None, 18)

        # Stream colors
        self.stream_colors = {
            StreamType.STATIC: (150, 150, 255),   # Blue
            StreamType.DYNAMIC: (255, 200, 100),  # Orange
            StreamType.EVENT: (100, 255, 100)     # Green
        }

    def draw(self, surface: pygame.Surface, token_store: MultiStreamStore,
             tokenizer: GameTokenizer) -> None:
        """Draw multi-stream visualization."""
        # Background
        panel_rect = pygame.Rect(self.x, self.y, self.width, self.height)
        pygame.draw.rect(surface, self.bg_color, panel_rect)
        pygame.draw.rect(surface, self.border_color, panel_rect, 2)

        # Title
        title = self.title_font.render("MULTI-STREAM TOKENS", True, (255, 255, 100))
        surface.blit(title, (self.x + 10, self.y + 10))

        # Statistics
        y_offset = self.y + 45
        stats = token_store.get_stats()

        stat_lines = [
            f"Total: {stats['total_generated']} tokens",
            f"Rate: {stats['tokens_per_second']:.1f} tok/s",
            f"Static: {stats['stream_counts']['static']} | " +
            f"Dynamic: {stats['stream_counts']['dynamic']} | " +
            f"Event: {stats['stream_counts']['event']}"
        ]

        for line in stat_lines:
            text = self.stat_font.render(line, True, (200, 200, 200))
            surface.blit(text, (self.x + 10, y_offset))
            y_offset += 20

        # Separator
        y_offset += 5
        pygame.draw.line(surface, self.border_color,
                        (self.x + 10, y_offset),
                        (self.x + self.width - 10, y_offset), 1)
        y_offset += 10

        # Three columns
        col_width = (self.width - 40) // 3
        stream_height = self.height - (y_offset - self.y) - 10

        self._draw_stream_column(surface, self.x + 10, y_offset, col_width,
                                 stream_height, "STATIC",
                                 token_store.get_static_tokens(),
                                 StreamType.STATIC, tokenizer)

        self._draw_stream_column(surface, self.x + 20 + col_width, y_offset,
                                 col_width, stream_height, "DYNAMIC",
                                 token_store.get_dynamic_tokens(20),
                                 StreamType.DYNAMIC, tokenizer)

        self._draw_stream_column(surface, self.x + 30 + col_width * 2, y_offset,
                                 col_width, stream_height, "EVENT",
                                 token_store.get_event_tokens(20),
                                 StreamType.EVENT, tokenizer)

    def _draw_stream_column(self, surface, x, y, width, height, title,
                           tokens, stream_type, tokenizer):
        """Draw a single stream column."""
        # Column border
        col_rect = pygame.Rect(x, y, width, height)
        pygame.draw.rect(surface, (30, 30, 30), col_rect)
        pygame.draw.rect(surface, self.stream_colors[stream_type], col_rect, 1)

        # Title
        title_text = self.stream_title_font.render(title, True, self.stream_colors[stream_type])
        surface.blit(title_text, (x + 5, y + 5))

        # Count
        count_text = self.token_font.render(f"({len(tokens)})", True, (150, 150, 150))
        surface.blit(count_text, (x + 5, y + 25))

        # Tokens
        token_y = y + 45
        for token in reversed(tokens):
            if token_y + 16 > y + height:
                break

            # Decode IDs to names
            entity_name = tokenizer.vocab.get_entity_name(token.entity_id)
            property_name = tokenizer.vocab.get_property_name(token.property_id)
            value = token.value

            # Format value
            if isinstance(value, tuple):
                value_str = f"({value[0]:.0f},{value[1]:.0f})"
            elif isinstance(value, bool):
                value_str = "↓" if value else "↑"
            else:
                value_str = str(value)

            token_str = f"{entity_name}.{property_name}={value_str}"
            if len(token_str) > width // 7:
                token_str = token_str[:width // 7 - 3] + "..."

            text = self.token_font.render(token_str, True, self.stream_colors[stream_type])
            surface.blit(text, (x + 8, token_y))
            token_y += 16


class MultiStreamGame:
    """
    Event-driven game with proper service architecture.

    The game loop is now simple - it just routes events.
    No knowledge of tokenization internals, just dependency injection.
    Enhanced with hybrid tokenization for Stage 3 completion.
    """

    def __init__(self, width: int = 1200, height: int = 600, enable_hybrid: bool = True):
        pygame.init()

        self.game_width = 800
        self.panel_width = width - self.game_width
        self.height = height

        self.screen = pygame.display.set_mode((width, height))
        caption = "Diabl0 - Stage 3: Multi-Stream"
        if enable_hybrid:
            caption += " + Hybrid"
        pygame.display.set_caption(caption)

        self.clock = pygame.font.Font(None, 24)
        self.running = True
        self.enable_hybrid = enable_hybrid

        # Configuration (data-driven, no hardcoded values)
        config = TokenizerConfig.create_default()

        # Core components (dependency injection)
        self.engine = GameEngine(arena_width=self.game_width, arena_height=height)
        self.tokenizer = GameTokenizer(config)
        self.token_store = MultiStreamStore()

        # Hybrid components (if enabled)
        self.episode_manager: Optional[EpisodeManager] = None
        self.continuous_store: Optional[ContinuousStore] = None
        self.provenance_tracker: Optional[ProvenanceTracker] = None
        self.hybrid_tokenizer: Optional[HybridTokenizer] = None
        self.current_episode_id: Optional[int] = None

        if enable_hybrid:
            # Initialize hybrid components (pass shared token store!)
            self.episode_manager = EpisodeManager(token_store=self.token_store)
            self.continuous_store = self.episode_manager.continuous_store
            self.provenance_tracker = ProvenanceTracker()

            # Create hybrid tokenizer
            hybrid_config = HybridTokenizerConfig()
            self.hybrid_tokenizer = HybridTokenizer(hybrid_config)

        # Service layer (single responsibility: events → tokens)
        self.tokenization_service = TokenizationService(
            self.tokenizer,
            self.token_store,
            enable_hybrid=enable_hybrid,
            episode_manager=self.episode_manager,
            continuous_store=self.continuous_store,
            provenance_tracker=self.provenance_tracker
        )

        # UI
        self.panel = MultiStreamPanel(
            x=self.game_width, y=0,
            width=self.panel_width, height=height
        )

        # Input tracking for events
        self.key_map = {
            pygame.K_w: 'W',
            pygame.K_a: 'A',
            pygame.K_s: 'S',
            pygame.K_d: 'D',
            pygame.K_e: 'E'  # E key to start new episode
        }
        self._last_keys: Set[str] = set()

        # Start first episode if hybrid mode
        if enable_hybrid:
            self._start_new_episode()

    def handle_events(self) -> None:
        """Process pygame events."""
        events = pygame.event.get()
        for event in events:
            if event.type == pygame.QUIT:
                self.running = False
            if event.type == pygame.KEYDOWN:
                if event.key == pygame.K_ESCAPE:
                    self.running = False
                # E key starts new episode in hybrid mode
                elif event.key == pygame.K_e and self.enable_hybrid:
                    self._complete_and_start_episode()

        self.engine.handle_input(events)

    def _start_new_episode(self) -> None:
        """Start a new episode for hybrid tokenization."""
        if not self.enable_hybrid:
            return

        self.current_episode_id = self.tokenization_service.start_episode(
            description=f"Episode at frame {self.engine.frame_count}"
        )
        print(f"\n[Episode {self.current_episode_id}] Started")

    def _complete_and_start_episode(self) -> None:
        """Complete current episode and start a new one."""
        if not self.enable_hybrid or not self.current_episode_id:
            return

        # Complete current episode
        stats = self.tokenization_service.complete_episode()
        print(f"[Episode {self.current_episode_id}] Completed - "
              f"{stats.get('total_tokens', 0)} tokens, "
              f"{stats.get('total_frames', 0)} frames")

        # Start new episode
        self._start_new_episode()

    def update(self, dt: float) -> None:
        """
        Update game - now just a simple event coordinator.

        No tokenization logic here - just route events to services.
        """
        # Get events from engine
        events = self.engine.update(dt)

        # Add input change events
        events.extend(self._get_input_events())

        # Services handle events (decoupled)
        for event in events:
            self.tokenization_service.handle_event(event)

    def _get_input_events(self):
        """Convert pygame input to InputChangeEvent."""
        keys = pygame.key.get_pressed()
        current = {self.key_map[k] for k in self.key_map if keys[k]}

        pressed = current - self._last_keys
        released = self._last_keys - current

        self._last_keys = current

        if pressed or released:
            return [InputChangeEvent(
                keys_pressed=pressed,
                keys_released=released,
                frame=self.engine.frame_count,
                timestamp=time.time()
            )]
        return []

    def draw(self) -> None:
        """Render game and UI."""
        self.screen.fill((0, 0, 0))

        # Draw game
        game_surface = self.screen.subsurface((0, 0, self.game_width, self.height))
        self.engine.draw(game_surface)

        # FPS
        fps = pygame.time.Clock().get_fps()
        fps_text = pygame.font.Font(None, 24).render(f"FPS: {fps:.1f}", True, (200, 200, 200))
        self.screen.blit(fps_text, (10, 10))

        # Hybrid mode indicator
        if self.enable_hybrid:
            mode_text = pygame.font.Font(None, 24).render(
                f"HYBRID MODE | Episode: {self.current_episode_id or 0}",
                True, (100, 255, 100)
            )
            self.screen.blit(mode_text, (10, 35))

        # Controls
        control_text = "WASD: Move | ESC: Quit"
        if self.enable_hybrid:
            control_text += " | E: New Episode"
        controls = pygame.font.Font(None, 24).render(control_text, True, (150, 150, 150))
        self.screen.blit(controls, (10, self.height - 30))

        # Separator
        pygame.draw.line(self.screen, (100, 100, 100),
                        (self.game_width, 0), (self.game_width, self.height), 2)

        # Panel
        self.panel.draw(self.screen, self.token_store, self.tokenizer)

        pygame.display.flip()

    def run(self) -> None:
        """Main game loop."""
        print("=" * 70)
        print("Diabl0 - Stage 3: Multi-Stream Token Architecture")
        if self.enable_hybrid:
            print("WITH HYBRID TOKENIZATION ENABLED")
        print("=" * 70)
        print("\nArchitecture improvements:")
        print("  ✓ Event-driven design")
        print("  ✓ Configuration-based tokenization")
        print("  ✓ Service-oriented (TokenizationService)")
        print("  ✓ Efficient storage (integer IDs)")
        if self.enable_hybrid:
            print("\nHybrid Features:")
            print("  ✓ Continuous data capture (HDF5)")
            print("  ✓ Episode management")
            print("  ✓ Token lineage tracking")
            print("  ✓ Symbolic + continuous fusion")
        print("\nControls:")
        print("  WASD - Move")
        print("  ESC  - Quit")
        if self.enable_hybrid:
            print("  E    - Start new episode")
        print("\n" + "=" * 70)

        clock = pygame.time.Clock()
        while self.running:
            dt = clock.tick(60) / 1000.0
            self.handle_events()
            self.update(dt)
            self.draw()

        # Complete final episode if hybrid mode
        if self.enable_hybrid and self.current_episode_id:
            stats = self.tokenization_service.complete_episode()
            print(f"\n[Episode {self.current_episode_id}] Final episode completed")

        # --- Summary ---
        print("\n" + "=" * 70)
        print("SESSION SUMMARY")
        print("=" * 70)
        stats = self.token_store.get_stats()
        print(f"\nTotal tokens: {stats['total_generated']}")
        print(f"Average rate: {stats['tokens_per_second']:.1f} tokens/sec")
        print(f"\nStream breakdown:")
        print(f"  STATIC:  {stats['stream_counts']['static']:>5d} tokens")
        print(f"  DYNAMIC: {stats['stream_counts']['dynamic']:>5d} tokens")
        print(f"  EVENT:   {stats['stream_counts']['event']:>5d} tokens")

        # Top 10 Tokens
        print("\nTop 10 Token Properties:")
        prop_counts = stats.get('property_counts', {})
        if not prop_counts:
            print("  No tokens were generated.")
        else:
            # Sort properties by count, descending
            sorted_props = sorted(prop_counts.items(), key=lambda item: item[1], reverse=True)
            
            # Print top 10
            for i, (prop_id, count) in enumerate(sorted_props[:10]):
                prop_name = self.tokenizer.vocab.get_property_name(prop_id)
                print(f"  {i+1:>2}. {prop_name:<25} {count:>6d} tokens")

        # Hybrid mode summary
        if self.enable_hybrid and self.episode_manager:
            print("\nHybrid Tokenization Summary:")
            episodes = self.episode_manager.list_episodes()
            print(f"  Episodes recorded: {len(episodes)}")
            if episodes:
                total_frames = sum(ep.total_frames for ep in episodes)
                print(f"  Total frames: {total_frames}")
                print(f"  Avg frames/episode: {total_frames // len(episodes)}")

            if self.provenance_tracker:
                hierarchy_stats = self.provenance_tracker.analyze_hierarchy()
                print(f"\nToken Hierarchy:")
                for level, count in hierarchy_stats['by_hierarchy'].items():
                    print(f"  {level}: {count} tokens")

        print("=" * 70)

        pygame.quit()
        sys.exit()


if __name__ == "__main__":
    import argparse

    parser = argparse.ArgumentParser(description="Diabl0 Stage 3 with Hybrid Tokenization")
    parser.add_argument('--no-hybrid', action='store_true',
                        help='Disable hybrid tokenization mode')
    args = parser.parse_args()

    game = MultiStreamGame(enable_hybrid=not args.no_hybrid)
    game.run()