from QuantConnect.Data.UniverseSelection import *
import math
import numpy as np
import pandas as pd
import decimal as d
from scipy import stats 


class FundamentalFactorAlgorithm(QCAlgorithm):
    def slope(self,ts):
        """
        Input: Price time series.
        Output: Annualized exponential regression slope, multipl
        """
        x = np.arange(len(ts))
        log_ts = np.log(ts)
        slope, intercept, r_value, p_value, std_err = stats.linregress(x, log_ts)
        annualized_slope = (np.power(np.exp(slope), 250) - 1) * 100
        return annualized_slope * (r_value ** 2)

    def Initialize(self):

        self.SetStartDate(2003, 1, 1)  #Set Start Date
        self.SetEndDate(2018, 3, 1)  #Set Start Date       
        self.SetCash(10000)            #Set Strategy Cash
    
        
        self.UniverseSettings.Resolution = Resolution.Daily
        self.AddUniverse(self.CoarseSelectionFunction, self.FineSelectionFunction)
        self.SetBrokerageModel(BrokerageName.InteractiveBrokersBrokerage, AccountType.Cash); 
        self.spy = self.AddEquity("SPY", Resolution.Minute).Symbol 
        self.holding_months = 1
        self.num_screener = 100
        self.num_stocks = 3
        self.formation_days = 200
        self.lowmom = False
        self.month_count = self.holding_months
        self.Schedule.On(self.DateRules.MonthStart("SPY"), self.TimeRules.At(0, 0), Action(self.monthly_rebalance))
        self.Schedule.On(self.DateRules.MonthStart("SPY"), self.TimeRules.At(10, 0), Action(self.rebalance))
        #self.Schedule.On(self.DateRules.EveryDay("SPY"), self.TimeRules.At(10, 0), Action(self.daily_rebalance))
        # rebalance the universe selection once a month
        self.rebalence_flag = 0
        # make sure to run the universe selection at the start of the algorithm even it's not the manth start
        self.first_month_trade_flag = 1
        self.trade_flag = 0 
        self.symbols = None
        self.open_stop_market_orders = {}
            
        # This version uses the average of two momentum slopes.
        # Want just one? Set them both to the same number.
        self.momentum_window = 60  # first momentum window.
        self.momentum_window2 = 90  # second momentum window
    
        # Limit minimum slope. Keep in mind that shorter momentum windows
        # yield more extreme slope numbers. Adjust one, and you may want
        # to adjust the other.
        self.minimum_momentum = 60  # momentum score cap
        self.exclude_days = 5
        self.size_method = 2
        self.use_bond_etf = True
        self.bond_etf = 'TLT'
 
    def CoarseSelectionFunction(self, coarse):
        if self.rebalence_flag or self.first_month_trade_flag:
            # drop stocks which have no fundamental data or have too low prices
            selected = [x for x in coarse if (x.HasFundamentalData) and (float(x.Price) > 5)]
            # rank the stocks by dollar volume 
            filtered = sorted(selected, key=lambda x: x.DollarVolume, reverse=True) 
    
            return [ x.Symbol for x in filtered[:200]]
        else:
            return self.symbols


    def FineSelectionFunction(self, fine):
      
        if self.rebalence_flag or self.first_month_trade_flag:
            try:
                filtered_fine = [x for x in fine if (x.ValuationRatios.EVToEBITDA > 0) 
                                                    and (x.EarningReports.BasicAverageShares.ThreeMonths > 0) 
                                                    and x.EarningReports.BasicAverageShares.ThreeMonths * (x.EarningReports.BasicEPS.TwelveMonths*x.ValuationRatios.PERatio) > 2e9]
            except:
                filtered_fine = [x for x in fine if (x.ValuationRatios.EVToEBITDA > 0) 
                                                and (x.EarningReports.BasicAverageShares.ThreeMonths > 0)] 

            top = sorted(filtered_fine, key = lambda x: x.ValuationRatios.EVToEBITDA, reverse=True)[:self.num_screener]
            self.symbols = [x.Symbol for x in top]
            
            self.rebalence_flag = 0
            self.first_month_trade_flag = 0
            self.trade_flag = 1
            return self.symbols
        else:
            return self.symbols
    
    def OnData(self, data):
        pass
    
    def monthly_rebalance(self):
        self.rebalence_flag = 1


    def inv_vola_calc(self,ts):
        """
        Input: Price time series.
        Output: Inverse exponential moving average standard deviation. 
        Purpose: Provides inverse vola for use in vola parity position sizing.
        """
        returns = np.log(ts).diff()
        stddev = returns.ewm(halflife=20, ignore_na=True, min_periods=0,
                         adjust=True).std(bias=False).dropna()
        return 1 / stddev.iloc[-1]

    def rebalance(self):
        # Get data
        hist_window = max(self.momentum_window,
                      self.momentum_window2) + self.exclude_days
        if self.symbols is None: return
        stocks = self.symbols
        #self.Debug("Stocks " + str(len(stocks)))
        #hist = self.History(context.security_list, "close", hist_window, "1d")
        hist = self.History(stocks, self.formation_days, Resolution.Daily)
        current = self.History(stocks, 1, Resolution.Minute)
        c_data = hist["close"].unstack(level=0)
        #self.Debug("c_data " + str(c_data))
        data_end = -1 * (self.exclude_days + 1 )  # exclude most recent data

        momentum1_start = -1 * (self.momentum_window + self.exclude_days)
        momentum_hist1 = c_data[momentum1_start:data_end]
        momentum2_start = -1 * (self.momentum_window2 + self.exclude_days)
        momentum_hist2 = c_data[momentum2_start:data_end]
    
        # Calculate momentum scores for all stocks.
        momentum_list = momentum_hist1.apply(self.slope)  # Mom Window 1
        #elf.Debug("momentum_list " + str((momentum_list)))
        momentum_list2 =  momentum_hist2.apply(self.slope)   # Mom Window 2
        # Combine the lists and make average
        momentum_concat = pd.concat((momentum_list, momentum_list2))
        mom_by_row = momentum_concat.groupby(momentum_concat.index)
        mom_means = mom_by_row.mean()
        
        # Sort the momentum list, and we've got ourselves a ranking table.
        ranking_table = mom_means.sort_values(ascending=False)
        #self.Debug("ranking_table " + str(len(ranking_table)))
        # Get the top X stocks, based on the setting above. Slice the dictionary.
        # These are the stocks we want to buy.
        buy_list = ranking_table[:self.num_stocks]
        self.Debug("==========================================================================")
        self.Debug(" buy list " + str(len(buy_list)))
        final_buy_list = buy_list[buy_list > self.minimum_momentum] # those who passed minimum slope requirement
        self.Debug("Final buy list " + str(len(final_buy_list))) 
        # Calculate inverse volatility, for position size.
        inv_vola_table = c_data[buy_list.index].apply(self.inv_vola_calc)
        # sum inv.vola for all selected stocks.
        sum_inv_vola = np.sum(inv_vola_table)


        spy_hist_150 = self.History([self.spy], 150, Resolution.Daily).loc[str(self.spy)]['close']
        spy_hist_200 = self.History([self.spy], 200, Resolution.Daily).loc[str(self.spy)]['close']
        if self.Securities[self.spy].Price > (spy_hist_150.mean() + spy_hist_200.mean())/2:
            can_buy = True
        else:
            can_buy = False
       
    
        equity_weight = 0.0 # for keeping track of exposure to stocks
        self.Debug("===================: " + str(self.open_stop_market_orders))
        # Sell positions no longer wanted.
        for security in self.Portfolio.Keys:
            #self.Debug("Portfolio key: " + str(security) + " value: " + str(security.Value))
           
            if (security not in final_buy_list):
                if (security.Value != self.bond_etf):
                    #self.Debug(" Sell security :" + str(security) + " Invested: " + str(self.Portfolio[security].Invested) + " [2] In open_stop_market_orders: " + str(str(security) in self.open_stop_market_orders.keys()))
                    #an tin eixa apo prin kai den tin thelw twra kai den exei ginei Stop Loss
                    if  self.Portfolio[security].Invested and str(security) in self.open_stop_market_orders.keys() :
                        shortOrder = self.open_stop_market_orders[str(security)]
                        shortOrder.Cancel("Short filled")
                        self.open_stop_market_orders.pop(str(security),None)
                        self.Debug("{0}: Short order is filled. ".format(shortOrder.OrderType) + str(security))
                    self.SetHoldings(security, 0)

        vola_target_weights = inv_vola_table / sum_inv_vola
        quantity = self.Portfolio.TotalPortfolioValue*  d.Decimal(.33)
        for security in final_buy_list.index:
            
            # allow rebalancing of existing, and new buys if can_buy, i.e. passed trend filter.
            self.Debug("** Invested: "+ str(self.Portfolio[security].Invested) + " or can_buy: " + str(can_buy) )
            if (self.Portfolio[security].Invested) or (can_buy): 
                if (self.size_method == 1):
                    weight = vola_target_weights[security]
                elif (self.size_method == 2):
                    
                    weight = (0.99 / len(final_buy_list))
             
                #self.Debug("Number of stocks " + str(self.num_stocks) + " Rebalance security " + str(security) + " with weight " + str(weight)  )
                #self.SetHoldings(security, weight)
                self.MarketOrder(security, quantity/self.Securities[security].Price)
                #self.Debug("333333")
                self.Debug("[3] In open_stop_market_orders " + str(str(security) in self.open_stop_market_orders.keys()) )
                if  str(security) in self.open_stop_market_orders.keys():
                    #self.Debug("4444")
                    longOrder = self.open_stop_market_orders[security]
                    updateOrderFields = UpdateOrderFields()
                    updateOrderFields.StopPrice = self.Securities[security].Price * d.Decimal(.9)
                    updateOrderFields.Tag = "Update #{0}".format(len(longOrder.UpdateRequests) + 1)
                    longOrder.Update(updateOrderFields)
                    self.Debug("Updated price " + str(longOrder.Get(OrderField.StopPrice)) + " security "+ str(security))
                else:
                    newTicket = self.StopMarketOrder(security, -quantity/self.Securities[security].Price , self.Securities[security].Price  * d.Decimal(.9))
                    #self.Debug("StopMarketOrder price " + str(self.Securities[security].Price  * d.Decimal(.9)))
                    self.open_stop_market_orders[str(security)] = newTicket
                    #self.Debug("===================: " + str(self.open_stop_market_orders))
                    #self.Debug("Open market orders size: " + str(len(self.open_stop_market_orders)))
                    self.Debug("{0}: Submitting StopMarketOrder. ".format(newTicket.OrderType) + str(security))
                equity_weight += weight
        
    
        # Fill remaining portfolio with bond ETF
        etf_weight = max(0.99 - equity_weight, 0.0)
    
        print ('equity exposure should be %s ' % str(equity_weight))
    
        if (self.use_bond_etf):
            self.Debug("Buy ETF" )
            self.AddEquity("TLT")
            self.SetHoldings(self.bond_etf, etf_weight)
  
    def OnOrderEvent(self, orderEvent):
        order = self.Transactions.GetOrderById(orderEvent.OrderId)
        self.Debug("{0}: {1}: {2}".format(self.Time, order.Type, orderEvent))