import numpy as np
import pandas as pd

class SymbolData:
    """
    This class stores data unique to each security in the universe.
    """

    def __init__(self, security, algorithm, num_days_per_sample=4, num_samples=100):
        """
        Input:
         - security
            Security object for the security
         - algorithm
            The algorithm instance running the backtest
         - num_days_per_sample
            The number of open-close intraday returns for each sample
         - num_samples
            The number of samples to train the model
        """
        self.exchange = security.Exchange
        self.symbol = security.Symbol
        self.algorithm = algorithm
        self.num_days_per_sample = num_days_per_sample
        self.num_samples = num_samples 
        self.previous_open = 0
        self.model = None
        
        # Setup consolidators
        self.consolidator = TradeBarConsolidator(timedelta(days=1))
        self.consolidator.DataConsolidated += self.CustomDailyHandler
        algorithm.SubscriptionManager.AddConsolidator(self.symbol, self.consolidator)
        
        # Warm up training set
        self.roc_window = np.array([])
        self.labels_by_day = pd.Series()
        
        data = {f'{self.symbol.ID}_(t-{i})' : [] for i in range(1, num_days_per_sample + 1)}
        self.features_by_day = pd.DataFrame(data)
        
        lookback = num_days_per_sample + num_samples + 1 
        history = algorithm.History(self.symbol, lookback, Resolution.Daily)
        if history.empty or 'close' not in history:
            algorithm.Log(f"Not enough history for {self.symbol} yet")    
            return
        
        history = history.loc[self.symbol]
        history['open_close_return'] = (history.close - history.open) / history.open
        
        start = history.shift(-1).open
        end = history.shift(-2).open
        history['future_return'] = (end - start) / start
        
        for day, row in history.iterrows():
            self.previous_open = row.open
            if self.update_features(day, row.open_close_return) and not pd.isnull(row.future_return):
                row = pd.Series([np.sign(row.future_return)], index=[day])
                self.labels_by_day = self.labels_by_day.append(row)[-self.num_samples:]
    
    
    def update_features(self, day, open_close_return):
        """
        Updates the training data features.
        
        Inputs
         - day
            Timestamp of when we're aware of the open_close_return
         - open_close_return
            Open to close intraday return
            
        Returns T/F, showing if the features are in place to start updating the training labels.
        """
        self.roc_window = np.append(open_close_return, self.roc_window)[:self.num_days_per_sample]
        
        if len(self.roc_window) < self.num_days_per_sample: 
            return False
            
        self.features_by_day.loc[day] = self.roc_window
        self.features_by_day = self.features_by_day[-(self.num_samples+2):]
        return True
        
        
    def CustomDailyHandler(self, sender, consolidated):
        """
        Updates the rolling lookback of training data.
        
        Inputs
         - sender
            Function calling the consolidator
         - consolidated
            Tradebar representing the latest completed trading day
        """
        time = consolidated.EndTime
        if time in self.features_by_day.index:
            return
        
        _open = consolidated.Open
        close = consolidated.Close
        
        open_close_return = (close - _open) / _open
        if self.update_features(time, open_close_return) and self.previous_open:
            day = self.features_by_day.index[-3]
            open_open_return = (_open - self.previous_open) / self.previous_open
            self.labels_by_day[day] = np.sign(open_open_return)
            self.labels_by_day = self.labels_by_day[-self.num_samples:]
            
        self.previous_open = _open


    def dispose(self):
        """
        Removes the consolidator subscription.
        """
        self.algorithm.SubscriptionManager.RemoveConsolidator(self.symbol, self.consolidator)
        
    
    @property
    def IsReady(self):
        return self.features_by_day.shape[0] == self.num_samples + 2

from Selection.FundamentalUniverseSelectionModel import FundamentalUniverseSelectionModel

