from math import ceil,floor,isnan
from datetime import datetime
import pandas as pd
import numpy as np
import cvxpy as cv

class AssetAllocationAlgorithm(QCAlgorithm):
    
    def Initialize(self):
        self.SetStartDate(2016, 1, 1)  #Set Start Date
        self.SetEndDate(2018, 1, 1)    #Set End Date
        self.SetCash(100000)            #Set Strategy Cash

        tickers = ["IEF", "TLT", "SPY", "EFA", "EEM", "JPXN", "XLK"]
        self.symbols = [] 
        for i in tickers:
            self.symbols.append(self.AddEquity(i, Resolution.Daily).Symbol)
        for syl in self.symbols:
            syl.window = RollingWindow[TradeBar](252) 

        self.Schedule.On(self.DateRules.MonthStart("SPY"), self.TimeRules.AfterMarketOpen("SPY"), Action(self.Rebalancing))
       

    def OnData(self, data):
        if data.ContainsKey("SPY"):
            for syl in self.symbols:
                syl.window.Add(data[syl])

    def Rebalancing(self):
        data = {}
        for syl in self.symbols:
            data[syl] = [float(i.Close) for i in syl.window]
        df_price = pd.DataFrame(data,columns=data.keys())
        # To order prices from oldest to latest
        
        # df_price = df_price.sort_values(by = 0, ascending = False)
        df_price = df_price.iloc[::-1]
        
        # Calculating returns
        # daily_return = df_price.pct_change().dropna().values
        daily_return = df_price.pct_change().dropna()

        if daily_return.size != 0:
            a = PortfolioOptimization(daily_return, 0, len(data))
            opt_weight = a.opt_portfolio()  
            
            if isnan(sum(opt_weight)): return
            self.Log(str(opt_weight))
            
            for i in range(len(data)):
                self.SetHoldings(df_price.columns[i], opt_weight[i])
                
                # equally weighted
                # self.SetHoldings(self.symbols[i], 1.0/len(data))
class PortfolioOptimization(object):

    import numpy as np
    import pandas as pd
    import cvxpy as cv

    def __init__(self, df_return, risk_free_rate, num_assets):
        
        self.daily_return = df_return
        self.risk_free_rate = risk_free_rate
        self.n = num_assets # numbers of risk assets in portfolio
        self.mu = np.matrix(df_return.mean(axis = 0).values)
        self.sigma = df_return.cov().values

    def opt_portfolio(self):
        
        w = cv.Variable((self.mu.shape[1],1))
        """
        MV model Variables
        """
        risk = cv.quad_form(w, self.sigma) 
        
        constraints = [cv.sum(w) == 1]
        constraints.append(w >= 0)
        constraints.append(w <= 0.25)

        prob = cv.Problem(cv.Minimize(risk), constraints)
        prob.solve()
        
        weights = np.matrix(w.value).T
        weights = weights/np.sum(weights)
        weights = np.squeeze(np.array(weights))

        return weights