from System import *
from QuantConnect import *
from QuantConnect.Algorithm import *
from QuantConnect.Python import PythonQuandl
from datetime import datetime, timedelta

class BasicTemplateAlgorithm(QCAlgorithm):

    def Initialize(self):

        self.SetStartDate(2015, 1, 1)
        self.SetEndDate(2017, 12, 31)
        self.SetCash(10000)
        
        self.returns = {}
        self.rank_returns = {}
        self.basket = []
        self.net = {}
        self.net_charge = 0
        
        # 'XLC', 'XLRE',
        
        self.etfs = ['XLK', 'XLI', 'XLE', 'XLC', 'XLF', 'XLY', 'XLB', 'XLU', 'XLP', 'XLV']
        for i in self.etfs:
            self.AddEquity(i, Resolution.Daily)
        
        self.AddEquity('BND', Resolution.Daily)
        
        self.tbill = 'USTREASURY/BILLRATES'
        self.AddData(TBill, self.tbill, Resolution.Daily)
        
        self.AddEquity('SPY', Resolution.Daily)
        self.Schedule.On(self.DateRules.MonthStart('SPY'), self.TimeRules.AfterMarketOpen('SPY', 1), Action(self.Signal))
        self.Schedule.On(self.DateRules.MonthStart('SPY'), self.TimeRules.AfterMarketOpen('SPY', 2), Action(self.Trade))

    def Signal(self):
        
        slices = self.History(self.etfs, 253, Resolution.Daily)
        
        self.returns = {}
        if not slices.empty:
            for i in self.etfs:
                df = slices.loc[i]
                etf_prev = df['close'][0]
                etf_curr = df['close'][-1]
                self.returns[i] = ((etf_curr - etf_prev) / etf_prev) * 100
        
        self.rank_returns = {key: rank for rank, key in enumerate(sorted(self.returns, key = self.returns.get, reverse = True), 1)}
        self.basket = [key for key, value in self.rank_returns.items() if value <= 3]
        
        self.net = {}
        for i in self.basket:
            if self.returns[i] > 0:
                self.net[i] = self.returns[i] - float(self.Securities[self.tbill].Price)
            else:
                self.net[i] = - 1
                
        self.Log(str(self.Time) + ' Net: ' + str(self.net))
    
    def Trade(self):
        
        etf_weight = 0.33
        
        self.net_charge = sum(1 for x in self.net if self.net[x] > 0)
        for kvp in self.Portfolio:
            security_hold = kvp.Value
            if security_hold.Invested and security_hold.Symbol.Value == 'BND' and self.net_charge == 3:
                self.Liquidate('BND')
        
        for kvp in self.Portfolio:
            security_hold = kvp.Value
            if security_hold.Invested and (security_hold.Symbol.Value != 'BND') and (security_hold.Symbol.Value not in self.basket or self.net[security_hold.Symbol.Value] < 0):
                self.Liquidate(security_hold.Symbol)
        
        added_symbols = []
        for key, value in self.net.items():
            if value > 0:
                if not self.Portfolio[key].Invested:
                    added_symbols.append(key)
        
        for added in added_symbols:
            self.SetHoldings(added, etf_weight)
        
        invested = []
        for kvp in self.Portfolio:
            security_hold = kvp.Value
            if security_hold.Invested:
                invested.append(security_hold.Symbol.Value)
        
        left_over = float(1 - (len(invested) * etf_weight))
        if self.net_charge < 3:
            self.SetHoldings('BND', left_over)
    
class TBill(PythonQuandl):
    
    def __init__(self):
        
        self.ValueColumnName = '52 Wk Coupon Equiv'