class BigTechUniverseSelectionModel(FundamentalUniverseSelectionModel):
    """
    This universe selection model contain the 10 largest securities in the technology sector.
    """
    
    def __init__(self, fine_size=10):
        """
        Input:
         - fine_size
            Maximum number of securities in the universe
        """
        self.fine_size = fine_size
        self.month = -1
        super().__init__(True)

    def SelectCoarse(self, algorithm, coarse):
        """
        Coarse universe selection is called each day at midnight.
        
        Input:
         - algorithm
            Algorithm instance running the backtest
         - coarse
            List of CoarseFundamental objects
            
        Returns the symbols that have fundamental data.
        """
        if algorithm.Time.month == self.month:
            return Universe.Unchanged
        return [ x.Symbol for x in coarse if x.HasFundamentalData ]
    
        
    def SelectFine(self, algorithm, fine):
        """
        Fine universe selection is performed each day at midnight after `SelectCoarse`.
        
        Input:
         - algorithm
            Algorithm instance running the backtest
         - fine
            List of FineFundamental objects that result from `SelectCoarse` processing
        
        Returns a list of symbols that are in the energy sector and have the largest market caps.
        """
        self.month = algorithm.Time.month
        
        tech_stocks = [ f for f in fine if f.AssetClassification.MorningstarSectorCode == MorningstarSectorCode.Technology ]
        sorted_by_market_cap = sorted(tech_stocks, key=lambda x: x.MarketCap, reverse=True)
        return [ x.Symbol for x in sorted_by_market_cap[:self.fine_size] ]

import pandas as pd
from sklearn.naive_bayes import GaussianNB
from dateutil.relativedelta import relativedelta

from symbol_data import SymbolData

class GaussianNaiveBayesAlphaModel(AlphaModel):
    """
    Emits insights in the direction of the prediction made by the SymbolData objects.
    """
    symbol_data_by_symbol = {}
    new_securities = False

    def Update(self, algorithm, data):
        """
        Called each time the alpha model receives a new data slice.
        
        Input:
         - algorithm
            Algorithm instance running the backtest
         - data
            A data structure for all of an algorithm's data at a single time step
        
        Returns a list of Insights to the portfolio construction model.
        """
        if self.new_securities:
            self.train()
            self.new_securities = False
        
        tradable_symbols = {}
        features = [[]]
        
        for symbol, symbol_data in self.symbol_data_by_symbol.items():
            if data.ContainsKey(symbol) and data[symbol] is not None and symbol_data.IsReady:
                tradable_symbols[symbol] = symbol_data
                features[0].extend(symbol_data.features_by_day.iloc[-1].values)

        insights = []
        if len(tradable_symbols) == 0:
            return []
            
        weight = 1 / len(tradable_symbols)
        for symbol, symbol_data in tradable_symbols.items():
            direction = symbol_data.model.predict(features)
            if direction:
                insights.append(Insight.Price(symbol, data.Time + timedelta(days=1, seconds=-1), 
                                              direction, None, None, None, weight))

        return insights
        
        
    def OnSecuritiesChanged(self, algorithm, changes):
        """
        Called each time the universe has changed.
        
        Input:
         - algorithm
            Algorithm instance running the backtest
         - changes
            The additions and removals of the algorithm's security subscriptions
        """
        for security in changes.AddedSecurities:
            self.symbol_data_by_symbol[security.Symbol] = SymbolData(security, algorithm)
            
            
        for security in changes.RemovedSecurities:
            symbol_data = self.symbol_data_by_symbol.pop(security.Symbol, None)
            if symbol_data:
                symbol_data.dispose()
        
        self.new_securities = True
    
    
    def train(self):
        """
        Trains the Gaussian Naive Bayes classifier model.
        """
        features = pd.DataFrame()
        labels_by_symbol = {}
        
        for symbol, symbol_data in self.symbol_data_by_symbol.items():
            if symbol_data.IsReady:
                features = pd.concat([features, symbol_data.features_by_day], axis=1)
                labels_by_symbol[symbol] = symbol_data.labels_by_day
        
        for symbol, symbol_data in self.symbol_data_by_symbol.items():
            if symbol_data.IsReady:
                symbol_data.model = GaussianNB().fit(features.iloc[:-2], labels_by_symbol[symbol])

from universe import BigTechUniverseSelectionModel
from alpha import GaussianNaiveBayesAlphaModel

class GaussianNaiveBayesClassificationAlgorithm(QCAlgorithm):

    def Initialize(self):
        self.SetStartDate(2015, 10, 1)
        self.SetEndDate(2020, 10, 13)
        self.SetCash(1000000)
        
        self.SetUniverseSelection(BigTechUniverseSelectionModel())
        self.UniverseSettings.Resolution = Resolution.Daily
        
        self.SetAlpha(GaussianNaiveBayesAlphaModel())
        
        self.SetPortfolioConstruction(InsightWeightingPortfolioConstructionModel())
        
        self.SetExecution(ImmediateExecutionModel())
        
        self.SetBrokerageModel(AlphaStreamsBrokerageModel())