QCAlgorithmFramework = QCAlgorithm
QCAlgorithmFrameworkBridge = QCAlgorithm

from QuantConnect import *
from QuantConnect.Parameters import *
from QuantConnect.Benchmarks import *
from QuantConnect.Brokerages import *
from QuantConnect.Util import *
from QuantConnect.Interfaces import *
from QuantConnect.Algorithm import *
from QuantConnect.Algorithm.Framework import *
from QuantConnect.Algorithm.Framework.Selection import *
from QuantConnect.Algorithm.Framework.Alphas import *
from QuantConnect.Algorithm.Framework.Portfolio import *
from QuantConnect.Algorithm.Framework.Execution import *
from QuantConnect.Algorithm.Framework.Risk import *
from QuantConnect.Indicators import *
from QuantConnect.Data import *
from QuantConnect.Data.Consolidators import *
from QuantConnect.Data.Custom import *
from QuantConnect.Data.Fundamental import *
from QuantConnect.Data.Market import *
from QuantConnect.Data.UniverseSelection import *
from QuantConnect.Notifications import *
from QuantConnect.Orders import *
from QuantConnect.Orders.Fees import *
from QuantConnect.Orders.Fills import *
from QuantConnect.Orders.Slippage import *
from QuantConnect.Scheduling import *
from QuantConnect.Securities import *
from QuantConnect.Securities.Equity import *
from QuantConnect.Securities.Forex import *
from QuantConnect.Securities.Interfaces import *
from datetime import date, datetime, timedelta
from QuantConnect.Python import *
from QuantConnect.Storage import *

from QuantConnect import Resolution
from QuantConnect.Algorithm import QCAlgorithm


# topics
# 1. use of CoarseSelectionFunction with indicator
# 2. creation of self-defined SelectionData class used in CoarseSelectionFunction
# 3. use of historical data to warm up indicator
# takeaways
# 1. avoid updating the indicator using the same price by specificying a time stamp
# 2. avoid placing invalid orders by checking if the current slice has data for a symbol
class EMAMomentumUniverse(QCAlgorithm):
    def Initialize(self):
        self.SetStartDate(2019, 1, 7)
        self.SetEndDate(2021, 1, 7)
        self.SetCash(100000)
        self.UniverseSettings.Resolution = Resolution.Daily
        self.AddUniverse(self.CoarseSelectionFunction)
        self.averages = {}

    def CoarseSelectionFunction(self, universe):
        selected = []
        universe = sorted(universe, key=lambda c: c.DollarVolume, reverse=True)
        universe = [c for c in universe if c.Price > 10][:100]

        for coarse in universe:
            symbol = coarse.Symbol

            if symbol not in self.averages:
                # 1. Call history to get an array of 200 days of history data
                history = self.History(symbol, 200, Resolution.Daily)

                # 2. Adjust SelectionData to pass in the history result
                self.averages[symbol] = SelectionData(history)
                # self.Debug(str(symbol))
                # self.Debug(history.index.get_level_values(1)[-1])  # last history date
                # self.Debug(f"history price: {history.close[-1]}")  # last history price
                # self.Debug(self.Time)  # algo date
                # self.Debug(f"algo price: {coarse.AdjustedPrice}")  # algo price
            self.averages[symbol].update(self.Time, coarse.AdjustedPrice)

            if (
                self.averages[symbol].is_ready()
                and self.averages[symbol].fast > self.averages[symbol].slow
            ):
                selected.append(symbol)

        return selected[:10]

    def OnSecuritiesChanged(self, changes):
        # self.Debug(self.Time)
        for security in changes.RemovedSecurities:
            self.Liquidate(security.Symbol)

        for security in changes.AddedSecurities:
            # self.Debug(str(security))
            if self.CurrentSlice.ContainsKey(
                security.Symbol
            ):  # check if the current slice has data for this symbol
                self.SetHoldings(security.Symbol, 0.10)


class SelectionData:
    # 3. Update the constructor to accept a history array
    def __init__(self, price_hist):
        self.slow = ExponentialMovingAverage(200)
        self.fast = ExponentialMovingAverage(50)
        # 4. Loop over the history data and update the indicators
        for bar in price_hist.itertuples():
            self.fast.Update(bar.Index[1], bar.close)
            self.slow.Update(bar.Index[1], bar.close)

    def is_ready(self):
        return self.slow.IsReady and self.fast.IsReady

    def update(self, time, price):
        self.fast.Update(time, price)
        self.slow.Update(time, price)