# region imports
from AlgorithmImports import *
# endregion

# Filter:
# Average daily dollar volume greater than $50 million over the last twenty days.
# Minimum price $5.00. I don't like trading stocks that have a price below $5
# because I believe they can be too volatile, but I do believe I have an
# edge trading stocks below $10, because the institutionals usually are not involved.

# Setup:
# Close of the SPY is above the 100-day simple moving average (SMA).
# This indicates a trend in the overall index.
# The close of the 25-day simple moving average crosses above the close of
# the 50-day simple moving average.

# Ranking:
# In case we have more setups than our position sizing allows,
# we rank by the highest rate of change over the last 200 trading days.
# This means the highest percentage price increase over the last 200 trading days.

# Entry:
# Next day market order on open. I'm not worried about slippage for a long trade,
# and I definitely want to be in.

# Stop-Loss:
# The day after entering the trade, place a stop-loss below the execution price
# of five times the average true range (ATR) of the last twenty days. That will definitely
# keep me outside the daily noise and give the trade room to develop.

# Reentry:
# If stopped out, renter the next day if all entry conditions apply again.

# Profit Protection:
# A trailing stop of 25 percent. This is in addition to the initial stop-loss.
# Eventually, as the stock rises, the trailing stop will move up above the stop-loss price.

# Profit taking:
# No profit target; the goal is to ride this as high as it will go.

# Position sizing:
# 2 percent risk and 10 percent maximum percentage size, with a maximum of ten positions.


