#!/usr/bin/env python3
"""
Complete v1.1.0 Workflow Example

Demonstrates all three major v1.1.0 features working together:
1. AVX-512 SIMD acceleration
2. TBL binary format
3. Prefetching pipeline

This example shows a complete end-to-end ML training pipeline
using all v1.1.0 optimizations for maximum performance.

New in v1.1.0:
- AVX-512 SIMD: 2x vector throughput on compatible hardware
- TBL Format: 12.4% smaller files, O(1) random access
- Prefetching: Overlapped I/O for reduced epoch time

Requirements:
- TurboLoader v1.1.0+
- PyTorch (optional, for tensor conversion)
- Dataset in TAR format
"""

import turboloader
import time
import os
import sys
import numpy as np


def check_v110_features():
    """Check which v1.1.0 features are available"""
    print("=" * 80)
    print("v1.1.0 FEATURE CHECK")
    print("=" * 80)

    # Check version
    version = turboloader.version()
    print(f"TurboLoader version: {version}")

    if version < "1.1.0":
        print("⚠️  WARNING: TurboLoader v1.1.0+ required for all features")
        print("   Please upgrade: pip install --upgrade turboloader")
        return False

    # Check SIMD support
    features = turboloader.features()
    print(f"\nSIMD features: {features}")

    if "AVX-512" in features:
        print("✅ AVX-512 SIMD: ENABLED (16-wide vectors)")
    elif "AVX2" in features:
        print("⚠️  AVX-512: Not available (using AVX2 fallback)")
    elif "NEON" in features:
        print("⚠️  AVX-512: Not available (using NEON fallback)")
    else:
        print("❌ SIMD: Not available (scalar fallback)")

    # Check TBL support
    print("\n✅ TBL Binary Format: AVAILABLE")
    print("   Command-line tool: tar_to_tbl")

    # Check Prefetching
    print("\n✅ Prefetching Pipeline: INTEGRATED")
    print("   Automatically enabled with DataLoader")

    print("=" * 80)
    print()

    return True


def step1_convert_to_tbl(tar_path: str, tbl_path: str):
    """
    Step 1: Convert TAR to TBL format

    Benefits:
    - 12.4% smaller file size
    - O(1) random access
    - Memory-mapped I/O
    """
    print("STEP 1: CONVERT TAR → TBL")
    print("=" * 80)

    if os.path.exists(tbl_path):
        print(f"TBL file already exists: {tbl_path}")
        tbl_size = os.path.getsize(tbl_path) / 1024 / 1024
        print(f"Size: {tbl_size:.2f} MB")
        print("Skipping conversion.")
        print("=" * 80)
        print()
        return tbl_path

    print(f"Input:  {tar_path}")
    print(f"Output: {tbl_path}")
    print()

    # Get TAR size
    tar_size = os.path.getsize(tar_path) / 1024 / 1024
    print(f"TAR size: {tar_size:.2f} MB")
    print()

    # Convert
    print("Converting...")
    start_time = time.time()

    def progress(current, total):
        if current % 500 == 0 or current == total:
            elapsed = time.time() - start_time
            rate = current / elapsed if elapsed > 0 else 0
            percent = (current / total) * 100 if total > 0 else 0
            print(f"  {current}/{total} ({percent:.1f}%) - {rate:.0f} samples/s")

    try:
        turboloader.convert_tar_to_tbl(tar_path, tbl_path, progress_callback=progress)
    except AttributeError:
        # Use command-line tool
        import subprocess

        print("Using command-line converter...")
        result = subprocess.run(["tar_to_tbl", tar_path, tbl_path], capture_output=True, text=True)
        if result.returncode != 0:
            print(f"ERROR: {result.stderr}")
            return None
        print(result.stdout)

    elapsed = time.time() - start_time

    # Get TBL size
    tbl_size = os.path.getsize(tbl_path) / 1024 / 1024
    savings = ((tar_size - tbl_size) / tar_size) * 100

    print()
    print(f"✅ Conversion complete in {elapsed:.2f}s")
    print(f"   TAR: {tar_size:.2f} MB")
    print(f"   TBL: {tbl_size:.2f} MB ({savings:.1f}% smaller)")
    print("=" * 80)
    print()

    return tbl_path


