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

from System import *
from QuantConnect import *
from QuantConnect.Data import *
from QuantConnect.Algorithm import *
from QuantConnect.Indicators import *
from System.Collections.Generic import List
import decimal as d
import numpy as np
import time
from datetime import datetime
import numpy as np
from scipy import stats
import pandas as pd


class AFCMOM(QCAlgorithm):
    '''Basic template algorithm simply initializes the date range and cash'''

    def Initialize(self):
        '''Initialise the data and resolution required, as well as the cash and start-end dates for your algorithm. All algorithms must initialized.'''
        self.first = -1
        self.bi_weekly = 0
        self.SetStartDate(2006,1,1)  
        self.SetEndDate(2009,1,1)    
        self.SetCash(100000)           
        
        self.SetBrokerageModel(BrokerageName.InteractiveBrokersBrokerage, AccountType.Margin)
        
        self.spy = self.AddEquity("SPY", Resolution.Daily).Symbol
        
        self.AddUniverse(self.CoarseSelectionFunction,self.FineSelectionFunction)
        self.UniverseSettings.Resolution = Resolution.Daily
        
        self.spy_200_sma = self.SMA("SPY",200,Resolution.Daily)
        
        self.Schedule.On(self.DateRules.Every(DayOfWeek.Wednesday,DayOfWeek.Wednesday), \
                 self.TimeRules.At(12, 0), \
                 Action(self.rebalnce))
                 
        self.stocks_to_trade = []

        self.SetWarmUp(201)
        
        self.splotName = 'Strategy Info'
        sPlot = Chart(self.splotName)
        sPlot.AddSeries(Series('Leverage',  SeriesType.Line, 0))
        self.AddChart(sPlot)
        
    def OnData(self, data):
        #if self.IsWarmingUp: return
    
        # assuming daily mode,dont chart in a smaller res and kill quota
        self.account_leverage = self.Portfolio.TotalAbsoluteHoldingsCost / self.Portfolio.TotalPortfolioValue
        self.Plot(self.splotName,'Leverage', float(self.account_leverage))
        
    def CoarseSelectionFunction(self, coarse):
        filtered_stocks = filter(lambda x: x.DollarVolume >250000,coarse)
        filtered_stocks = filter(lambda x: x.HasFundamentalData,filtered_stocks)
        filtered_stocks = filter(lambda x: x.Price >=20,filtered_stocks)
        filtered_stocks = filtered_stocks[:100]
        return [stock.Symbol for stock in filtered_stocks]
        
    def FineSelectionFunction(self, fine):
        filtered_stocks = filter(lambda x: x.SecurityReference.IsPrimaryShare,fine)
        return [stock.Symbol for stock in filtered_stocks]
        
    def OnSecuritiesChanged(self, changes):
        dt = datetime(self.Time.year,self.Time.month,self.Time.day)
        if dt.weekday() != 3 or self.Securities[self.spy].Price < self.spy_200_sma.Current.Value:
            return 
        
        # manage stocks_to_trade(add and remove items)
        for security in changes.RemovedSecurities:
            # update stocks_to_trade list
            if security.Symbol in self.stocks_to_trade:
                self.stocks_to_trade.remove(security.Symbol)
            if security.Invested:
                self.Liquidate(security.Symbol)
        for security in changes.AddedSecurities:
            #security.MarginModel = PatternDayTradingMarginModel()
            if security.Symbol not in self.stocks_to_trade:
                self.stocks_to_trade.append(security.Symbol)
        
        if self.stocks_to_trade:
            for stock in self.stocks_to_trade:
                ATR = self.my_ATR(stock,14)
                self.stocks_to_trade.sort(key = lambda x: self.get_slope(stock,90),reverse= True)
                maximum_range = int(round(len(self.stocks_to_trade) * 0.10))
                self.stocks_to_trade[:maximum_range]
                cash = float(self.Portfolio.Cash)
                oo = len(self.Transactions.GetOpenOrders(stock))
                if self.Securities[stock].Price >self.moving_average(stock,100) and not self.gapper(stock,90) and cash >0 and not oo:
                    self.SetHoldings(stock,self.weight(stock,ATR))
    
        
    def rebalnce(self):
        self.bi_weekly +=1
        if self.bi_weekly%2 == 0:
            for stock in self.Portfolio.Values:
                if stock.Invested:
                    symbol = stock.Symbol
                    shares_held = float(self.Portfolio[symbol].Quantity)
                    if (self.Securities[symbol].Price < self.moving_average(symbol,100) and shares_held >0) or (self.gapper(symbol,90) and shares_held>0):
                        self.Liquidate(symbol)
                    else:
                        if shares_held >0: 
                            ATR = self.my_ATR(symbol,20)
                            cost_basis = float(self.Portfolio[symbol].AveragePrice)
                            shares_held = float(self.Portfolio[symbol].Quantity)
                            percent_of_p = ((cost_basis * shares_held )/ float(self.Portfolio.TotalPortfolioValue))
                            weight= self.weight(symbol,ATR)
                            diff_in_desired_weight = weight -percent_of_p
                            if diff_in_desired_weight < 0:
                                order_amount = shares_held * diff_in_desired_weight
                                self.MarketOrder(symbol,order_amount)
                
    
    def GetHistory(self, security, period):
        security_data = self.History([security],period,Resolution.Daily)
        
        # check if candle has close component(assume others?)
        if 'close' not in security_data.columns:
            self.Log("History had no Close for %s"%security)
            return None
        
        # we need enough for the np.diff which removes 1 from length
        if len(security_data.close.index) < period: 
            self.Log("Close test had too few or no values for %s with period %d"%(security, period))
            return None
            
        return security_data
        
    def gapper(self,security,period):
        if not self.Securities.ContainsKey(security):
            return 0
            
        security_data = self.GetHistory(security, period)
        if security_data is None: return 0
    
        close_data = security_data.close.values
        
        return np.max(np.abs(np.diff(close_data))/close_data[:-1])>=0.15
        
    def get_slope(self,security,period):
        if not self.Securities.ContainsKey(security):
            return 0
            
        security_data = self.GetHistory(security, period)
        if security_data is None: return 0
    
        y = np.log(security_data.close.values)
        x = np.arange(len(y))
        slope, intercept, r_value, p_value, std_err = stats.linregress(x,y)
        
        return ((np.exp(slope)**252)-1)*(r_value**2)
        
    def my_ATR(self,security,period):
        if not self.Securities.ContainsKey(security):
            return 0
        self.first+=1
        
        security_data = self.GetHistory(security, period)
        if security_data is None: return 0
        
        c_data = security_data.close.values
        l_data= security_data.low.values
        h_data = security_data.high.values
        
        true_range = h_data-l_data
        average_true_range = np.mean(true_range)
        average_true_range_smooted = ((average_true_range*13)+true_range[-1])/14
        
        return average_true_range_smooted if not self.first else average_true_range
    
    def weight(self,security,atr):
        risk = float(self.Portfolio.TotalPortfolioValue)*0.0001
        
        # prevent div by zero
        if atr == 0.:
            return 0.
            
        return (((risk/atr) * float(self.Securities[security].Price))/float(self.Portfolio.TotalPortfolioValue)*100)
        
    def moving_average(self,security,period):
        if not self.Securities.ContainsKey(security):
            return 0
        
        security_data = self.GetHistory(security, period)
        if security_data is None: return 0
        
        return security_data.close.mean()