from AlgorithmImports import *
from collections import defaultdict
from datetime import timedelta, time

class VwmaStrategy(QCAlgorithm):

    def Initialize(self):
        # ---------------------------------------------------------
        #   Set basic parameters and framework
        # ---------------------------------------------------------
        self.SetStartDate(2023, 4, 1)
        self.SetEndDate(2024, 12, 30)
        self.SetCash(1000000)
        self.SetSecurityInitializer(lambda x: x.SetDataNormalizationMode(DataNormalizationMode.Raw))

        self.trade_weight       = 0.4      # e.g. 0.8 => 80% of portfolio equity
        self.minimum_threshold  = 0.002
        self.vwma_days          = 27
        self.days_rate_of_change= 6

        self.total_minutes = self.vwma_days * 960  

        # ---------------------------------------------------------
        #   Symbol configuration
        # ---------------------------------------------------------
        self.AddEquity("SPY", Resolution.Daily)
        self.symbol = "QQQ"
        self.AddEquity(self.symbol, Resolution.Minute)

        # ---------------------------------------------------------
        #   Store daily VWMA values
        # ---------------------------------------------------------
        self.vwma_history = []

        # ---------------------------------------------------------
        #   Schedule a daily event at market close
        # ---------------------------------------------------------
        self.Schedule.On(
            self.DateRules.EveryDay("SPY"),
            self.TimeRules.At(16, 0, 0),
            self.DayStart
        )


    def DayStart(self):
        """
        This gets called once per day at 16:00 (market close).
        We calculate the current VWMA, then decide whether to go
        long, short, or stay out based on the (N-day) rate of
        change of the VWMA, where N = days_rate_of_change.
        """
        today_vwma = self.CalculateVWMA(self.symbol, self.total_minutes)
        if today_vwma is None:
            return
        
        self.Debug(f"Date: {self.Time}, VWMA: {today_vwma}")
        
        # Store the latest VWMA in our history
        self.vwma_history.append(today_vwma)

        # Keep the vwma_history from growing too large; 30 is arbitrary
        if len(self.vwma_history) > 30:
            self.vwma_history.pop(0)

        # We need at least N+1 daily VWMA points to compute an N-day rate of change
        if len(self.vwma_history) < self.days_rate_of_change + 1:
            return

        # ---------------------------------------------------------
        #   Compute the N-day rate of change:
        #   roc = (VWMA_current - VWMA_n_days_ago) / VWMA_n_days_ago
        # ---------------------------------------------------------
        past_vwma    = self.vwma_history[-(self.days_rate_of_change + 1)]
        current_vwma = self.vwma_history[-1]

        if abs(past_vwma) < 1e-12:
            return

        roc = (current_vwma - past_vwma) / past_vwma

        # ---------------------------------------------------------
        #   Trading logic (using portfolio-weighted sizing)
        # ---------------------------------------------------------
        current_position = self.Portfolio[self.symbol].Quantity
        price            = self.Securities[self.symbol].Price

        # 1) If ROC > minimum_threshold => Long
        if roc > self.minimum_threshold:
            # How many shares correspond to `trade_weight` * total_portfolio_value
            desired_value_long = self.Portfolio.TotalPortfolioValue * self.trade_weight
            target_shares_long = int(desired_value_long / price)

            # Calculate how many shares we need to buy/sell to get from current_position to target_shares_long
            delta_shares = target_shares_long - current_position

            # If delta_shares > 0, we’re buying; if delta_shares < 0, we’re selling
            if delta_shares != 0:
                self.MarketOrder(self.symbol, delta_shares)
            return

        # 2) If ROC < -minimum_threshold => Short
        if roc < -self.minimum_threshold:
            # How many shares correspond to `trade_weight` * total_portfolio_value
            desired_value_short = self.Portfolio.TotalPortfolioValue * self.trade_weight
            # Convert that to shares, but for a short position we aim for negative holdings
            target_shares_short = int(desired_value_short / price)
            new_target_shares   = -target_shares_short  # Make it negative for short

            delta_shares = new_target_shares - current_position
            if delta_shares != 0:
                self.MarketOrder(self.symbol, delta_shares)
            return

        # 3) If |ROC| <= minimum_threshold => Liquidate
        if current_position != 0:
            self.Liquidate(self.symbol)


    def CalculateVWMA(self, symbol, bars_count):
        """
        Calculate a Volume-Weighted Moving Average (VWMA) over the last
        `bars_count` minute-bars (excluding today's data up to the current time).
        """
        history = self.history[TradeBar](symbol, bars_count, Resolution.Minute, extended_market_hours=True)
        if not history:
            return None

        # Exclude today's bars
        today = self.Time.date()
        history = [bar for bar in history if bar.Time.date() != today]
        if not history:
            return None

        vwma_sum = 0
        volume_sum = 0
        for bar in history:
            vwma_sum += bar.Close * bar.Volume
            volume_sum += bar.Volume

        if volume_sum == 0:
            return None

        return vwma_sum / volume_sum