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

class M_StrengthAlgorithm(QCAlgorithm):

    def Initialize(self):
        self.SetUniverseSelection(ManualUniverseSelectionModel(["SPY"]))
        self.SetPortfolioConstruction(EqualWeightingPortfolioConstructionModel())
        self.SetExecution(ImmediateExecutionModel())
        # self.add_alpha(RsiAlphaModel())
        # self.add_risk_management(NullRiskManagementModel())

        ## included stop-loss at 4% ##

        self.SetStartDate(2019, 1, 1)
        self.SetEndDate(2024, 9, 1)
        self.SetCash(100000)
        self.spy = self.AddEquity("SPY", Resolution.Minute).Symbol
        self.lookback = timedelta(minutes=5)
        self.position_count = 0
        self.positions = []
        self.halt_trades = True
        self.initial_price = None
        self.stop_loss_percent = 0.04
        self.Settings.MinimumOrderMarginPortfolioValue = 0
        self.SetWarmUp(self.lookback)

        self.Schedule.On(
            self.DateRules.EveryDay(),
            self.TimeRules.AfterMarketOpen(self.spy, 1),
            self.InitialInvestment
        )

        self.Schedule.On(
            self.DateRules.EveryDay(),
            self.TimeRules.At(10, 0),
            self.CheckInitialSlope
        )

        self.Schedule.On(
            self.DateRules.EveryDay(),
            self.TimeRules.Every(timedelta(minutes=5)),
            self.EvaluateMarket
        )

        self.Schedule.On(
            self.DateRules.EveryDay(),
            self.TimeRules.At(12, 15),
            self.SellAllPositions
        )

    def InitialInvestment(self):
        if self.IsWarmingUp:
            return
        if not self.Portfolio[self.spy].Invested:
            self.SetHoldings(self.spy, 0.4)

    def CheckInitialSlope(self):
        history = self.History(self.spy, self.lookback, Resolution.Minute)
        if history.empty or self.spy not in history.index or 'close' not in history.columns:
            return
        closes = history.loc[self.spy]["close"].values
        if len(closes) < self.lookback.total_seconds() / 60:
            return
        x = np.arange(len(closes))
        m, _ = np.polyfit(x, closes, 1)
        if abs(m) <= np.tan(np.radians(30)):
            self.halt_trades = False
            self.initial_price = closes[-1]
        else:
            self.halt_trades = True

    def EvaluateMarket(self):
        if self.IsWarmingUp or self.halt_trades:
            return
        history = self.History(self.spy, self.lookback, Resolution.Minute)
        if history.empty or self.spy not in history.index or 'close' not in history.columns:
            return
        closes = history.loc[self.spy]["close"].values
        if len(closes) < self.lookback.total_seconds() / 60:
            return
        current_price = closes[-1]
        allocation = (self.Portfolio.TotalPortfolioValue - self.Portfolio[self.spy].HoldingsValue) / current_price
        quantity = int(allocation)
        if quantity > 0 and self.CanPlaceOrder(quantity, current_price):
            if not self.Portfolio[self.spy].Invested:
                self.MarketOrder(self.spy, quantity)
                self.positions.append(current_price)
                self.position_count += 1
                self.SetStopLoss(self.spy, self.stop_loss_percent)
            elif len(self.positions) > 0 and self.position_count < 3:
                if current_price > self.positions[-1]:
                    self.MarketOrder(self.spy, quantity)
                    self.positions.append(current_price)
                    self.position_count += 1
                    self.SetStopLoss(self.spy, self.stop_loss_percent)
                elif current_price < self.initial_price:
                    self.halt_trades = True
            elif len(self.positions) == 3:
                if current_price > self.positions[-1]:
                    self.MarketOrder(self.spy, quantity)
                    sell_quantity = self.Portfolio[self.spy].Quantity / len(self.positions)
                    if sell_quantity > 0:
                        self.MarketOrder(self.spy, -sell_quantity)
                    self.positions.pop(0)
                    self.positions.append(current_price)
                    self.SetStopLoss(self.spy, self.stop_loss_percent)
                elif current_price < self.initial_price:
                    self.halt_trades = True

    def CanPlaceOrder(self, quantity, price):
        cost = quantity * price
        margin_remaining = self.Portfolio.GetMarginRemaining(self.spy, OrderDirection.Buy)
        if self.Portfolio.Cash >= cost and margin_remaining >= cost:
            return True
        return False

    def SetStopLoss(self, symbol, stop_loss_percent):
        quantity = self.Portfolio[symbol].Quantity
        if quantity > 0:
            stop_price = self.Portfolio[symbol].AveragePrice * (1 - stop_loss_percent)
            self.StopMarketOrder(symbol, -quantity, stop_price)

    def SellAllPositions(self):
        if self.Portfolio.Invested:
            history = self.History(self.spy, self.lookback, Resolution.Minute)
            if history.empty or self.spy not in history.index or 'close' not in history.columns:
                self.Liquidate(self.spy)
                self.positions = []
                self.position_count = 0
                return
            closes = history.loc[self.spy]["close"].values
            if len(closes) < 2:
                self.Liquidate(self.spy)
                self.positions = []
                self.position_count = 0
                return
            m, _ = np.polyfit(np.arange(len(closes)), closes, 1)
            if m > 0:
                return
            self.Liquidate(self.spy)
            self.positions = []
            self.position_count = 0