import logging
from datetime import datetime
from typing import Dict, List, Tuple

import pandas as pd
import torch

from config import config
from src.models.gnn_model import CorporateActionAwareGNN

logger = logging.getLogger(__name__)


class GNNBacktester:
    def __init__(
        self,
        model: CorporateActionAwareGNN,
        price_data: Dict,
        initial_capital: float = config.INITIAL_CAPITAL,
    ):
        self.model = model
        self.price_data = price_data
        self.initial_capital = initial_capital
        self.portfolio_value = initial_capital
        self.portfolio = {}  # {ticker: shares}
        self.trade_log = []
        self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

    def run_backtest(self, dataset: List) -> Tuple[pd.Series, pd.DataFrame]:
        """
        Run backtest on the given dataset

        Parameters:
        dataset: List of PyG Data objects

        Returns:
        Tuple of (portfolio_values, trade_log)
        """
        portfolio_values = []
        dates = []

        for data in dataset:
            date = data.date
            current_tickers = data.tickers

            # Get current prices
            current_prices = {}
            for ticker in current_tickers:
                if ticker in self.price_data and date in self.price_data[ticker].index:
                    current_prices[ticker] = self.price_data[ticker].loc[date][
                        "Adj Close"
                    ]

            # Calculate current portfolio value
            current_value = sum(
                shares * current_prices[ticker]
                for ticker, shares in self.portfolio.items()
                if ticker in current_prices
            )
            cash = self.portfolio_value - current_value
            current_portfolio_value = current_value + cash

            # Store portfolio value
            portfolio_values.append(current_portfolio_value)
            dates.append(date)

            # Get predictions from GNN
            with torch.no_grad():
                data = data.to(self.device)
                predictions = self.model(data).squeeze().cpu().numpy()

            # Create trading signals
            signals = {}
            for i, ticker in enumerate(current_tickers):
                if ticker in current_prices:
                    # Get corporate action flag (last feature)
                    corporate_action_flag = data.x[i, -1].item()

                    # Buy if predicted return > threshold and no upcoming corporate action
                    if predictions[i] > 0.005 and corporate_action_flag == 0:
                        signals[ticker] = "buy"
                    # Sell if predicted return < threshold or upcoming corporate action
                    elif predictions[i] < -0.005 or corporate_action_flag > 0:
                        signals[ticker] = "sell"

            # Execute trades
            for ticker, signal in signals.items():
                if signal == "buy" and cash > 0:
                    # Buy with 10% of available cash
                    price = current_prices[ticker]
                    shares_to_buy = int(
                        (cash * 0.1) / (price * (1 + config.TRANSACTION_COST))
                    )
                    if shares_to_buy > 0:
                        cost = shares_to_buy * price * (1 + config.TRANSACTION_COST)
                        self.portfolio[ticker] = (
                            self.portfolio.get(ticker, 0) + shares_to_buy
                        )
                        cash -= cost
                        self.trade_log.append(
                            (date, ticker, "buy", shares_to_buy, price)
                        )
                        logger.debug(
                            f"Bought {shares_to_buy} shares of {ticker} at {price:.2f}"
                        )

                elif signal == "sell" and ticker in self.portfolio:
                    # Sell all shares
                    shares = self.portfolio.pop(ticker)
                    proceeds = (
                        shares * current_prices[ticker] * (1 - config.TRANSACTION_COST)
                    )
                    cash += proceeds
                    self.trade_log.append(
                        (date, ticker, "sell", shares, current_prices[ticker])
                    )
                    logger.debug(
                        f"Sold {shares} shares of {ticker} at {current_prices[ticker]:.2f}"
                    )

            # Update portfolio value
            new_value = sum(
                shares * current_prices[ticker]
                for ticker, shares in self.portfolio.items()
                if ticker in current_prices
            )
            self.portfolio_value = new_value + cash

        # Create portfolio value series
        portfolio_series = pd.Series(portfolio_values, index=dates)

        # Create trade log DataFrame
        if self.trade_log:
            trade_log_df = pd.DataFrame(self.trade_log)
            trade_log_df.columns = ["date", "ticker", "action", "shares", "price"]
        else:
            trade_log_df = pd.DataFrame()
            trade_log_df.columns = ["date", "ticker", "action", "shares", "price"]

        return portfolio_series, trade_log_df

    def get_portfolio_composition(self, date: datetime) -> Dict[str, float]:
        """
        Get portfolio composition at a specific date

        Parameters:
        date: Date to get composition for

        Returns:
        Dictionary of {ticker: weight} where weight is the percentage of portfolio
        """
        # Get current prices
        current_prices = {}
        for ticker in self.portfolio:
            if ticker in self.price_data and date in self.price_data[ticker].index:
                current_prices[ticker] = self.price_data[ticker].loc[date]["Adj Close"]

        # Calculate current value
        current_value = sum(
            shares * current_prices[ticker]
            for ticker, shares in self.portfolio.items()
            if ticker in current_prices
        )
        cash = self.portfolio_value - current_value

        # Calculate weights
        composition = {}
        for ticker, shares in self.portfolio.items():
            if ticker in current_prices:
                composition[ticker] = (
                    shares * current_prices[ticker]
                ) / self.portfolio_value

        # Add cash
        composition["cash"] = cash / self.portfolio_value

        return composition