# region imports
from AlgorithmImports import *
# endregion

class BuyHoldMoneyMarket(QCAlgorithm):

    def Initialize(self):
        # Phase 1: Initial buy-and-hold strategy
        self.SetStartDate(2022, 12, 12)
        self.SetCash(2000000)
        self.AddEquity("VUSB", Resolution.Minute)
        self.AddEquity("VMFXX", Resolution.Minute)

        # Schedule liquidation and transition to the new strategy
        self.Schedule.On(self.DateRules.Today, self.TimeRules.AfterMarketOpen("VUSB", 1), self.LiquidateAndTransition)

        # Flag to track the strategy phase
        self.transitioned = False

        # Set the benchmark during initialization
        self.SetBenchmark("QQQ")

    def OnData(self, data: Slice):
        if not self.transitioned:
            # Phase 1: Buy and hold logic
            if not self.Portfolio.Invested:
                self.SetHoldings("VUSB", 1)
                self.SetHoldings("VMFXX", 0)
        else:
            # Phase 2: New trading strategy logic
            self.ExecuteNewStrategy(data)

    def LiquidateAndTransition(self):
        # Liquidate positions in VUSB and VMFXX
        self.Liquidate("VUSB")
        self.Liquidate("VMFXX")
        self.Debug("Liquidated VUSB and VMFXX positions")

        # Transition to the new strategy
        self.TransitionToNewStrategy()

    def TransitionToNewStrategy(self):
        # Phase 2: Initialize stocks and indicators for the new strategy
        self.Equities = ['WPM', 'GFI', 'AGI', 'HMY', 'EGO', 'AMLP', 'TSLA', 'LMT', 'GD']
        self.symbol_objs = {}
        
        # Initialize indicators and rolling windows for each symbol
        self.smaLong = {}
        self.smaMid = {}
        self._atr = {}
        self.high50 = {}
        self.low50 = {}

        for symbol in self.Equities:
            symbol_obj = self.AddEquity(symbol, Resolution.Daily).Symbol
            self.symbol_objs[symbol] = symbol_obj

            self.smaLong[symbol] = SimpleMovingAverage(100)
            self.smaMid[symbol] = SimpleMovingAverage(50)
            self._atr[symbol] = AverageTrueRange(symbol, 100, MovingAverageType.EXPONENTIAL)
            self.high50[symbol] = Maximum(50)
            self.low50[symbol] = Minimum(50)
        
        # Schedule daily trades based on the new strategy
        self.Schedule.On(self.DateRules.EveryDay(), self.TimeRules.At(12, 17), Action(self.TradeOnMACDSignal))

        # Mark that the transition to the new strategy is complete
        self.transitioned = True

    def ExecuteNewStrategy(self, data: Slice):
        for symbol in self.Equities:
            if self._atr[symbol].Current.Value != 0:
                price = self.Securities[symbol].Close
                holdings = self.Portfolio[symbol].Quantity

                # Realize profits if UnrealizedProfitPercent > 10%
                if self.Portfolio[symbol].UnrealizedProfitPercent > 0.10:
                    self.Liquidate(symbol)
                    self.Log(f"Sold {symbol} at {self.Securities[symbol].Price} to realize gains.")
                    continue

                if holdings == 0:
                    if self.smaMid[symbol].Current.Value > self.smaLong[symbol].Current.Value:
                        if price >= self.high50[symbol].Current.Value:
                            self.SetHoldings(self.symbol_objs[symbol], 1 / len(self.Equities))
                            self.Log("1BUY " + str(symbol) + " @" + str(price))
                    elif self.smaMid[symbol].Current.Value < self.smaLong[symbol].Current.Value:
                        if price <= self.low50[symbol].Current.Value:
                            self.SetHoldings(self.symbol_objs[symbol], 1 / len(self.Equities))
                            self.Log("2Short " + str(symbol) + " @" + str(price))

                elif holdings > 0:
                    if price <= self.Securities[symbol].Price - 3 * self._atr[symbol].Current.Value:
                        self.Liquidate(symbol)
                        self.Log("close long " + str(symbol) + " @" + str(price))

                elif holdings < 0:
                    if price >= self.Securities[symbol].Price + 3 * self._atr[symbol].Current.Value:
                        self.Liquidate(symbol)
                        self.Log("close short " + str(symbol) + " @" + str(price))

    def TradeOnMACDSignal(self):
        self.Debug("TradeOnMACDSignal")
        self.ExecuteNewStrategy(None)

    def OnData(self, data: Slice):
        if self.transitioned:
            # Update indicators in OnData after transitioning to the new strategy
            for symbol in self.Equities:
                history = self.History[TradeBar](symbol, 120, Resolution.Minute)
                for bar in history:
                    self.smaLong[symbol].Update(bar.EndTime, bar.Close)
                    self.smaMid[symbol].Update(bar.EndTime, bar.Close)
                    self._atr[symbol].Update(bar)
                    self.low50[symbol].Update(bar.EndTime, bar.Close)
                    self.high50[symbol].Update(bar.EndTime, bar.Close)