from QuantConnect.Data import SubscriptionDataSource
from QuantConnect.Python import PythonData
from QuantConnect.Algorithm import QCAlgorithm
from QuantConnect.Data.UniverseSelection import *
import decimal as d
import numpy as np
import pandas as pd
import time
from scipy.stats import linregress, zscore
import talib
from datetime import timedelta, date, datetime
import datetime as dt
from pandas.tseries.offsets import BDay

class TechnicalMultiFactorAlgo(QCAlgorithm):

    def Initialize(self):

        self.SetStartDate(2009, 1, 1)  #Set Start Date
        self.SetEndDate(2018, 1, 1)  #Set Start Date

        self.SetCash(100000)           #Set Strategy Cash

        self.UniverseSettings.Resolution = Resolution.Daily
        self.UniverseSettings.Leverage   = 1

        self.max_per_side_assets     = 4 # so 2*N is total assets at L=2

        self.SetBrokerageModel(BrokerageName.InteractiveBrokersBrokerage, AccountType.Margin)
        #self.AddUniverse(self.CoarseSelectionFunction)
        self.AddUniverse(StockDataSource, "sp500", self.stockDataSource)

        self.AddEquity("SPY", Resolution.Minute)

        if self.LiveMode:
            
            self.Schedule.On(self.DateRules.MonthStart("SPY"),
                self.TimeRules.At(8,05),
                Action(self.Downloader))
            
            self.Schedule.On(self.DateRules.MonthStart("SPY"),
                self.TimeRules.At(8,15),
                Action(self.Downloader))
            
            self.Schedule.On(self.DateRules.MonthStart("SPY"),
                self.TimeRules.At(8,25),
                Action(self.Downloader))
            
            self.Schedule.On(self.DateRules.MonthStart("SPY"),
                self.TimeRules.At(8,35),
                Action(self.Downloader))
                
            self.Schedule.On(self.DateRules.MonthStart("SPY"),
                self.TimeRules.At(8,45),
                Action(self.Downloader))
            
            self.Schedule.On(self.DateRules.MonthStart("SPY"),
                self.TimeRules.At(8,55),
                Action(self.Downloader))
            
            self.Schedule.On(self.DateRules.MonthStart("SPY"),
                self.TimeRules.At(9,5),
                Action(self.Downloader))
            
            self.Schedule.On(self.DateRules.MonthStart("SPY"),
                self.TimeRules.At(9,15),
                Action(self.Downloader))

            
            self.Schedule.On(self.DateRules.MonthStart("SPY"),
                self.TimeRules.At(9,25),
                Action(self.Strategy))
                
            self.Schedule.On(self.DateRules.MonthStart("SPY"),
                self.TimeRules.AfterMarketOpen("SPY", 0),
                Action(self.Rebalance))
        else:
            # Note we force shift into market hours, otherwise these run out of order if set before 9:31
            
            for i in range(2):
                self.Schedule.On(self.DateRules.MonthStart("SPY"),
                    self.TimeRules.At(9,25+i),
                    Action(self.Downloader))
            
            self.Schedule.On(self.DateRules.MonthStart("SPY"),
                self.TimeRules.At(9,32),
                Action(self.Strategy))
                
            self.Schedule.On(self.DateRules.MonthStart("SPY"),
                self.TimeRules.At(9,33),
                Action(self.Rebalance))

        self.Schedule.On(self.DateRules.MonthStart("SPY"), \
            self.TimeRules.BeforeMarketClose("SPY", 0), \
            Action(self.Reset_Baskets))
            
        self.universe              = []
        self.Reset_Baskets()

        self.n_hist_items = 253 * 3
        self.SetWarmUp(self.n_hist_items)
        
        self.splotName = 'Strategy Info'
        sPlot = Chart(self.splotName)
        sPlot.AddSeries(Series('Leverage',  SeriesType.Line, 0))
        self.AddChart(sPlot)
        
        self.last_month_fired_coarse    = None #we cannot rely on Day==1 like before
        self.kama = self.KAMA("SPY", 200, Resolution.Daily)

    def Reset_Baskets(self):
        
        self.current_symbols_long  = []

        self.split_universe        = []
        self.current_subset        = 0
        self.asset_hist            = {}
        self.dfs                   = {}
        
    def Downloader(self):
        
        self.Log("Downloader %d "%self.current_subset + str(self.Time))
        
        if len(self.split_universe) == 0:
            return
        
        self.dfs = {}
        for symbol in self.split_universe[self.current_subset]:
            self.asset_hist[symbol] = self.History([symbol,], self.n_hist_items, Resolution.Daily).astype(np.float32)
            hist = self.asset_hist[symbol].unstack(level=0)
            hist.index = pd.to_datetime(hist.index)
            if len(hist.index) < self.n_hist_items:
                pass
            else:
                self.dfs[symbol] = hist
        self.current_subset += 1
        
        
    def z_scoring(self,date):
        tickers = list(self.dfs.keys())
        all_factors = FactorCollection().allFactors(list(self.dfs.values()),date,'close')[0]
        z_factors = []
        zf_dict = {}
        for factor in all_factors:  # Calculate factor z-scores => normalize for comparison...
            z_factor = (factor - np.mean(factor)) / np.std(factor)
            z_factors.append(z_factor)
        z_factors = np.array(z_factors)
        where_are_nans = np.isnan(z_factors)
        z_factors[where_are_nans] = 0.
        for i in range(len(tickers)):
            zf_dict[tickers[i]] = np.dot(z_factors.T[i], np.array([0.25,-0.5,0.5,-0.5,0.5,0.5]))
        return zf_dict

       
    def Strategy(self):
        
        self.Log("Strategy " + str(self.Time))
        if len(self.universe) == 0:
            return

        date = pd.to_datetime(self.Time)
        z_dict = self.z_scoring(date)
            
        sorted_stock = sorted(z_dict.items(), key=lambda x: x[1],reverse=True)
        sorted_symbol = [sorted_stock[i][0] for i in xrange(len(sorted_stock))]
        self.current_symbols_long = sorted_symbol[:self.max_per_side_assets]

        for symbol in self.universe:
            if symbol not in self.current_symbols_long:
                self.RemoveSecurity(symbol)
        
    
    def Rebalance(self):
        self.Log("Rebalance " + str(self.Time))
        #if self.Securities["SPY"].Price > self.kama.Current.Value:
        if len(self.current_symbols_long) > 0:
            for i in self.Portfolio.Values:
                if (i.Invested) and (i.Symbol not in self.current_symbols_long):
                    self.Liquidate(i.Symbol)

            for sym in self.current_symbols_long:
                self.SetHoldings(sym, 1./float(len(self.current_symbols_long)))
        #else:
            #self.Liquidate()

        self.account_leverage = self.Portfolio.TotalAbsoluteHoldingsCost / self.Portfolio.TotalPortfolioValue
        self.Plot(self.splotName,'Leverage', float(self.account_leverage))
        
    def stockDataSource(self, data):
        list = []
        for item in data:
            for symbol in item["Symbols"]:
                symbolString = Symbol.Create(symbol, SecurityType.Equity, Market.USA)
                list.append(symbolString)
        
        result = [x for x in list] 

        self.split_universe = np.array_split(result, 2) 
            
        return result


    # this event fires whenever we have changes to our universe
    def OnSecuritiesChanged(self, changes):
        # liquidate removed securities
        for security in changes.RemovedSecurities:
            if security.Symbol in self.universe:
                self.universe.remove(security.Symbol)

        # we want equal allocation in each security in our universe
        for security in changes.AddedSecurities:
            if security.Symbol not in self.universe:
                self.universe.append(security.Symbol)
                    
    def OnEndOfAlgorithm(self):
        for trade in self.TradeBuilder.ClosedTrades:
            self.Log(str(trade.Symbol)+
                    ", MAE: "+ str(trade.MAE)+
                    ", MFE: "+ str(trade.MFE)+
                    ", EOT Drawdown: "+ str(trade.EndTradeDrawdown)+
                    ", Entry Price: "+str(trade.EntryPrice)+
                    ", Entry Time: "+str(trade.EntryTime)+
                    ", Exit Price: "+str(trade.ExitPrice)+
                    ", Exit Time: "+str(trade.ExitTime)+
                    ", Quantity: "+str(trade.Quantity)+
                    ", Profit Loss: "+str(trade.ProfitLoss)+
                    ", Direction: "+str(trade.Direction)+
                    ", Duration: "+str(trade.Duration)
                    )