def step2_create_simd_pipeline():
    """
    Step 2: Create SIMD-accelerated transform pipeline

    All transforms use AVX-512 SIMD when available:
    - Resize (bilinear interpolation)
    - RandomHorizontalFlip (memory operations)
    - ColorJitter (brightness/contrast/saturation)
    - Normalize (FMA operations)
    """
    print("STEP 2: CREATE SIMD-ACCELERATED PIPELINE")
    print("=" * 80)

    # Create transforms (all SIMD-accelerated)
    transforms = {
        "resize": turboloader.Resize(224, 224, turboloader.InterpolationMode.BILINEAR),
        "flip": turboloader.RandomHorizontalFlip(p=0.5),
        "color_jitter": turboloader.ColorJitter(brightness=0.2, contrast=0.2, saturation=0.2),
        "normalize": turboloader.ImageNetNormalize(to_float=True),
    }

    print("Created SIMD transforms:")
    for name in transforms:
        print(f"  ✅ {name}")

    print()
    print("All transforms will use:")
    features = turboloader.features()
    if "AVX-512" in features:
        print("  • AVX-512 (16-wide vectors)")
    elif "AVX2" in features:
        print("  • AVX2 (8-wide vectors)")
    elif "NEON" in features:
        print("  • NEON (4-wide vectors)")
    else:
        print("  • Scalar operations")

    print("=" * 80)
    print()

    return transforms


def step3_benchmark_full_pipeline(
    dataset_path: str, transforms: dict, num_workers: int = 8, batch_size: int = 64
):
    """
    Step 3: Benchmark complete v1.1.0 pipeline

    Combines:
    - TBL format (12.4% smaller, O(1) access)
    - SIMD transforms (AVX-512/AVX2/NEON)
    - Prefetching pipeline (overlapped I/O)
    """
    print("STEP 3: BENCHMARK COMPLETE v1.1.0 PIPELINE")
    print("=" * 80)
    print(f"Dataset: {dataset_path}")
    print(f"Workers: {num_workers}")
    print(f"Batch size: {batch_size}")
    print()

    # Create loader (prefetching automatically enabled)
    loader = turboloader.DataLoader(dataset_path, batch_size=batch_size, num_workers=num_workers)

    # Warm-up
    print("Warming up...")
    for i, batch in enumerate(loader):
        if i >= 2:
            break
        for sample in batch:
            img = sample["image"]
            # Apply all transforms
            img = transforms["resize"].apply(img)
            img = transforms["flip"].apply(img)
            img = transforms["color_jitter"].apply(img)
            img = transforms["normalize"].apply(img)

    # Benchmark
    print("Running benchmark...")
    start_time = time.time()
    total_images = 0
    total_batches = 0

    for batch in loader:
        total_batches += 1
        for sample in batch:
            img = sample["image"]

            # Apply full SIMD pipeline
            img = transforms["resize"].apply(img)
            img = transforms["flip"].apply(img)
            img = transforms["color_jitter"].apply(img)
            img = transforms["normalize"].apply(img)

            total_images += 1

        # Limit to first 1000 images for quick demo
        if total_images >= 1000:
            break

    elapsed = time.time() - start_time
    throughput = total_images / elapsed

    print()
    print("=" * 80)
    print("BENCHMARK RESULTS")
    print("=" * 80)
    print(f"Images processed: {total_images}")
    print(f"Batches: {total_batches}")
    print(f"Time: {elapsed:.2f}s")
    print(f"Throughput: {throughput:.1f} img/s")
    print()
    print("Pipeline stages:")
    print("  1. TBL Reader (memory-mapped, O(1) access)")
    print("  2. JPEG Decode (libjpeg-turbo)")
    print("  3. Resize (SIMD interpolation)")
    print("  4. RandomFlip (SIMD memory ops)")
    print("  5. ColorJitter (SIMD brightness/contrast)")
    print("  6. Normalize (SIMD FMA)")
    print()
    print("v1.1.0 features in use:")
    print("  ✅ TBL Format (12.4% smaller)")
    print("  ✅ SIMD Acceleration (AVX-512/AVX2/NEON)")
    print("  ✅ Prefetching Pipeline (overlapped I/O)")
    print("=" * 80)
    print()

    return throughput


