from math import ceil,floor,isnan
from datetime import datetime
import pandas as pd
import numpy as np
from scipy.optimize import minimize



class AssetAllocationAlgorithm(QCAlgorithm):
    
    def Initialize(self):
        self.SetStartDate(2007, 01, 01)  #Set Start Date
        self.SetEndDate(datetime.now())    #Set End Date
        self.SetCash(100000)            #Set Strategy Cash

        tickers = ["IEF", "TLT", "SPY", "EFA", "EEM", "JPXN", "VGT"]
        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()) 
        daily_return = (df_price / df_price.shift(1)).dropna()
            
        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

    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.target_vol = 0.05

    def annual_port_return(self, weights):
        # calculate the annual return of portfolio
        return np.sum(self.daily_return.mean() * weights) * 252

    def annual_port_vol(self, weights):
        # calculate the annual volatility of portfolio
        return np.sqrt(np.dot(weights.T, np.dot(self.daily_return.cov() * 252, weights)))

    def min_func(self, weights):
        
        # method 1: maximize sharp ratio
        # return - self.annual_port_return(weights) / self.annual_port_vol(weights)
        
        # method 2: maximize the return with target volatility 
        return - self.annual_port_return(weights) / self.target_vol

    def opt_portfolio(self):
        # maximize the sharpe ratio to find the optimal weights
        cons = ({'type': 'eq', 'fun': lambda x: np.sum(x) - 1})
        bnds = tuple((0, 1) for x in range(2)) + tuple((0, 0.25) for x in range(self.n - 2))
        opt = minimize(self.min_func,                               # object function
                       np.array(self.n * [1. / self.n]),            # initial value
                       method='SLSQP',                              # optimization method
                       bounds=bnds,                                 # bounds for variables 
                       constraints=cons)                            # constraint conditions
                      
        opt_weights = opt['x']
 
        return opt_weights