import logging
from datetime import datetime

import matplotlib.pyplot as plt
import pandas as pd
from src.evaluation.backtester import GNNBacktester
from src.evaluation.metrics import calculate_performance_metrics, compare_to_benchmark
from src.utils.visualization import (
    plot_feature_importance,
    plot_performance,
    plot_trade_log,
)

from config import config
from src.amd.optimizations import AMDOptimizer
from src.data.pipeline import StockDataPipeline
from src.models.gnn_model import CorporateActionAwareGNN
from src.models.trainer import GNNTrainer
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("stock_gnn_r9700.log"),
        logging.StreamHandler(),
    ],
)
logger = logging.getLogger(__name__)


def main():
    # Initialize memory manager
    memory_manager = MemoryManager()
    logger.info(memory_manager.get_memory_stats())

    # Initialize AMD optimizer
    amd_optimizer = AMDOptimizer()

    # Initialize data pipeline
    logger.info("Initializing data pipeline")
    pipeline = StockDataPipeline()

    # Update all data
    logger.info("Updating all data sources")
    pipeline.update_all_data()

    # Create datasets
    logger.info("Creating training and validation datasets")
    train_dataset = pipeline.create_training_dataset()
    val_dataset = pipeline.create_validation_dataset(
        start_date=config.TRAIN_END_DATE, end_date=config.VAL_END_DATE
    )

    logger.info(f"Training dataset size: {len(train_dataset)}")
    logger.info(f"Validation dataset size: {len(val_dataset)}")

    # Initialize model
    logger.info("Initializing GNN model")
    # x has shape (num_stocks, seq_len, num_features); features are in the last dim
    num_features = train_dataset[0].x.shape[2]
    model = CorporateActionAwareGNN(num_features)

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

    # Train model
    logger.info("Training model with AMD optimizations")
    trainer = GNNTrainer(model)
    train_losses, val_losses = trainer.train(train_dataset, val_dataset)

    # Plot training curves
    plt.figure(figsize=(10, 5))
    plt.plot(train_losses, label="Training Loss")
    plt.plot(val_losses, label="Validation Loss")
    plt.title("Training and Validation Loss")
    plt.xlabel("Epoch")
    plt.ylabel("Loss")
    plt.legend()
    plt.savefig("training_curves.png")
    plt.close()

    # Benchmark model
    logger.info("Benchmarking model performance")
    sample_data = train_dataset[0].to(config.DEVICE)
    benchmark_results = trainer.benchmark(sample_data)
    logger.info(f"Benchmark Results: {benchmark_results}")

    # Load best model
    trainer.load_model()

    # Run backtest on validation set
    logger.info("Running backtest on validation set")
    backtester = GNNBacktester(model, pipeline)
    portfolio_values, trade_log = backtester.run_backtest(val_dataset)

    # Get benchmark data (auto_adjust=True means 'Close' already contains adjusted prices)
    benchmark_data = pipeline.price_data[config.INDEX_TICKER]
    benchmark_values = benchmark_data.loc[portfolio_values.index]["Close"]

    # Calculate performance metrics
    logger.info("Calculating performance metrics")
    portfolio_returns = portfolio_values.pct_change().dropna()
    benchmark_returns = benchmark_values.pct_change().dropna()

    metrics = calculate_performance_metrics(portfolio_returns, benchmark_returns)
    comparison = compare_to_benchmark(portfolio_values, benchmark_values)

    # Print metrics
    logger.info("\nPerformance Metrics:")
    for metric, value in metrics.items():
        if isinstance(value, float):
            logger.info(f"{metric.replace('_', ' ').title()}: {value:.4f}")
        else:
            logger.info(f"{metric.replace('_', ' ').title()}: {value}")

    logger.info("\nComparison to Benchmark:")
    for metric, value in comparison.items():
        if isinstance(value, float):
            logger.info(f"{metric.replace('_', ' ').title()}: {value:.4f}")
        else:
            logger.info(f"{metric.replace('_', ' ').title()}: {value}")

    # Plot performance
    plot_performance(portfolio_values, benchmark_values, "portfolio_performance.png")

    # Plot trade log
    plot_trade_log(trade_log, "trade_log.png")

    # Save results
    results = {
        "portfolio_values": portfolio_values,
        "benchmark_values": benchmark_values,
        "trade_log": trade_log,
        "metrics": metrics,
        "comparison": comparison,
    }

    results_df = pd.DataFrame(
        {
            "date": portfolio_values.index,
            "portfolio_value": portfolio_values.values,
            "benchmark_value": benchmark_values.values,
        }
    )
    results_df.to_csv("backtest_results.csv", index=False)

    # Memory cleanup
    memory_manager.empty_cache()
    logger.info("Backtest completed. Results saved to backtest_results.csv")
    logger.info(memory_manager.get_memory_stats())


if __name__ == "__main__":
    main()