from Execution.ImmediateExecutionModel import ImmediateExecutionModel
from HealthcareUniverse import HealthcareUniverse
from datetime import timedelta
from enum import Enum

class TachyonModulatedCircuit(QCAlgorithm):

    def Initialize(self):
        # Set Start Date so that backtest has 5+ years of data
        # The blocked section of code below is to remain UNCHANGED for the weekly competitions. 
        # 
        # Insight-weighting portfolio construction model:
        # - You can change the rebalancing date rules or portfolio bias
        # - For more info see https://github.com/QuantConnect/Lean/blob/master/Algorithm.Framework/Portfolio/InsightWeightingPortfolioConstructionModel.py
        # 
        # Use the Alpha Streams Brokerage Model: 
        # - Developed in conjunction with funds to model their actual fees, costs, etc. Please do not modify other models.
        ###############################################################################################################################
        self.SetStartDate(2019, 3, 1)   # 5 years up to the submission date
        self.SetCash(1000000)           # Set $1m Strategy Cash to trade significant AUM
        self.SetBenchmark('SPY')        # SPY Benchmark
        self.SetBrokerageModel(AlphaStreamsBrokerageModel())  
        self.SetExecution(ImmediateExecutionModel()) 
        self.SetPortfolioConstruction(InsightWeightingPortfolioConstructionModel())
        ###############################################################################################################################
        # Do not change the code above 

        # Add the alpha model and anything else you want below
        self.AddAlpha(MyCompetitionAlphaModel())
        
        # Add a universe selection model        
        self.SetUniverseSelection(HealthcareUniverse())

class MyCompetitionAlphaModel(AlphaModel):
    
    def __init__(self,
                 period = 14,
                 resolution = Resolution.Daily):# *args, **kwargs):
        '''Initializes a new default instance of your Alpha Model class.'''
        
        self.period = period
        self.resolution = resolution
        self.insightPeriod = Time.Multiply(Extensions.ToTimeSpan(resolution), period)
        self.symbolDataBySymbol ={}


    def Update(self, algorithm, data):
        '''Updates this alpha model with the latest data from the algorithm.
        This is called each time the algorithm receives data for subscribed securities
        Args:
            algorithm: The algorithm instance
            data: The new data available
        Returns:
            The new insights generated'''
        insights = []

        # This is where insights are returned, which are then passed to the
        # Portfolio Construction, Risk, and Execution models.
        
        # The following Insight properties MUST be set before returning
        #   - Symbol        -- Secuirty Symbol
        #   - Duration      -- Time duration that the Insight is in effect 
        #   - Direction     -- Direction of predicted price movement 
        #   - Weight        -- Proportion of algorithm capital to be allocated to this Insight

        for symbol, symbolData in self.symbolDataBySymbol.items():
            rsi = symbolData.RSI
            previous_state = symbolData.State
            state = self.GetState(rsi, previous_state)
            wt = 1/len(self.symbolDataBySymbol.items())
            if state != previous_state and rsi.IsReady:
                if state == State.TrippedLow:
                    insights.append(Insight.Price(symbol, self.insightPeriod, InsightDirection.Up,wt))
                if state == State.TrippedHigh:
                    insights.append(Insight.Price(symbol, self.insightPeriod, InsightDirection.Down,wt))

            symbolData.State = state
            
        return insights

    def OnSecuritiesChanged(self, algorithm, changes):
        '''Event fired each time the we add/remove securities from the data feed
        Args:
            algorithm: The algorithm instance that experienced the change in securities
            changes: The security additions and removals from the algorithm'''
        symbols = [ x.Symbol for x in changes.RemovedSecurities ]
        if len(symbols) > 0:
            for subscription in algorithm.SubscriptionManager.Subscriptions:
                if subscription.Symbol in symbols:
                    self.symbolDataBySymbol.pop(subscription.Symbol, None)
                    subscription.Consolidators.Clear()

        # initialize data for added securities

        addedSymbols = [ x.Symbol for x in changes.AddedSecurities if x.Symbol not in self.symbolDataBySymbol]
        if len(addedSymbols) == 0: return

        history = algorithm.History(addedSymbols, self.period, self.resolution)

        for symbol in addedSymbols:
            rsi = algorithm.RSI(symbol, self.period, MovingAverageType.Wilders, self.resolution)

            if not history.empty:
                ticker = SymbolCache.GetTicker(symbol)

                if ticker not in history.index.levels[0]:
                    continue

                for tuple in history.loc[ticker].itertuples():
                    rsi.Update(tuple.Index, tuple.close)

            self.symbolDataBySymbol[symbol] = SymbolData(symbol, rsi)
    
    def GetState(self, rsi, previous):
        ''' Determines the new state. This is basically cross-over detection logic that
        includes considerations for bouncing using the configured bounce tolerance.'''
        if rsi.Current.Value > 70:
            return State.TrippedHigh
        if rsi.Current.Value < 30:
            return State.TrippedLow
        if previous == State.TrippedLow:
            if rsi.Current.Value > 35:
                return State.Middle
        if previous == State.TrippedHigh:
            if rsi.Current.Value < 65:
                return State.Middle

        return previous


