import logging
import time

import numpy as np
import torch
import torch.nn as nn
from torch_geometric.data import Data

from config import config
from src.amd.optimizations import AMDOptimizer
from src.models.intraday_gnn import IntradayGNN
from src.utils.memory_manager import MemoryManager

# Configure logging
logging.basicConfig(
    level=logging.INFO,
    format="%(asctime)s - %(name)s - %(levelname)s - %(message)s",
    handlers=[
        logging.FileHandler("benchmark_r9700.log"),
        logging.StreamHandler(),
    ],
)
logger = logging.getLogger(__name__)


def benchmark_model():
    """Benchmark the GNN model on AMD Radeon R9700 AI Pro"""
    # Initialize memory manager
    memory_manager = MemoryManager()
    logger.info(memory_manager.get_memory_stats())

    # Initialize AMD optimizer
    amd_optimizer = AMDOptimizer()

    # Create a sample model
    num_features = len(config.INTRADAY_FEATURES) + 5  # +5 for price features
    model = IntradayGNN(num_features, config.SEQUENCE_LENGTH)

    # Optimize model for AMD GPU
    model = amd_optimizer.optimize_model(model)

    # Create sample data
    batch_size = config.BATCH_SIZE
    sequence_length = config.SEQUENCE_LENGTH
    num_stocks = 50  # Number of stocks in the graph

    # Create random data
    num_edges = 200
    x = torch.randn(num_stocks, sequence_length, num_features).to(config.DEVICE)
    edge_index = torch.randint(0, num_stocks, (2, num_edges)).to(config.DEVICE)
    edge_attr = torch.randn(num_edges, 1).to(config.DEVICE)
    y = torch.randn(num_stocks, 1).to(config.DEVICE)
    sample_data = Data(x=x, edge_index=edge_index, edge_attr=edge_attr, y=y)

    # Warm-up
    logger.info("Warming up...")
    for _ in range(10):
        with torch.no_grad():
            _ = model(sample_data)

    # Benchmark inference
    logger.info("Benchmarking inference...")
    start_time = time.time()
    num_runs = 100

    for _ in range(num_runs):
        with torch.no_grad():
            _ = model(sample_data)

    inference_time = (time.time() - start_time) / num_runs
    logger.info(f"Average inference time: {inference_time:.6f} seconds")

    # Benchmark training
    logger.info("Benchmarking training...")
    model.train()
    optimizer = torch.optim.Adam(model.parameters(), lr=config.LEARNING_RATE)
    criterion = nn.MSELoss()

    start_time = time.time()

    for _ in range(num_runs):
        optimizer.zero_grad()
        out = model(sample_data)
        loss = criterion(out, y)
        loss.backward()
        optimizer.step()

    training_time = (time.time() - start_time) / num_runs
    logger.info(f"Average training time: {training_time:.6f} seconds")

    # Memory usage
    memory_info = memory_manager.check_memory()
    logger.info(
        f"GPU Memory Usage: {memory_info['allocated'] / 1024**3:.2f}GB / {memory_info['total'] / 1024**3:.2f}GB"
    )

    # Throughput
    logger.info(f"Inference throughput: {1 / inference_time:.2f} samples/second")
    logger.info(f"Training throughput: {1 / training_time:.2f} samples/second")

    # Detailed benchmark with different batch sizes
    logger.info("\nDetailed benchmark with different configurations:")

    batch_sizes = [32, 64, 128, 256]
    sequence_lengths = [30, 60, 120]

    for batch_size in batch_sizes:
        for seq_len in sequence_lengths:
            # Create data for this configuration
            bx = torch.randn(num_stocks, seq_len, num_features).to(config.DEVICE)
            bei = torch.randint(0, num_stocks, (2, num_edges)).to(config.DEVICE)
            bea = torch.randn(num_edges, 1).to(config.DEVICE)
            by = torch.randn(num_stocks, 1).to(config.DEVICE)
            bdata = Data(x=bx, edge_index=bei, edge_attr=bea, y=by)

            # Benchmark inference
            start_time = time.time()
            for _ in range(10):  # Fewer runs for detailed benchmark
                with torch.no_grad():
                    _ = model(bdata)
            inf_time = (time.time() - start_time) / 10

            # Benchmark training
            start_time = time.time()
            for _ in range(10):
                optimizer.zero_grad()
                out = model(bdata)
                loss = criterion(out, by)
                loss.backward()
                optimizer.step()
            train_time = (time.time() - start_time) / 10

            logger.info(
                f"Batch: {batch_size}, Seq Len: {seq_len}, "
                f"Inf Time: {inf_time:.6f}s, Train Time: {train_time:.6f}s, "
                f"Inf Tput: {1 / inf_time:.2f} samples/s, Train Tput: {1 / train_time:.2f} samples/s"
            )

    # Memory benchmark
    logger.info("\nMemory benchmark:")

    # Test different model sizes
    hidden_channels_list = [64, 128, 256, 512]

    for hidden_channels in hidden_channels_list:
        # Create a model with this configuration
        model = IntradayGNN(num_features, config.SEQUENCE_LENGTH)
        model.feature_processor = nn.Sequential(
            nn.Linear(num_features, hidden_channels),
            nn.SiLU(),
            nn.Linear(hidden_channels, hidden_channels),
            nn.LayerNorm(hidden_channels),
        )
        model.linear = nn.Linear(hidden_channels, 1)

        # Optimize model
        model = amd_optimizer.optimize_model(model)

        # Estimate memory usage
        estimated_memory = memory_manager.estimate_model_memory(model)
        logger.info(
            f"Hidden Channels: {hidden_channels}, Estimated Memory: {estimated_memory / 1024**3:.2f}GB"
        )

        # Clean up
        del model
        memory_manager.empty_cache()

    # Final memory stats
    logger.info("\nFinal Memory Stats:")
    logger.info(memory_manager.get_memory_stats())


if __name__ == "__main__":
    benchmark_model()