#!/usr/bin/env python3
"""
PyTorch DistributedDataParallel (DDP) Example with TurboLoader

This example demonstrates multi-GPU training with PyTorch DDP and TurboLoader.

Features:
- Automatic data sharding across GPUs
- Synchronized batch normalization
- Gradient synchronization
- Model checkpointing on rank 0
- Per-GPU metrics logging

Requirements:
    pip install torch torchvision turboloader

Usage:
    # Single machine, multiple GPUs
    python distributed_ddp.py --data-path /path/to/data.tar --gpus 4

    # Multi-node (requires MASTER_ADDR and MASTER_PORT environment variables)
    python distributed_ddp.py --data-path /path/to/data.tar --gpus 8 --nodes 2 --node-rank 0
"""

import argparse
import os
import time

import torch
import torch.nn as nn
import torch.distributed as dist
import torch.multiprocessing as mp
from torch.nn.parallel import DistributedDataParallel as DDP

try:
    from torchvision.models import resnet50
except ImportError:
    print("Error: torchvision not installed")
    print("Install with: pip install torchvision")
    exit(1)

import turboloader


def setup_dist(rank, world_size):
    """Initialize distributed training environment."""
    os.environ['MASTER_ADDR'] = os.environ.get('MASTER_ADDR', 'localhost')
    os.environ['MASTER_PORT'] = os.environ.get('MASTER_PORT', '12355')

    # Initialize process group
    dist.init_process_group(
        backend='nccl',
        init_method='env://',
        world_size=world_size,
        rank=rank
    )

    # Set device
    torch.cuda.set_device(rank)


def cleanup_dist():
    """Clean up distributed training."""
    dist.destroy_process_group()


def train_worker(rank, world_size, args):
    """Training worker function for each GPU."""
    print(f"[Rank {rank}] Starting training worker")

    # Setup distributed training
    setup_dist(rank, world_size)

    # Set device
    device = torch.device(f'cuda:{rank}')

    # Create model
    model = resnet50(num_classes=args.num_classes)
    model = model.to(device)

    # Wrap with DDP
    model = DDP(
        model,
        device_ids=[rank],
        output_device=rank,
        find_unused_parameters=False
    )

    # Loss and optimizer
    criterion = nn.CrossEntropyLoss()
    optimizer = torch.optim.SGD(
        model.parameters(),
        lr=args.lr * world_size,  # Scale learning rate with world size
        momentum=0.9,
        weight_decay=1e-4
    )

    # Learning rate scheduler
    scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(
        optimizer,
        T_max=args.epochs
    )

    # Create TurboLoader with distributed sharding
    loader = turboloader.DataLoader(
        args.data_path,
        batch_size=args.batch_size,  # Per-GPU batch size
        num_workers=args.num_workers,
        shuffle=True,
        enable_distributed=True,  # Enable automatic sharding
        drop_last=True  # Ensure equal batches across GPUs
    )

    # Create transforms
    transforms = turboloader.Compose([
        turboloader.Resize(256, 256),
        turboloader.RandomCrop(224, 224),
        turboloader.RandomHorizontalFlip(0.5),
        turboloader.ColorJitter(0.2, 0.2, 0.2, 0.1),
        turboloader.ImageNetNormalize(),
        turboloader.ToTensor()
    ])

    # Training loop
    for epoch in range(args.epochs):
        model.train()

        epoch_start = time.time()
        running_loss = 0.0
        correct = 0
        total = 0
        num_batches = 0

        for batch_idx, batch in enumerate(loader):
            # Process batch
            images = []
            labels = []

            for sample in batch:
                img = transforms.apply(sample['image'])
                images.append(torch.from_numpy(img).float())
                labels.append(sample.get('label', 0))

            # Stack into batch tensors
            images = torch.stack(images).to(device, non_blocking=True)
            labels = torch.tensor(labels, dtype=torch.long).to(device, non_blocking=True)

            # Forward pass
            optimizer.zero_grad()
            outputs = model(images)
            loss = criterion(outputs, labels)

            # Backward pass
            loss.backward()
            optimizer.step()

            # Statistics
            running_loss += loss.item()
            _, predicted = outputs.max(1)
            total += labels.size(0)
            correct += predicted.eq(labels).sum().item()
            num_batches += 1

            # Log progress (only rank 0)
            if rank == 0 and batch_idx % args.log_interval == 0:
                avg_loss = running_loss / (batch_idx + 1)
                acc = 100. * correct / total
                print(f'[Rank {rank}] Epoch {epoch} [{batch_idx}/{num_batches}] '
                      f'Loss: {avg_loss:.3f} Acc: {acc:.2f}%')

        # End of epoch
        epoch_time = time.time() - epoch_start

        # Synchronize metrics across all ranks
        avg_loss = torch.tensor(running_loss / num_batches).to(device)
        acc = torch.tensor(100. * correct / total).to(device)

        dist.all_reduce(avg_loss, op=dist.ReduceOp.AVG)
        dist.all_reduce(acc, op=dist.ReduceOp.AVG)

        # Log epoch summary (only rank 0)
        if rank == 0:
            samples_per_sec = total * world_size / epoch_time
            print(f'\n{"="*70}')
            print(f'Epoch {epoch} Summary:')
            print(f'  Loss: {avg_loss.item():.4f}')
            print(f'  Accuracy: {acc.item():.2f}%')
            print(f'  Time: {epoch_time:.2f}s')
            print(f'  Throughput: {samples_per_sec:.1f} samples/sec')
            print(f'  Learning Rate: {optimizer.param_groups[0]["lr"]:.6f}')
            print(f'{"="*70}\n')

        # Step scheduler
        scheduler.step()

        # Save checkpoint (only rank 0)
        if rank == 0 and (epoch + 1) % args.save_interval == 0:
            checkpoint = {
                'epoch': epoch,
                'model_state_dict': model.module.state_dict(),
                'optimizer_state_dict': optimizer.state_dict(),
                'scheduler_state_dict': scheduler.state_dict(),
                'loss': avg_loss.item(),
                'accuracy': acc.item(),
            }
            checkpoint_path = f'{args.checkpoint_dir}/checkpoint_epoch_{epoch}.pt'
            os.makedirs(args.checkpoint_dir, exist_ok=True)
            torch.save(checkpoint, checkpoint_path)
            print(f'[Rank {rank}] Saved checkpoint: {checkpoint_path}')

        # Synchronize before next epoch
        dist.barrier()

    # Final cleanup
    if rank == 0:
        print(f'\n{"="*70}')
        print('Training completed!')
        print(f'Final accuracy: {acc.item():.2f}%')
        print(f'{"="*70}')

    cleanup_dist()