class SymbolData:
    '''Contains data specific to a symbol required by this model'''
    def __init__(self, symbol, rsi):
        self.Symbol = symbol
        self.RSI = rsi
        self.State = State.Middle


class State(Enum):
    '''Defines the state. This is used to prevent signal spamming and aid in bounce detection.'''
    TrippedLow = 0
    Middle = 1
    TrippedHigh = 2

from Selection.FundamentalUniverseSelectionModel import FundamentalUniverseSelectionModel

class HealthcareUniverse(FundamentalUniverseSelectionModel):
    '''
    This module selects the most liquid stocks listed on the Nasdaq Stock Exchange.
    '''

    def __init__(self, filterFineData = True, universeSettings = None, securityInitializer = None):
        '''Initializes a new default instance of the TechnologyUniverseModule'''
        super().__init__(filterFineData, universeSettings, securityInitializer)
        self.numberOfSymbolsCoarse = 1000
        self.numberOfSymbolsFine = 100
        self.dollarVolumeBySymbol = {}
        self.lastMonth = -1

    def SelectCoarse(self, algorithm, coarse):
        '''
        Performs a coarse selection:
        
        -The stock must have fundamental data
        -The stock must have positive previous-day close price
        -The stock must have positive volume on the previous trading day
        '''
        if algorithm.Time.month == self.lastMonth: 
            return Universe.Unchanged

        sortedByDollarVolume = sorted([x for x in coarse if x.HasFundamentalData and x.Volume > 0 and x.Price > 0],
            key = lambda x: x.DollarVolume, reverse=True)[:self.numberOfSymbolsCoarse]

        self.dollarVolumeBySymbol = {x.Symbol:x.DollarVolume for x in sortedByDollarVolume}
        
        # If no security has met the QC500 criteria, the universe is unchanged.
        if len(self.dollarVolumeBySymbol) == 0:
            return Universe.Unchanged

        return list(self.dollarVolumeBySymbol.keys())

    def SelectFine(self, algorithm, fine):
        '''
        Performs a fine selection for companies in the Morningstar Banking Sector
        '''
        # Filter stocks and sort on dollar volume
        sortedByDollarVolume = sorted([x for x in fine if x.AssetClassification.MorningstarSectorCode == MorningstarSectorCode.Healthcare],
            key = lambda x: self.dollarVolumeBySymbol[x.Symbol], reverse=True)

        if len(sortedByDollarVolume) == 0:
            return Universe.Unchanged
            
        self.lastMonth = algorithm.Time.month

        return [x.Symbol for x in sortedByDollarVolume[:self.numberOfSymbolsFine]]