from clr import AddReference
AddReference("System")
AddReference("QuantConnect.Algorithm")
AddReference("QuantConnect.Algorithm.Framework")
AddReference("QuantConnect.Common")

import pandas as pd 
import numpy as np

from sklearn.decomposition import PCA
from sklearn.preprocessing import StandardScaler
import talib as ta
from sklearn.svm import SVC
from sklearn import metrics 

from QuantConnect.Indicators import *
from QuantConnect.Algorithm import *
from QuantConnect import *
from QuantConnect.Algorithm.Framework import *
from QuantConnect.Algorithm.Framework.Alphas import *
from QuantConnect.Algorithm.Framework.Execution import *


class myAlgo(QCAlgorithm):

    def Initialize(self):
        self.SetStartDate(2016, 1, 1)  # Set Start Date
        self.SetEndDate(2020, 1, 1)
        self.SetCash(100000)  # Set Strategy Cash
        
        ## Helper variables for universe selection and checking for conditions to update only at the
        ## beginning of each month
        self.num_coarse = 60
        self.curmonth = self.Time.month
        
        # spy = self.AddEquity("SPY", Resolution.Daily)
        # aapl = self.AddEquity("AAPL", Resolution.Daily)
        # msft = self.AddEquity("MSFT", Resolution.Daily)
        # amzn = self.AddEquity("AMZN", Resolution.Daily)
        # googl = self.AddEquity("GOOGL", Resolution.Daily)
        # self.symbols = [spy.Symbol, aapl.Symbol, msft.Symbol, amzn.Symbol, googl.Symbol]
        
        self.AddUniverse(self.CoarseSelectionFunction)
        self.spy = self.AddEquity("SPY", Resolution.Daily).Symbol
        
        # self.SetAlpha(ConstantAlphaModel(InsightType.Price, InsightDirection.Up, TimeSpan.FromMinutes(20), 0.025, None))
        # self.SetExecution(ImmediateExecutionModel())
        
        self.Schedule.On(self.DateRules.Every(DayOfWeek.Monday), self.TimeRules.AfterMarketOpen(self.spy, 60), Action(self.trade))
        # self.Schedule.On(self.DateRules.EveryDay(), self.TimeRules.AfterMarketOpen(self.spy, 5), Action(self.trade))
    
        self.historical_data = 276
        # self.feature_window = 250
        self.long_pos = 5
        self.short_pos = 5
        self.short_period = 20
        self.long_period = 200
        
        self.symbols = []

    def CoarseSelectionFunction(self, coarse):

        sym = sorted([x for x in coarse if x.HasFundamentalData and x.Price>20], 
                                key = lambda x: x.DollarVolume, reverse = True)
        
        syms = sym[:self.num_coarse]
        
        for i in syms:
            if not self.Securities.ContainsKey(i.Symbol):
                # j = self.AddEquity(i.Symbol, Resolution.Daily)
                self.symbols.append(i.Symbol)
            
        return self.symbols
     
    def OnData(self, data):
        '''OnData event is the primary entry point for your algorithm. Each new data point will be pumped in here.
            Arguments:
                data: Slice object keyed by symbol containing the stock data
        '''
        pass
        
        
        
    def trade(self):
        # prices = self.History(self.symbols, self.historical_data, Resolution.Daily)
        longs = []
        shorts = []
        thismonth = self.Time.month
        this = 'this month: ' + str(thismonth)
        self.Debug(this)
        
        if self.curmonth != thismonth:
            self.Liquidate()
            cur = 'cur_month: ' + str(self.curmonth)
            self.Debug(cur)
            self.curmonth = thismonth
            
        for symbol in self.symbols:
            security = self.Securities[symbol]
            if security.IsTradable:
                # price = prices.loc[symbol.Value]
                price = self.History(symbol, self.historical_data, Resolution.Daily)
                # s_sma = self.History(symbol, self.short_period, Resolution.Daily)
                # l_sma = self.History(symbol, self.long_period, Resolution.Daily)
                # price_hist = self.History(symbol, 10, Resolution.Daily)
                if (not price.empty): #and (not s_sma.empty) and (not l_sma.empty) and (not price_hist.empty):
                    price_list = list(price['close'])
                    
                    short_sma = price['close'][-self.short_period:].mean()
                    long_sma = price['close'][-self.long_period:].mean()
                    
                    close_10_days = price_list[-10:-9]
                    current_close = price_list[-1:]
                    
                    # price_list = list(prices.loc[symbol.Value]['close'])
                    
                    price_df = pd.DataFrame(price_list)
                    
                    price_df = price_df.rename(columns = {0:'close'})
                    price_df['sma_5'] = price_df['close'].rolling(window = 5).mean()
                    price_df['rsi'] = self.rsi(price_df, 14)
                    
                    price_df['ema_12'] = self.ema(price_df['close'], 12)
                    price_df['ema_26'] = self.ema(price_df['close'], 26)
                    price_df['macd'] = price_df['ema_12'] - price_df['ema_26']
                    price_df['macd_signal_line'] = self.ema(price_df['macd'], 9) #macd signal line
                    price_df['macd_crossover'] = price_df['macd'] - price_df['macd_signal_line'] #diff between macd and its signal line
                    
                    price_df.dropna(inplace=True)
                    
                    if len(price_df) != 250:
                        continue
                    
                    price_df['label'] = price_df.apply(self.label, axis=1)
                    
                    X = price_df[['rsi', 'macd', 'macd_crossover']]
                    y = price_df['label']
                    
                    #train for the first 249 days
                    X_train = X[:-1]
                    y_train = y[:-1]
                    
                    #predict the 250th day
                    X_test = X[-1:]
                    y_test = y[-1:]
                    
                    ss = StandardScaler()
                    
                    #usig principal component analysis machine learning to ensure no overfitting of variables
                    pca = PCA(n_components = 2)
                    
                    X_train = ss.fit_transform(X_train)
                    X_test = ss.transform(X_test)
                    
                    X_train = pca.fit_transform(X_train)
                    X_test = pca.transform(X_test)
                    
                    #fitting the trainset to support vector classifier machine learning
                    clf = SVC(kernel = 'rbf', C=10, max_iter = 5, random_state = 42)
                    clf.fit(X_train, y_train)
                    
                    #predict the 250th day (whether to buy or sell on the last day)
                    y_pred = clf.predict(X_test)
                    
                    train_r2 = clf.score(X_train, y_train)
                    test_r2 = metrics.accuracy_score(y_test, y_pred)
                    
                    pred = "prediction: " + str(y_pred)
                    train_str = "train r^2: " + str(train_r2)
                    test_str = "test r^2: " + str(test_r2)
                    
                    self.Debug(pred)
                    self.Debug(train_str)
                    self.Debug(test_str)
                    
                    #momentum used to rank the stocks
                    #so that the top 5 will be will enter long while the bottom 5 will enter short
                    momentum = current_close[0] - close_10_days[0]
                    
                    if y_pred == 1 and short_sma>long_sma:
                        longs.append((symbol, momentum))
                        i = "long: " + str(symbol) + " " + str(momentum)
                        self.Debug(i)
                        
                    elif y_pred == -1 and short_sma<long_sma:
                        shorts.append((symbol, momentum))
                        i = "short: " + str(symbol) + " " + str(momentum)
                        self.Debug(i)
            
        long_list = self.list_(longs, 'l')
        short_list = self.list_(shorts, 's')
          
        for i in self.Portfolio.Values:
            # check if stock is in portfolio
            # exit/liquidate if it is currently in holding but not found in our long/short list
            if i.Invested:
                self.Debug(i.Price)
                if i.IsLong:
                    if (i.Symbol not in long_list):
                            self.Liquidate(i.Symbol)
                            ii = "liquidate: " + str(i.Symbol)
                            self.Debug(ii)
                
                elif i.IsShort:
                    if (i.Symbol not in short_list): 
                        self.Liquidate(i.Symbol)
                        ii = "liquidate: " + str(i.Symbol)
                        self.Debug(ii)
        
        for l in long_list:
            if self.Portfolio[l].Quantity == 0:
            #buy if stock is not in holding but found in our long list
                self.SetHoldings(l,1.0/(len(long_list)+len(short_list)))
                ll = "buy: " + str(l)
                self.Debug(ll)
                
        for s in short_list:  
            if self.Portfolio[s].Quantity == 0:
            #sell if stock is not in holding but found in our short list
                self.SetHoldings(s,-1.0/(len(long_list)+len(short_list)))
                ss = "sell: " + str(s)
                self.Debug(ss)
    
    def ema(self, price, window):
        return pd.Series.ewm(price, span=window).mean()[window-1:]
        
    def rsi(self, df, window):
        i = 1
        pos_period = []
        neg_period = []
        close = list(df['close'])
        while i <= (len(df)-1):
            diff = close[i] - close[i-1]
            if diff>0:
                pos_period.append(diff)
                neg_period.append(0)
            else:
                pos_period.append(0)
                neg_period.append(abs(diff))
            i = i+1
            
        pos_period = pd.Series(pos_period)
        neg_period = pd.Series(neg_period)
        
        ema_u = self.ema(pos_period, window)
        ema_d = self.ema(neg_period, window)
        
        rsi = 100 - (100/(1+(ema_u/ema_d)))
        
        return rsi
            
    def label(self, df):
        if df['close'] > df['sma_5']:
            return 1 #buy
        else:
            return -1 #sell
          
            
    def list_(self, lists, long_short):
        if lists == []:
            return
        dtypes = [('symbol', object), ('momentum', float)]
        list_ = np.array(lists, dtype = dtypes)
        list_ = (np.sort(list_, order = 'momentum'))[::-1]
        if long_short == 'l':
            return list_[:self.long_pos]['symbol']
        else:
            return list_[-self.short_pos:]['symbol']