def main():
    """Main function."""
    parser = argparse.ArgumentParser(
        description='PyTorch DDP Training with TurboLoader'
    )
    parser.add_argument(
        '--data-path',
        type=str,
        required=True,
        help='Path to training data (TAR or TBL format)'
    )
    parser.add_argument(
        '--batch-size',
        type=int,
        default=128,
        help='Per-GPU batch size (default: 128)'
    )
    parser.add_argument(
        '--num-workers',
        type=int,
        default=4,
        help='Number of data loading workers per GPU (default: 4)'
    )
    parser.add_argument(
        '--epochs',
        type=int,
        default=90,
        help='Number of training epochs (default: 90)'
    )
    parser.add_argument(
        '--lr',
        type=float,
        default=0.1,
        help='Base learning rate (default: 0.1)'
    )
    parser.add_argument(
        '--num-classes',
        type=int,
        default=1000,
        help='Number of classes (default: 1000)'
    )
    parser.add_argument(
        '--gpus',
        type=int,
        default=None,
        help='Number of GPUs (default: all available)'
    )
    parser.add_argument(
        '--log-interval',
        type=int,
        default=100,
        help='Logging interval in batches (default: 100)'
    )
    parser.add_argument(
        '--save-interval',
        type=int,
        default=10,
        help='Checkpoint save interval in epochs (default: 10)'
    )
    parser.add_argument(
        '--checkpoint-dir',
        type=str,
        default='./checkpoints',
        help='Checkpoint directory (default: ./checkpoints)'
    )
    parser.add_argument(
        '--nodes',
        type=int,
        default=1,
        help='Number of nodes (default: 1)'
    )
    parser.add_argument(
        '--node-rank',
        type=int,
        default=0,
        help='Node rank (default: 0)'
    )

    args = parser.parse_args()

    # Determine world size
    if args.gpus is None:
        args.gpus = torch.cuda.device_count()

    if args.gpus == 0:
        print("Error: No CUDA devices available")
        print("This script requires GPUs for distributed training")
        exit(1)

    world_size = args.gpus * args.nodes

    print("=" * 70)
    print("PyTorch DDP Training with TurboLoader")
    print("=" * 70)
    print(f"Configuration:")
    print(f"  Data path: {args.data_path}")
    print(f"  GPUs per node: {args.gpus}")
    print(f"  Nodes: {args.nodes}")
    print(f"  World size: {world_size}")
    print(f"  Per-GPU batch size: {args.batch_size}")
    print(f"  Global batch size: {args.batch_size * world_size}")
    print(f"  Workers per GPU: {args.num_workers}")
    print(f"  Epochs: {args.epochs}")
    print(f"  Base learning rate: {args.lr}")
    print(f"  Scaled learning rate: {args.lr * world_size}")
    print("=" * 70 + "\n")

    # Spawn training processes
    mp.spawn(
        train_worker,
        args=(world_size, args),
        nprocs=args.gpus,
        join=True
    )


if __name__ == '__main__':
    main()
