#!/usr/bin/env python3
"""
ARES Context Compression Layer - Demo Script

This script demonstrates how to use the compression layer in your orchestration stack.
Run with: python run_demo.py
"""

import torch

from ares_compression import CompressionConfig, CompressionMode, ContextCompressor
from ares_compression.utils import create_sample_context, format_compression_report


def resolve_demo_device() -> str:
    """Prefer CUDA when available, otherwise fall back to CPU."""
    return "cuda" if torch.cuda.is_available() else "cpu"


def demo_basic_usage(device: str) -> None:
    """Demonstrate basic usage with minimal setup."""
    print("\n" + "=" * 60)
    print("DEMO 1: Basic Usage (3-line integration)")
    print("=" * 60)

    compressor = ContextCompressor.from_preset(CompressionMode.BALANCED, device=device)
    context = create_sample_context(seq_len=2048, hidden_dim=768, device=device)
    result = compressor.compress(context)
    print(format_compression_report(result))


def demo_orchestration_integration(device: str) -> None:
    """Demonstrate integration into an orchestration pipeline."""
    print("\n" + "=" * 60)
    print("DEMO 2: Orchestration Stack Integration")
    print("=" * 60)

    config = CompressionConfig(
        mode=CompressionMode.BALANCED,
        max_memory_reduction=0.5,
        max_latency_overhead=0.1,
    )
    compressor = ContextCompressor.from_config(config, device=device)
    prompts = [
        create_sample_context(1024, 768, device),
        create_sample_context(2048, 768, device),
        create_sample_context(4096, 768, device),
    ]

    print("\nProcessing multiple prompts through compression layer:\n")
    for index, prompt_context in enumerate(prompts, start=1):
        print(f"Prompt {index}: {prompt_context.shape[0]} tokens")
        compressed_result = compressor.compress(prompt_context)
        print(f"  -> Compressed to {compressed_result.compressed_context.shape[0]} tokens")
        print(f"  -> Memory saved: {compressed_result.memory_saved:.2f} MB")
        print(f"  -> Latency: {compressed_result.latency_ms:.2f} ms")
        print()

    stats = compressor.get_stats()
    print("Overall Statistics:")
    print(f"  Total compressions: {stats['total_compressions']}")
    print(f"  Avg compression ratio: {stats['avg_compression_ratio']:.1%}")
    print(f"  Total memory saved: {stats['total_memory_saved_mb']:.2f} MB")
    print(f"  Avg latency: {stats['total_latency_ms'] / stats['total_compressions']:.2f} ms")


def demo_custom_configuration(device: str) -> None:
    """Demonstrate custom configuration for specific use cases."""
    print("\n" + "=" * 60)
    print("DEMO 3: Custom Configuration")
    print("=" * 60)

    config = CompressionConfig(
        mode=CompressionMode.CUSTOM,
        semantic_pruning={
            "enabled": True,
            "target_ratio": 0.3,
            "min_tokens": 64,
            "use_attention_scores": True,
            "use_position_aware": True,
        },
        dynamic_precision={
            "enabled": True,
            "default_precision": "int8",
            "precision_thresholds": {
                "fp16": 0.9,
                "int8": 0.6,
                "int4": 0.3,
            },
        },
        entropy_gating={
            "enabled": True,
            "low_entropy_threshold": 0.4,
            "high_entropy_threshold": 0.8,
        },
        tiered_cache={
            "enabled": True,
            "hot_cache_size": 500,
            "warm_cache_size": 3000,
            "cold_cache_size": 8000,
            "eviction_policy": "lru",
        },
        max_memory_reduction=0.6,
    )

    compressor = ContextCompressor.from_config(config, device=device)
    context = create_sample_context(seq_len=4096, hidden_dim=768, device=device)
    result = compressor.compress(context)

    print("Custom Aggressive Configuration Results:")
    print(f"  Original: {result.original_length} tokens")
    print(f"  Compressed: {result.compressed_length} tokens")
    print(f"  Compression ratio: {result.compression_ratio:.1%}")
    print(f"  Memory saved: {result.memory_saved:.2f} MB")


def demo_comparison(device: str) -> None:
    """Compare different compression modes."""
    print("\n" + "=" * 60)
    print("DEMO 4: Mode Comparison")
    print("=" * 60)

    modes = [
        CompressionMode.CONSERVATIVE,
        CompressionMode.BALANCED,
        CompressionMode.AGGRESSIVE,
    ]
    context = create_sample_context(seq_len=2048, hidden_dim=768, device=device)

    print("\nComparing compression modes:\n")
    print(f"{'Mode':<15} {'Ratio':<10} {'Memory (MB)':<15} {'Latency (ms)'}")
    print("-" * 60)

    for mode in modes:
        compressor = ContextCompressor.from_preset(mode, device=device)
        result = compressor.compress(context)
        print(
            f"{mode.value:<15} "
            f"{result.compression_ratio:<10.1%} "
            f"{result.memory_saved:<15.2f} "
            f"{result.latency_ms:.2f}"
        )


def demo_aiiq_bridge() -> None:
    """Demonstrate the orchestration bridge entry point."""
    print("\n" + "=" * 60)
    print("DEMO 5: AIIQ Orchestration Bridge")
    print("=" * 60)

    print(
        """
from pathlib import Path
from ares_compression import build_aiiq_task_bridge

bridge = build_aiiq_task_bridge()
result = bridge.run_task_sync(
    goal="Implement the researched optimization in this repo",
    repo_path=Path(r"Q:\\ARES"),
    timeout_seconds=900,
)
print(result.status, result.summary)
""".strip()
    )


def main() -> None:
    """Run all demos."""
    device = resolve_demo_device()
    print("\n" + "=" * 60)
    print("ARES Context Compression Layer - Demo Suite")
    print("=" * 60)
    print(f"Device: {device}")

    try:
        if not torch.cuda.is_available():
            print("\nWarning: CUDA not available. Falling back to CPU.")
            print("Performance will be slower. For best results, use GPU.\n")

        demo_basic_usage(device)
        demo_orchestration_integration(device)
        demo_custom_configuration(device)
        demo_comparison(device)
        demo_aiiq_bridge()

        print("\n" + "=" * 60)
        print("All demos completed successfully!")
        print("=" * 60)
        print("\nNext Steps:")
        print("1. Install: pip install -e .")
        print("2. Import: from ares_compression import ContextCompressor")
        print("3. Use: compressor = ContextCompressor.from_preset('balanced')")
        print("4. Integrate: result = compressor.compress(your_context)")
        print("\nSee README.md for detailed documentation.\n")
    except Exception as exc:
        print(f"\nError running demo: {exc}")
        print("\nTroubleshooting:")
        print("- Ensure PyTorch is installed: pip install torch")
        print("- Ensure CUDA is available (for GPU acceleration)")
        print("- Check dependencies: pip install -e .")


if __name__ == "__main__":
    main()
