import statsmodels.api as sm
import numpy

MAX_LEVERAGE = 2.0

beta_samples = 252

alpha_samples = 15

volatility_samples = 5

MAX_VOLATILITY = 0.05

class BasicTemplateAlgorithm(QCAlgorithm):
    
    def __init__(self):
        
        pass
    
    def Initialize(self):
    
        self.SetCash(100000)
        
        self.SetStartDate(2008,1,1)
    
        #self.SetEndDate(2007,12,31)
        
        self.SPY = self.AddEquity("SPY", Resolution.Minute).Symbol
        self.QQQ = self.AddEquity("QQQ", Resolution.Minute).Symbol
        self.TLT = self.AddEquity("TLT", Resolution.Minute).Symbol
        
        #self.SetBrokerageModel(BrokerageName.InterativeBrokersBrokerage, AccountType.Margin)
        
        self.Schedule.On(self.DateRules.EveryDay("SPY"), self.TimeRules.AfterMarketOpen("SPY", 15), Action(self.Rebalance))
    
    def OnData(self, slice):
        
        pass
    
    def min_variance(self, returns):
    
        cov = numpy.cov(returns)
        vu = numpy.ones(cov.shape[0])
        if (numpy.linalg.det(cov) != 0):
            num = numpy.linalg.solve(cov,vu)
            den = numpy.dot(vu, num)
            w = num / den
            return w / numpy.sum(numpy.absolute(w))
        else:
            return None
            
    def Rebalance(self):
        
        try:

            self.spy_price = float(self.Securities["SPY"].Price)
            if (self.spy_price == 0):
                self.Log("Missing SPY price")
                return
        
            self.qqq_price = float(self.Securities["QQQ"].Price)
            if (self.qqq_price == 0):
                self.Log("Missing QQQ price")
                return
        
            self.tlt_price = float(self.Securities["TLT"].Price)
            if (self.tlt_price == 0):
                self.Log("Missing TLT price")
                return
    
            spy_prices = numpy.asarray([float(x.Price) for x in self.History("SPY",beta_samples,Resolution.Daily)] + [self.spy_price])
            spy_returns = numpy.nan_to_num(spy_prices[1:] / spy_prices[:-1] - 1.0)
            
            qqq_prices = numpy.asarray([float(x.Price) for x in self.History("QQQ",beta_samples,Resolution.Daily)] + [self.qqq_price])
            qqq_returns = numpy.nan_to_num(qqq_prices[1:] / qqq_prices[:-1] - 1.0)
    
            tlt_prices = numpy.asarray([float(x.Price) for x in self.History("TLT",beta_samples,Resolution.Daily)] + [self.tlt_price])
            tlt_returns = numpy.nan_to_num(tlt_prices[1:] / tlt_prices[:-1] - 1.0)
            
            if (len(spy_returns) == len(qqq_returns) == len(tlt_returns)):
                security_returns = numpy.vstack([qqq_returns,tlt_returns])
            else:
                self.Log("Missing data")    
                return
                    
            pfolio_weights = self.min_variance(security_returns[:,-alpha_samples:])
            if (pfolio_weights is None):
                self.Log("Min_variance error")    
                return
        
            pfolio_returns = numpy.nan_to_num((pfolio_weights[:,None]*security_returns).sum(axis=0))
        
            pfolio_beta = numpy.nan_to_num(sm.OLS(pfolio_returns[-beta_samples:],sm.add_constant(spy_returns[-beta_samples:], prepend=True)).fit().params[1])
        
            benchmark_split = -pfolio_beta / (1.0 + abs(pfolio_beta))
    
            pfolio_split = 1.0 / (1.0 + abs(pfolio_beta))
        
            pfolio_std = numpy.std(pfolio_split * pfolio_returns[-volatility_samples:] + benchmark_split * spy_returns[-volatility_samples:] )
        
            leverage = numpy.clip(MAX_VOLATILITY / (pfolio_std * numpy.sqrt(252)),0.0,MAX_LEVERAGE)
        
            self.SetHoldings("SPY", benchmark_split * leverage)
            self.Log("SPY  {0: 4.2f} % {1: 4.2f} ".format((benchmark_split * leverage) / MAX_LEVERAGE * 100.0, self.spy_price))
        
            self.SetHoldings("QQQ", pfolio_split * pfolio_weights[0] * leverage)
            self.Log("QQQ  {0: 4.2f} % {1: 4.2f} ".format((pfolio_split * pfolio_weights[0] * leverage) / MAX_LEVERAGE * 100.0, self.qqq_price))
        
            self.SetHoldings("TLT", pfolio_split * pfolio_weights[1] * leverage)
            self.Log("TLT {0: 4.2f} % {1: 4.2f} ".format((pfolio_split * pfolio_weights[1] * leverage) / MAX_LEVERAGE * 100.0, self.tlt_price))

            self.Log("CASH {0: 4.2f} %".format((MAX_LEVERAGE - leverage) / MAX_LEVERAGE * 100.0))
        
        except Exception as e: pass