# region imports
from AlgorithmImports import *

# endregion



class CryptoMomentum(QCAlgorithm):


    def Initialize(self):
        self.SetStartDate(2020, 1, 1)
        self.SetEndDate(2021, 1, 1)

        self.SetCash(100000)

        self.SetBrokerageModel(BrokerageName.GDAX, AccountType.Cash)

        self.UniverseSettings.Resolution = Resolution.Hour
        self.Settings.FreePortfolioValuePercentage = 0.05


        self.AddUniverse(CryptoCoarseFundamentalUniverse(Market.GDAX, self.UniverseSettings, self.universe_filter))

        
        self.symbol_data_by_symbol = {}

        #symbol over which other crypto prices will be divided by; either BTC or ETH usually
        self.base_ticker = "BTCUSD"
        
        self.base_symbol = self.AddCrypto(self.base_ticker, Resolution.Hour).Symbol

        self.fast_ma = self.EMA(self.base_symbol, 24*15)
        self.slow_ma = self.EMA(self.base_symbol, 24*50)

        self.SetWarmUp(24*20, Resolution.Hour)



        self.SetBenchmark(self.base_symbol)



    def universe_filter(self, crypto_coarse: List[CryptoCoarseFundamental]) -> List[Symbol]:
        
        restricted_tickers = ['DAIUSD', 'DAIUSDC', 'USDC']
        universe = []
        #universe = [cf.Symbol for cf in crypto_coarse if "ETHUSD" ==  cf.Symbol.Value]

        for cf in crypto_coarse:
            if cf.Symbol.Value not in restricted_tickers and cf.Volume > 100000:
                if "USD" in cf.Symbol.Value and "USDC" not in cf.Symbol.Value:
                    universe.append(cf.Symbol)


        return universe




    def OnSecuritiesChanged(self, changes: SecurityChanges) -> None:

        for security in changes.AddedSecurities:
            self.symbol_data_by_symbol[security.Symbol] = SymbolData(self, security.Symbol, self.base_symbol)

        # If you have a dynamic universe, track removed securities
        for security in changes.RemovedSecurities:
            self.Liquidate(security.Symbol)

            symbol_data = self.symbol_data_by_symbol.pop(security.Symbol, None)
            if symbol_data:
                symbol_data.dispose()


    def OnData(self, data: Slice):


        #daily trigger
        if data.Time.hour != 12: return


        buy_symbols = []
        sell_symbols = []


        sorted_symbol_data = sorted(self.symbol_data_by_symbol.values(),key=lambda x: x.rate_of_change, reverse=True)


        #top 5 momentum

        for symbol_data in sorted_symbol_data[:5]:
            
            symbol = symbol_data.symbol

            if self.base_ticker in symbol.Value: continue

            if symbol_data.sma_fast.Current.Value > symbol_data.sma_slow.Current.Value and symbol_data.sma_slow.IsReady:
                buy_symbols.append(symbol)

            #self.Plot("moving averages", symbol_data.sma_fast, symbol_data.sma_slow)
        
        for symbol in self.symbol_data_by_symbol.keys():
            if symbol not in buy_symbols:
                self.Liquidate(symbol)


        if len(buy_symbols) > 0:

            '''
            refine logic so minor buys/sells aren't triggered (i.e. within +/- few percentage points of target)
            '''

            port_value = self.Portfolio.TotalPortfolioValue
            target_weight = (1/len(buy_symbols))*.95
            target_value = target_weight * port_value

            if target_weight > 0.2: target_weight = 0.2

            portfolio_targets = []

            for symbol in buy_symbols:
                portfolio_targets.append(PortfolioTarget(symbol, target_weight))

            self.SetHoldings(portfolio_targets)


        else:
            self.Liquidate()



class SymbolData:
    def __init__(self, algorithm, symbol, base_symbol):

        self.algorithm = algorithm
        self.symbol = symbol
        self.base_symbol = base_symbol


        self.sma_fast = ExponentialMovingAverage(24*1)
        self.sma_slow = ExponentialMovingAverage(24*10)


        self.rate_of_change = 0

        # Create a consolidator to update the indicator
        self.consolidator = TradeBarConsolidator(1) 

        self.consolidator.DataConsolidated += self.OnDataConsolidated

        # Register the consolidator to update the indicator
        self.algorithm.SubscriptionManager.AddConsolidator(symbol, self.consolidator)



    def OnDataConsolidated(self, sender: object, consolidated_bar: TradeBar) -> None:

        base_close = self.algorithm.Securities[self.base_symbol].Close

        #calculate moving averages relative to the base symbol
        self.sma_fast.Update(consolidated_bar.EndTime, consolidated_bar.Close/base_close)
        self.sma_slow.Update(consolidated_bar.EndTime, consolidated_bar.Close/base_close)

        self.rate_of_change= self.sma_fast.Current.Value/self.sma_slow.Current.Value

    # If you have a dynamic universe, remove consolidators for the securities removed from the universe
    def dispose(self) -> None:
        self.algorithm.SubscriptionManager.RemoveConsolidator(self.symbol, self.consolidator)