class StockDataSource(PythonData):
    
    def GetSource(self, config, date, isLive):
        url = "https://www.dropbox.com/s/rucwr4ph5w02a84/sp900universe.csv?dl=1" if isLive else \
                "https://www.dropbox.com/s/1hk2oxeban0lt5u/ETFS.csv?dl=1"
                
        return SubscriptionDataSource(url, SubscriptionTransportMedium.RemoteFile)
    
    def Reader(self, config, line, date, isLive):
        if not (line.strip() and line[0].isdigit()): return None
        
        stocks = StockDataSource()
        stocks.Symbol = config.Symbol
        
        csv = line.split(',')
        if isLive:
            stocks.Time = date
            stocks["Symbols"] = csv
        else:   
            stocks.Time = datetime.strptime(csv[0], "%m/%d/%Y")
            stocks["Symbols"] = csv[1:]
        return stocks      
                    
        
class FactorCollection(object):
    
    def cut_dataframe(self,df,start_date=None,end_date=None):
        all_dates = df.index.values
        if not end_date:
            end_date = all_dates.max()
        if not start_date:
            start_date = all_dates.min()
        cut_df = df.ix[df.index.searchsorted(start_date):(1+df.index.searchsorted(end_date))]
        return cut_df
         
    def slopeWeekly(self,df,date,column_name):
        day = df.index.asof(pd.to_datetime(date))
        cut_df = self.cut_dataframe(df,end_date=day)
        ema = pd.ewma(cut_df[column_name], span=50).ix[-25:]
        ema_dates = ema.index.values
        ema_time_delta_days = pd.Timedelta(ema_dates.max()-ema_dates.min()).total_seconds()/3600/24
        slope = (ema.values[-1] - ema.values[0]) / ema_time_delta_days
        return slope

    def volumentumWeekly(self,df,date,column_name):
        day = df.index.asof(pd.to_datetime(date))
        friday1 = df.index.asof(day - timedelta(days=(day.weekday() - 4) % 7, weeks=0))
        friday2 = df.index.asof(day - timedelta(days=(day.weekday() - 4) % 7, weeks=1))
        one_week = [day - timedelta(i) for i in range(7)]
        six_months = [day - timedelta(i) for i in range(180)]
        avg_week_volume = df.loc[one_week].dropna()[column_name].mean()
        avg_six_months_volume = df.loc[six_months].dropna()[column_name].mean()
        return (df[column_name].loc[friday1] - df[column_name].loc[friday2]) * avg_week_volume / avg_six_months_volume

    def volumentumMonthly(self,df,date,column_name):
        day = df.index.asof(pd.to_datetime(date))
        end_of_month1 = pd.to_datetime(date(day.year,day.month,1)) - timedelta(days=1)
        end_of_month1 = df.index.asof(pd.to_datetime(end_of_month1))
        end_of_month2 = pd.to_datetime(date(end_of_month1.year,end_of_month1.month,1)) - timedelta(days=1)
        end_of_month2 = df.index.asof(pd.to_datetime(end_of_month2))
        one_month = [df.index.asof(day - timedelta(i)) for i in range(30)]
        twelve_months = [df.index.asof(day - timedelta(i)) for i in range(360)]
        avg_monthly_volume = df.loc[one_month].dropna()[column_name].mean()
        avg_twelve_months_volume = df.loc[twelve_months].dropna()[column_name].mean()
        return (df[column_name].loc[end_of_month1] - df[column_name].loc[end_of_month2]) \
               * avg_monthly_volume / avg_twelve_months_volume

    def momentumNMo(self,df,date,column_name,number_of_months):
        end_date = df.index.asof(pd.to_datetime(date))
        start_date = end_date - timedelta(days=30*number_of_months)
        start_date = df.index.asof(start_date)
        cut_df = self.cut_dataframe(df,start_date=start_date,end_date=date)[column_name]
        cut_df_minus1 = self.cut_dataframe(df,start_date=start_date-BDay(1),end_date=end_date-BDay(1))[column_name]
        len_df = len(cut_df)
        len_df_minus1 = len(cut_df_minus1)
        len_min = min(len_df,len_df_minus1)
        daily_returns = cut_df.values[:len_min] / cut_df_minus1.values[:len_min]
        return daily_returns.mean()

    def meanReversion(self,df,date,column_name,n_days,N_days):
        day = df.index.asof(pd.to_datetime(date))
        n_day_range = [df.index.asof(day - timedelta(i)) for i in range(n_days)]
        N_day_range = [df.index.asof(day - timedelta(i)) for i in range(N_days)]
        avg_n_days = df.loc[n_day_range].dropna()[column_name].mean()
        avg_N_days = df.loc[N_day_range].dropna()[column_name].mean()
        return avg_n_days / avg_N_days - 1.0

    def highLowRange(self,df,date,column_name):
        day = df.index.asof(pd.to_datetime(date))
        fifty_two_day_range = [df.index.asof(day - timedelta(i)) for i in range(52*5)]
        fifty_two_week_low = df.loc[fifty_two_day_range].dropna()['low'].min()
        fifty_two_week_high = df.loc[fifty_two_day_range].dropna()['high'].max()
        current_price = df.loc[day][column_name]
        return (current_price - fifty_two_week_low) / (fifty_two_week_high - fifty_two_week_low)

    def moneyFlow(self,df,date):
        day = df.index.asof(pd.to_datetime(date))
        close = df.loc[day]['close']
        low = df.loc[day]['low']
        high = df.loc[day]['high']
        volume = df.loc[day]['volume']
        return (((close - low) - (high - close)) / (high - low)) * volume

    def moneyFlowPersistency(self,df,date,number_of_months):
        day = df.index.asof(pd.to_datetime(date))
        day_range = [df.index.asof(day - timedelta(i)) for i in range(number_of_months*30)]
        money_flows = np.array([self.moneyFlow(df,day1 - BDay(1)) for day1 in day_range])
        signs_of_money_flows = np.sign(money_flows)
        return (signs_of_money_flows[signs_of_money_flows>0]).sum() / (number_of_months*30)

    def slopeDaily(self,df,date,column_name):
        day = df.index.asof(pd.to_datetime(date))
        cut_df = self.cut_dataframe(df,end_date=day)
        ema = pd.ewma(cut_df[column_name], span=10).ix[-5:]
        ema_dates = ema.index.values
        ema_time_delta_days = pd.Timedelta(ema_dates.max()-ema_dates.min()).total_seconds()/3600/24
        slope = (ema.values[-1] - ema.values[0]) / ema_time_delta_days
        return slope

    def slopeMonthly(self,df,date,column_name):
        day = df.index.asof(pd.to_datetime(date))
        cut_df = self.cut_dataframe(df,end_date=day)
        ema = pd.ewma(cut_df[column_name], span=300).ix[-150:]
        ema_dates = ema.index.values
        ema_time_delta_days = pd.Timedelta(ema_dates.max()-ema_dates.min()).total_seconds()/3600/24
        slope = (ema.values[-1] - ema.values[0]) / ema_time_delta_days
        return slope

    def pxRet(self,df,date,column_name, number_of_days):
        day = df.index.asof(pd.to_datetime(date))
        day_minus_n_days = df.index.asof(day - timedelta(days=number_of_days))
        cut_df = self.cut_dataframe(df,end_date=day,start_date=day_minus_n_days)[column_name]
        return (cut_df.values[-1] - cut_df.values[0])/cut_df.values[0]

    def currPxRet(self,df,date,column_name):
        day = df.index.asof(pd.to_datetime(date))
        price = df.loc[day][column_name]
        day_start = df.index.asof(day - timedelta(days=3*360))
        price_mean = self.cut_dataframe(df,start_date=day_start, end_date=day)[column_name].mean()
        return 1.0 - price_mean / price

    def nDayADR(self,df,date,column_name,number_of_days):
        day = df.index.asof(pd.to_datetime(date))
        day_minus_n_days = df.index.asof(day - timedelta(days=number_of_days))
        cut_df = self.cut_dataframe(df,end_date=day,start_date=day_minus_n_days)[column_name]
        cut_df_minus1 = self.cut_dataframe(df,end_date=day,start_date=day_minus_n_days)[column_name]
        return (cut_df.values / cut_df_minus1.values).mean()

    def nDayADP(self,df,date,column_name,number_of_days):
        day = df.index.asof(pd.to_datetime(date))
        day_minus_n_days = df.index.asof(day - timedelta(days=number_of_days))
        cut_df = self.cut_dataframe(df,end_date=day,start_date=day_minus_n_days)[column_name]
        cut_df_minus1 = self.cut_dataframe(df,end_date=day,start_date=day_minus_n_days)[column_name]
        return (cut_df.values - cut_df_minus1.values).mean()

    def pxRet2(self,df,date,column_name,N_days,n_days):
        return -0.5 * self.pxRet(df,date,column_name,N_days) + 0.5 * self.pxRet(df,date,column_name,n_days)

    def currPxRetSlope(self,df,date,column_name):
        return -0.5 * self.currPxRet(df,date,column_name) + 0.5 * self.slopeDaily(df,date,column_name)

    def allFactors(self,df_list,date,column_name):
        """All Factors as an array"""
        f1 = []
        f10 = []
        f12 = []
        f14 = []
        f20 = []
        f28 = []
        for df in df_list:
            if not df.empty:
                f1.append(float(self.slopeWeekly(df, date, column_name)))
                f10.append(float(self.highLowRange(df, date, column_name)))
                f12.append(float(self.moneyFlowPersistency(df, date, 1)))
                f14.append(float(self.moneyFlowPersistency(df, date, 6)))
                f20.append(float(self.pxRet(df, date, column_name, 90)))
                f28.append(float(self.currPxRetSlope(df, date, column_name)))

        factor_names = ['f1', 'f10', 'f12', 'f14', 'f20', 'f28']
                         
        return [f1,f10,f12,f14,f20,f28], factor_names