# region imports
from AlgorithmImports import *
# endregion

from Selection.FundamentalUniverseSelectionModel import FundamentalUniverseSelectionModel

class RevisedBuyOnDip(QCAlgorithm):
    def Initialize(self):
        self.SetStartDate(2000, 1, 1)
        self.SetEndDate(2025, 1, 1)
        self.SetCash(100_000)
        
        self.spy = self.AddEquity("SPY", Resolution.DAILY).Symbol
        
        # Dictionaries to track lots and total allocation per symbol:
        # self.lots: key = symbol, value = list of tuples (lot_allocation, entry_price)
        self.lots = {}
        # self.totalAllocated: key = symbol, value = cumulative allocation percentage
        self.totalAllocated = {}
        
        self.UniverseSettings.Resolution = Resolution.DAILY
        self._universe = BiggestMarketCapUniverse(self)
        self.SetUniverseSelection(self._universe)
        
        # Schedule our daily rebalancing call
        self.Schedule.On(self.DateRules.EveryDay(self.spy), 
                         self.TimeRules.AfterMarketOpen(self.spy, 1), 
                         self.Rebalance)

    def Rebalance(self):
        # 1. Check for profit taking on any held positions
        for symbol in list(self.lots.keys()):
            # Ensure we have a valid price
            if symbol not in self.Securities or self.Securities[symbol].Price == 0:
                continue
            currentPrice = self.Securities[symbol].Price
            lotsToKeep = []
            totalReduction = 0.0
            # Check each lot for a 5% gain over its entry price
            for lotAllocation, entryPrice in self.lots[symbol]:
                if (currentPrice - entryPrice) / entryPrice >= 0.05:
                    totalReduction += lotAllocation
                    self.Log(f"Liquidating {lotAllocation*100:.1f}% of {symbol} at {currentPrice} (entry was {entryPrice})")
                else:
                    lotsToKeep.append((lotAllocation, entryPrice))
            # If any lot qualifies, update the holdings for that symbol
            if totalReduction > 0:
                newAllocation = self.totalAllocated[symbol] - totalReduction
                self.totalAllocated[symbol] = newAllocation
                self.SetHoldings(symbol, newAllocation)
                if newAllocation == 0:
                    # Remove symbol entirely if fully liquidated
                    del self.lots[symbol]
                    del self.totalAllocated[symbol]
                else:
                    self.lots[symbol] = lotsToKeep

        # 2. Universe selection and buying on dip
        selectedSymbol = self._universe.last_selected_symbol
        if not selectedSymbol:
            return
        
        # Add the symbol if it is not already in our Securities
        if selectedSymbol not in self.Securities:
            self.AddEquity(selectedSymbol, Resolution.DAILY)
        
        # Retrieve recent price history to calculate the price drop
        history = self.History(selectedSymbol, 2, Resolution.DAILY)
        if history.empty or len(history) < 2:
            return
        
        prevClose = history.iloc[-2]['close']
        currClose = history.iloc[-1]['close']
        priceDrop = (prevClose - currClose) / prevClose

        # Check if price drop is significant (>= 5%) and we have capacity to add more allocation.
        # We assume a maximum of 99% allocation for any given stock.
        currentAlloc = self.totalAllocated.get(selectedSymbol, 0)
        if priceDrop >= 0.05 and currentAlloc + 0.09 <= 0.99:
            # Reduce our SPY allocation from 100% to 90%
            self.SetHoldings(self.spy, 0.9)
            # Increase the allocation for the selected stock by 9%
            newAllocation = currentAlloc + 0.09
            self.SetHoldings(selectedSymbol, newAllocation)
            self.totalAllocated[selectedSymbol] = newAllocation
            
            # Record the new purchase as a lot
            if selectedSymbol not in self.lots:
                self.lots[selectedSymbol] = []
            self.lots[selectedSymbol].append((0.09, currClose))
            self.Log(f"Bought additional 9% of {selectedSymbol} at {currClose}, total allocation now {newAllocation*100:.1f}%")
            
class BiggestMarketCapUniverse(FundamentalUniverseSelectionModel):
    def __init__(self, algorithm):
        super().__init__(True)  # True = use dynamic universe selection
        self.algorithm = algorithm
        self.last_selected_symbol = None

    def SelectCoarse(self, algorithm, coarse):
        filtered = [x for x in coarse if x.HasFundamentalData and x.Market == Market.USA]
        return [x.Symbol for x in filtered]

    def SelectFine(self, algorithm, fine):
        if not fine:
            return []
        fine = sorted(fine, key=lambda x: x.MarketCap, reverse=True)
        self.last_selected_symbol = fine[0].Symbol
        return [self.last_selected_symbol]