from AlgorithmImports import *
import numpy as np
from datetime import timedelta

class MarketCapWeightedSP500Tracker(QCAlgorithm):

    def Initialize(self):
        self.SetStartDate(2012, 1, 1)
        self.SetEndDate(2025, 1, 1)
        self.SetCash(100000)
        self.UniverseSettings.Resolution = Resolution.Daily

        self.spy = self.AddEquity("SPY", Resolution.Daily).Symbol
        self.bil = self.AddEquity("BIL", Resolution.Daily).Symbol

        self.AddUniverse(self.CoarseSelectionFunction, self.FineSelectionFunction)

        self.selected_by_market_cap = []
        self.rebalance_flag = False
        self.spy_200day_window = RollingWindow[float](200)
        self.entry_prices = {}
        self.previous_bil_allocation = 0.0

        self.Schedule.On(self.DateRules.MonthStart(self.spy), 
                        self.TimeRules.AfterMarketOpen(self.spy, 30), 
                        self.SetRebalanceFlag)
        self.Schedule.On(self.DateRules.WeekStart(self.spy, DayOfWeek.Wednesday), 
                        self.TimeRules.AfterMarketOpen(self.spy, 30), 
                        self.MonthlyRebalance)

        # Initialize rolling window with historical data
        history = self.History(self.spy, 200, Resolution.Daily)
        if not history.empty:
            for time, row in history.loc[self.spy].iterrows():
                self.spy_200day_window.Add(row["close"])

    def CoarseSelectionFunction(self, coarse):
        filtered = [x for x in coarse if x.HasFundamentalData 
                   and x.Price > 5 
                   and x.Market == Market.USA]
        return [x.Symbol for x in filtered]

    def FineSelectionFunction(self, fine):
        filtered = [x for x in fine if x.MarketCap > 1e10
                   and x.ValuationRatios.PERatio < 20 
                   and x.ValuationRatios.PERatio > 0
                   and x.SecurityReference.SecurityType == "ST00000001"]

        sorted_by_cap = sorted(filtered, key=lambda x: x.MarketCap, reverse=True)[:30]
        self.selected_by_market_cap = [(x.Symbol, x.MarketCap) for x in sorted_by_cap]
        return [x.Symbol for x in sorted_by_cap]

    def SetRebalanceFlag(self):
        if self.Time.weekday() == 2:  # Wednesday
            self.rebalance_flag = True

    def OnData(self, data):
        # Update price window
        if not data.Bars.ContainsKey(self.spy): return
        self.spy_200day_window.Add(data.Bars[self.spy].Close)
        
        # Track if any stop-loss was triggered
        stop_loss_triggered = False
        
        # Check for stop-loss triggers on all holdings
        for kvp in self.Portfolio:
            symbol = kvp.Key
            holding = kvp.Value
            
            if holding.Invested and symbol != self.bil:
                current_price = self.Securities[symbol].Price
                
                if symbol not in self.entry_prices:
                    self.entry_prices[symbol] = current_price
                
                price_drop = (self.entry_prices[symbol] - current_price) / self.entry_prices[symbol]
                
                if price_drop >= 0.10:
                    stop_loss_triggered = True
                    # Liquidate the position
                    self.Liquidate(symbol)
                    self.Debug(f"Stop-loss triggered for {symbol} at {current_price}")
        
        # If any stop-loss was triggered, invest all available cash in BIL
        if stop_loss_triggered:
            # Calculate total available cash (including unsettled cash from sales)
            available_cash = self.Portfolio.Cash + self.Portfolio.UnsettledCash
            if available_cash > 0:
                # Calculate how much BIL we can buy with available cash
                bil_price = self.Securities[self.bil].Price
                bil_quantity = available_cash / bil_price
                self.MarketOrder(self.bil, bil_quantity)
                self.Debug(f"Invested ${available_cash:0.2f} in BIL after stop-loss")

    def MonthlyRebalance(self):
        if not self.rebalance_flag: return
        self.rebalance_flag = False
        self.entry_prices.clear()  # Reset entry prices at rebalance

        if self.spy_200day_window.Count < 200:
            self.Debug("Waiting for enough SPY history.")
            return

        spy_price = self.Securities[self.spy].Price
        sma_200 = sum(self.spy_200day_window) / 200

        # Calculate new BIL allocation based on SMA
        bil_weight = 0.0
        if spy_price < sma_200:
            bil_weight = min((sma_200 - spy_price) / sma_200, 1.0)

        # Apply 5% minimum reduction rule from previous month
        min_bil_allocation = self.previous_bil_allocation * 0.95
        bil_weight = max(bil_weight, min_bil_allocation)

        equity_weight = 1.0 - bil_weight

        if not self.selected_by_market_cap:
            self.Debug("No stocks selected for investment.")
            return

        total_market_cap = sum([x[1] for x in self.selected_by_market_cap])
        weights = {x[0]: (x[1] / total_market_cap) * equity_weight for x in self.selected_by_market_cap}

        invested = set()
        for symbol, weight in weights.items():
            if weight > 0:
                self.SetHoldings(symbol, weight)
                invested.add(symbol)
                self.entry_prices[symbol] = self.Securities[symbol].Price

        if bil_weight > 0:
            self.SetHoldings(self.bil, bil_weight)
            invested.add(self.bil)
        else:
            self.Liquidate(self.bil)

        # Store current BIL allocation for next month's minimum
        self.previous_bil_allocation = self.Portfolio[self.bil].HoldingsValue / self.Portfolio.TotalPortfolioValue
        self.Debug(f"New BIL allocation: {bil_weight*100:0.2f}% (Minimum was {min_bil_allocation*100:0.2f}%)")

        # Liquidate positions not in current selection
        for kvp in self.Portfolio:
            symbol = kvp.Key
            if kvp.Value.Invested and symbol not in invested and symbol != self.spy:
                self.Liquidate(symbol)