# SPDX-License-Identifier: Apache-2.0
import argparse
import os

import torch
import torch.nn as nn

from fastvideo.distributed.parallel_state import (
    cleanup_dist_env_and_memory, destroy_distributed_environment,
    destroy_model_parallel, get_tp_rank,
    get_tp_world_size, maybe_init_distributed_environment_and_model_parallel, get_world_group)
from fastvideo.layers.linear import ColumnParallelLinear, RowParallelLinear
from fastvideo.logger import init_logger

logger = init_logger(__name__)


class SimpleTPModel(nn.Module):
    """A simple model that uses tensor parallelism."""

    def __init__(self, hidden_size=1024, intermediate_size=4096):
        super().__init__()
        # Column parallel linear layer (splits output dimension)
        self.fc1 = ColumnParallelLinear(
            input_size=hidden_size,
            output_size=intermediate_size,
            bias=True,
            gather_output=
            False,  # Don't gather output since we're passing to row parallel
            skip_bias_add=False)

        # Row parallel linear layer (splits input dimension)
        self.fc2 = RowParallelLinear(
            input_size=intermediate_size,
            output_size=hidden_size,
            bias=True,
            input_is_parallel=True,  # Input is already split from previous layer
            skip_bias_add=False)

        self.activation = nn.GELU()

    def forward(self, x):
        # Forward through column parallel layer
        hidden_states, _ = self.fc1(x)

        # Apply activation
        hidden_states = self.activation(hidden_states)

        # Forward through row parallel layer
        output, _ = self.fc2(hidden_states)

        return output


def initialize_random_weights(model, seed=42):
    """Initialize the model with random weights using a fixed seed for reproducibility."""
    # Set seed for reproducibility
    torch.manual_seed(seed)

    # Initialize weights for each layer
    with torch.no_grad():
        # For ColumnParallelLinear layers
        if hasattr(model, 'fc1'):
            nn.init.normal_(model.fc1.weight, mean=0.0, std=0.02)
            if model.fc1.bias is not None:
                nn.init.zeros_(model.fc1.bias)

        # For RowParallelLinear layers
        if hasattr(model, 'fc2'):
            nn.init.normal_(model.fc2.weight, mean=0.0, std=0.02)
            if model.fc2.bias is not None:
                nn.init.zeros_(model.fc2.bias)

    logger.info("Model initialized with random weights")
    return model


def setup_args():
    parser = argparse.ArgumentParser(
        description='Simple Tensor Parallelism Example')
    parser.add_argument('--tensor-model-parallel-size',
                        type=int,
                        default=8,
                        help='Degree of tensor model parallelism')
    parser.add_argument('--batch-size',
                        type=int,
                        default=8,
                        help='Batch size for the example')
    parser.add_argument('--hidden-size',
                        type=int,
                        default=1024,
                        help='Hidden size for the model')
    parser.add_argument('--intermediate-size',
                        type=int,
                        default=4096,
                        help='Intermediate size for the model')
    return parser.parse_args()


def main():
    args = setup_args()

    maybe_init_distributed_environment_and_model_parallel(args.tensor_model_parallel_size, args.tensor_model_parallel_size)
    rank = get_world_group().rank
    local_rank = get_world_group().local_rank
    # Get tensor parallel info
    tp_rank = get_tp_rank()
    tp_world_size = get_tp_world_size()

    logger.info(
        f"Process rank {rank} initialized with TP rank {tp_rank} in TP world size {tp_world_size}"
    )

    # Create a simple model
    model = SimpleTPModel(hidden_size=args.hidden_size,
                          intermediate_size=args.intermediate_size)

    # Initialize with random weights
    model = initialize_random_weights(model)

    # Create a random input tensor
    batch_size = args.batch_size
    hidden_size = args.hidden_size
    x = torch.randn(batch_size, hidden_size, dtype=torch.float)

    # Move to GPU if available
    device = torch.device(
        f"cuda:{local_rank}" if torch.cuda.is_available() else "cpu")
    model = model.to(device)
    x = x.to(device)

    # Forward pass
    logger.info(f"Running forward pass on TP rank {tp_rank}")
    with torch.no_grad():
        output = model(x)

    # Print output shape and statistics
    logger.info(f"Output shape: {output.shape}")
    logger.info(
        f"Output mean: {output.mean().item()}, std: {output.std().item()}")

    # Clean up
    logger.info("Cleaning up distributed environment")
    destroy_model_parallel()
    destroy_distributed_environment()
    cleanup_dist_env_and_memory()

    logger.info("Example completed successfully")


if __name__ == "__main__":
    main()