class LongTrendHighMomentum(QCAlgorithm):

    def Initialize(self):
        # self.Debug("Start Init algo")

        self.SetStartDate(1999, 1, 2)  # Set Start Date
        # self.SetStartDate(2023, 1, 2)  # Set Start Date
        self.SetEndDate(2002, 7, 10)  # Set End Date
        self.SetCash(100000)  # Set Strategy Cash

        self.UniverseSettings.Resolution = Resolution.Daily
        self.universe = self.AddUniverse(self.CoarseSelectionFunction, self.FineSelectionFunction)
        self.spy = self.AddEquity("SPY", Resolution.Daily)
        self.spySma = self.SMA("SPY", 100, Resolution.Daily)
        self.WarmUpIndicator("SPY", self.spySma)
        self.SetBenchmark("SPY")
        self.securities = []
        self.coarse_averages = {}
        self.market_order_tickets = []

        self.AddRiskManagement(TrailingStopRiskManagementModel(0.25))
        # self.Schedule.On(self.DateRules.EveryDay("SPY"),
        #                  self.TimeRules.AfterMarketOpen("SPY", 0.0001, False),
        #                  self.DailyCheck)

        # self.Debug("Finish Init algo")


    def OnData(self, slice):
        if self.spy.Close < self.spySma.Current.Value:
            return

        self.Debug(f"OnData Fired at: {self.Time}")

        todays_setups = []
        for security in self.securities:
            twenty_five_sma = security.twentyFiveSma.Current.Value
            fifty_sma = security.fiftySma.Current.Value
            if twenty_five_sma > fifty_sma:
                todays_setups.append(security)

        todays_setups.sort(
            key=lambda s: s.twoHundredRoc.Current.Value, reverse=True)

        # filter securities in our portfolio we hold only if x.Invested
        current_open_positions = [x for x in self.Portfolio.Values if x.Invested]
        current_open_positions_symbols = [x.Symbol.Value for x in current_open_positions]

        num_setups_to_buy = 10 - len(current_open_positions)
        if (num_setups_to_buy) < 1:
            return

        # log number of setups
        self.Log(f"Number of setups: {num_setups_to_buy}")

        setups_passed = todays_setups[:num_setups_to_buy]
        setups_to_buy = []
        for setup_passed in setups_passed:
            symbol_value = setup_passed.Symbol.Value
            if symbol_value not in current_open_positions_symbols:
                setups_to_buy.append(setup_passed)

        for setup in setups_to_buy:
            current_price_of_setup = setup.Price
            stop_price = current_price_of_setup - \
                (setup.twentyAtr.Current.Value * 5)

            risk_per_share = current_price_of_setup - stop_price

            # Calculate portfolio risk per trade
            portfolio_risk_per_trade = self.Portfolio.TotalPortfolioValue * 0.02

            # Calculate shares to buy, round down to nearest integer
            total_shares = portfolio_risk_per_trade / risk_per_share // 1
            # Check if total_shares exceeds 10% of portfolio, if so, adjust total_shares
            if total_shares * current_price_of_setup > self.Portfolio.TotalPortfolioValue * 0.1:
                total_shares = (self.Portfolio.TotalPortfolioValue * 0.1) // current_price_of_setup

            if total_shares > 0:
                ticket = self.MarketOrder(setup.Symbol, total_shares)


    def OnSecuritiesChanged(self, changes: SecurityChanges) -> None:
        # self.Log("OnSecuritiesChanged:")
        # self.Debug(f"Start OnSecuritiesChanged")

        for security in changes.AddedSecurities:
            # self.Debug(f"Adding {security.Symbol.Value} to securities")

            # Don't add spy to tradeable securities
            if (security != self.spy):

                # check if security exists in averages

                security.twentyFiveSma = self.SMA(security.Symbol, 25)
                security.fiftySma = self.SMA(security.Symbol, 50)
                security.twoHundredRoc = self.ROC(security.Symbol, 200)
                security.twentyAtr = self.ATR(security.Symbol, 20)

                self.WarmUpIndicator(
                    security.Symbol, security.twentyFiveSma)
                self.WarmUpIndicator(
                    security.Symbol, security.fiftySma)
                self.WarmUpIndicator(
                    security.Symbol, security.twoHundredRoc)
                self.WarmUpIndicator(
                    security.Symbol, security.twentyAtr)

                self.securities.append(security)

        for security in changes.RemovedSecurities:
            if security in self.securities:
                self.DeregisterIndicator(security.twentyFiveSma)
                self.DeregisterIndicator(security.fiftySma)
                self.DeregisterIndicator(security.twoHundredRoc)
                self.DeregisterIndicator(security.twentyAtr)
                self.securities.remove(security)

        self.securities = [s for s in self.securities if not s.IsDelisted]


        # self.Debug(f"End OnSecuritiesChanged")


    def OnOrderEvent(self, orderEvent: OrderEvent) -> None:
        order = self.Transactions.GetOrderById(orderEvent.OrderId)

        if orderEvent.Status == OrderStatus.Invalid:
            x = 1

        if orderEvent.Status == OrderStatus.Filled:
            # self.pending_tickets = [x for x in self.pending_tickets if x.Symbol.Value != order.Symbol.Value]
            # Stop market order is 5*atr below execution price
            holdings_quantity = self.Portfolio[order.Symbol].AbsoluteQuantity

            self.Log(f"portfolio size = ")

            if holdings_quantity > 0:
                average_price = self.Portfolio[order.Symbol].AveragePrice
                # self.Debug(f"OrderStatus.Filled: {order.Symbol.Value} holdings_quantity {holdings_quantity} average_price {average_price}")

                twenty_day_atr = self.ATR(order.Symbol, 20)
                self.WarmUpIndicator(order.Symbol, twenty_day_atr)

                twenty_day_atr_value = twenty_day_atr.Current.Value
                stop_price = average_price - 5 * twenty_day_atr_value

                self.StopMarketOrder(
                    order.Symbol, -holdings_quantity, stop_price)

            

    def CoarseSelectionFunction(self, coarse):
        selected = []
        # some initial screening filter
        universe = [x for x in coarse if x.Price > 5]
        # self.Debug("start coarse selection function")

        for coarse in universe:
            symbol = coarse.Symbol

            if symbol not in self.coarse_averages:
                history = self.History(symbol, 20, Resolution.Daily)
                self.coarse_averages[symbol] = SelectionData(history)

            self.coarse_averages[symbol].update(self.Time, coarse.DollarVolume)

            dollarVolumeAvg = self.coarse_averages[symbol].dollarVolumeAvg.Current.Value
            if self.coarse_averages[symbol].is_ready() and dollarVolumeAvg >= 50_000_000 and coarse.HasFundamentalData:
                selected.append(symbol)

        # self.Debug("end coarse selection function")
        return selected

    def FineSelectionFunction(self, fine):
        # self.Debug(f"Start FineSelectionFunction")

        filtered = []
        for security in fine:            
            if security.SecurityReference.ExchangeId == "NAS" or security.SecurityReference.ExchangeId == "NYS" or security.SecurityReference.ExchangeId == "ASE":
                    filtered.append(security.Symbol)
        # self.Debug(f"End FineSelectionFunction")

        return filtered

class SelectionData():
    # 3. Update the constructor to accept a history array
    def __init__(self, history):
        self.dollarVolumeAvg = SimpleMovingAverage(20)
        # 4. Loop over the history data and update the indicators
        for bar in history.itertuples():
            volume = getattr(bar, 'volume', 0)
            close = getattr(bar, 'close', 0)
            dollar_volume = volume * close if volume and close else 0
            self.dollarVolumeAvg.Update(bar.Index[1], dollar_volume)

    def is_ready(self):
        return self.dollarVolumeAvg.IsReady

    def update(self, time, value):
        self.dollarVolumeAvg.Update(time, value)