import logging

import matplotlib.pyplot as plt
import pandas as pd

from config import config
from src.data.pipeline import StockDataPipeline
from src.evaluation.metrics import calculate_performance_metrics, compare_to_benchmark
from src.models.backtester import GNNBacktester
from src.models.gnn_model import CorporateActionAwareGNN
from src.models.trainer import GNNTrainer
from src.utils.visualization import plot_performance, plot_trade_log

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


def main():
    # 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()

    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")
    # Get number of features from first data point
    num_features = train_dataset[0].x.shape[1]
    model = CorporateActionAwareGNN(num_features)

    # Train model
    logger.info("Training model")
    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()

    # Load best model
    trainer.load_model()

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

    # Get benchmark data
    benchmark_data = pipeline.price_data[config.BENCHMARK_TICKER]
    benchmark_values = benchmark_data.loc[portfolio_values.index]["Adj 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_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)

    logger.info("Backtest completed. Results saved to backtest_results.csv")


if __name__ == "__main__":
    main()