def step4_pytorch_integration(dataset_path: str, transforms: dict):
    """
    Step 4: PyTorch integration example

    Shows how to use v1.1.0 features with PyTorch training
    """
    print("STEP 4: PYTORCH INTEGRATION")
    print("=" * 80)

    try:
        import torch
        import torch.nn as nn

        # Create loader
        loader = turboloader.DataLoader(dataset_path, batch_size=32, num_workers=8)

        # Create tensor converter
        to_tensor = turboloader.ToTensor(
            format=turboloader.TensorFormat.PYTORCH_CHW, normalize=True
        )

        # Simulate training loop
        print("Simulating PyTorch training loop...")
        print()

        num_epochs = 2
        for epoch in range(num_epochs):
            print(f"Epoch {epoch + 1}/{num_epochs}")

            epoch_start = time.time()
            batch_count = 0
            sample_count = 0

            for batch in loader:
                images = []

                for sample in batch:
                    img = sample["image"]

                    # Apply SIMD transforms
                    img = transforms["resize"].apply(img)
                    img = transforms["flip"].apply(img)
                    img = transforms["color_jitter"].apply(img)
                    img = transforms["normalize"].apply(img)

                    # Convert to PyTorch tensor (CHW format)
                    img = to_tensor.apply(img)
                    images.append(torch.from_numpy(img))

                # Stack into batch tensor
                batch_tensor = torch.stack(images)

                # Simulate training step
                # model(batch_tensor)
                # loss.backward()
                # optimizer.step()

                batch_count += 1
                sample_count += len(batch)

                # Limit for demo
                if sample_count >= 200:
                    break

            epoch_time = time.time() - epoch_start
            epoch_throughput = sample_count / epoch_time

            print(
                f"  Processed {sample_count} images in {epoch_time:.2f}s ({epoch_throughput:.1f} img/s)"
            )

        print()
        print("✅ PyTorch integration complete")
        print("   All v1.1.0 features work seamlessly with PyTorch")
        print("=" * 80)
        print()

    except ImportError:
        print("PyTorch not installed. Skipping PyTorch integration demo.")
        print("Install with: pip install torch")
        print("=" * 80)
        print()


def main():
    """
    Complete v1.1.0 workflow demonstration

    This example demonstrates all three major v1.1.0 features:
    1. TBL Binary Format (12.4% smaller, O(1) access)
    2. AVX-512 SIMD Acceleration (2x vector throughput)
    3. Prefetching Pipeline (overlapped I/O)

    Usage:
        python complete_v110_workflow.py /path/to/dataset.tar
    """

    # Check features
    if not check_v110_features():
        sys.exit(1)

    # Configuration
    if len(sys.argv) < 2:
        print("Usage: python complete_v110_workflow.py <tar_file>")
        print()
        print("Example:")
        print("  python complete_v110_workflow.py /tmp/imagenet.tar")
        sys.exit(1)

    tar_path = sys.argv[1]
    tbl_path = tar_path.replace(".tar", ".tbl")

    if not os.path.exists(tar_path):
        print(f"ERROR: Dataset not found: {tar_path}")
        sys.exit(1)

    print("=" * 80)
    print("COMPLETE v1.1.0 WORKFLOW")
    print("=" * 80)
    print()
    print("This demo will:")
    print("  1. Convert TAR → TBL format (12.4% smaller)")
    print("  2. Create SIMD-accelerated pipeline (AVX-512/AVX2/NEON)")
    print("  3. Benchmark complete v1.1.0 pipeline")
    print("  4. Demonstrate PyTorch integration")
    print()
    print("=" * 80)
    print()

    # Step 1: Convert to TBL
    tbl_path = step1_convert_to_tbl(tar_path, tbl_path)
    if not tbl_path:
        print("Conversion failed. Exiting.")
        sys.exit(1)

    # Step 2: Create SIMD pipeline
    transforms = step2_create_simd_pipeline()

    # Step 3: Benchmark full pipeline
    throughput = step3_benchmark_full_pipeline(tbl_path, transforms)

    # Step 4: PyTorch integration
    step4_pytorch_integration(tbl_path, transforms)

    # Summary
    print()
    print("=" * 80)
    print("v1.1.0 WORKFLOW COMPLETE")
    print("=" * 80)
    print()
    print(f"Achieved throughput: {throughput:.1f} img/s")
    print()
    print("v1.1.0 features demonstrated:")
    print("  ✅ TBL Binary Format")
    print("     • 12.4% smaller file size")
    print("     • O(1) random access")
    print("     • Memory-mapped I/O")
    print()
    print("  ✅ AVX-512 SIMD Acceleration")
    print("     • 16-wide vector operations (or AVX2/NEON fallback)")
    print("     • 2x transform throughput")
    print("     • All transforms SIMD-optimized")
    print()
    print("  ✅ Prefetching Pipeline")
    print("     • Overlapped I/O with computation")
    print("     • Reduced epoch time")
    print("     • Automatic in DataLoader")
    print()
    print("Your optimized dataset: " + tbl_path)
    print()
    print("Ready for production ML training!")
    print("=" * 80)


if __name__ == "__main__":
    main()
