from math import ceil,floor,isnan
from datetime import datetime
import pandas as pd
import numpy as np
from scipy.optimize import minimize
#from Risk.MaximumDrawdownPercentPerSecurity import MaximumDrawdownPercentPerSecurity
# Based on the following discussion - https://www.quantconnect.com/forum/discussion/3216/classical-asset-allocation-with-mean-variance-optimization/p1
# add data normalization
# no SL first to see how it respond 

# Method
# Each month we estimate the optimal mix of assets weights based on information from the prior 252 trading days and use that mix for the next month. For the covariance matrix, we used the historical covariance matrix of returns for the trailing twelve months. 
# As the mean-variance optimization seeks any optimal set of portfolio weights, There is the potential for the portfolio to become quite concentrated at times. To reduce this possibility, we imposed caps (max weights) on assets to enforce greater diversification as indicated in the paper. E.g. impose a cap of 25% for all risky assets and no cap (i.e. a cap of 100%) for all cash-like assets.
# Optimization
#  For optimization, we tried the different methods like maximizing the Sharpe Ratio, maximizing the return given target volatility stays unchanged (We use 15% as the target volatility). 



class AssetAllocationAlgorithm(QCAlgorithm):
    
    # Risky - US Equities (S&P 500) = SPY, VOO, IVV, EEM, JPXN
    # Risky - Crypto eft - BLK, 
    # Risky Foreign Equities (Foreign Large Cap) = VEA, IEFA, EFA
    # Risky Commodities (Broad Diversified) = DBC, PDBC, GSG
    # Mid - REITs = VNQ (North America), VNQI (Global excluding US), REET (Global)
    # growth - VGT, SCHG, VONG
    # Low - Treasury Bond - TLT, IEF
    # Low - Bonds (Total Bond Market) = AGG, BND, BNDX
    
    # mix of 1 high, 1 low and 1 mid - SPY, TLT, BND, VGT, 
    
    def Initialize(self):
        # frame 1
        # self.SetStartDate(2016, 1, 1)
        # self.SetEndDate(2020, 12, 31)
        # #frame 2
        # self.SetStartDate(2011, 1, 1)
        # self.SetEndDate(2015, 12, 31)
        # # #frame 3
        self.SetStartDate(2006, 1, 1)
        self.SetEndDate(2010, 12, 31)
        # #frame 4
        # self.SetStartDate(2001, 1, 1)
        # self.SetEndDate(2005, 12, 31)
        #11 year test
        # self.SetStartDate(2010, 1, 1)
        # self.SetEndDate(2020, 12, 31)
        self.SetCash(100000)             # Set Strategy Cash
        self.SetSecurityInitializer(lambda x: x.SetDataNormalizationMode(DataNormalizationMode.TotalReturn))
        
        #note without data normalization - it will be factoring splits and dividends 
        # DataNormalizationMode.Adjusted //Factoring in splits and dividends, default setting if no normalization
        # .SplitAdjusted // Just factoring splits, paying dividends as cash
        # .TotalReturn //Adding dividends to asset price
        # .Raw // Price as raw, dividends paid as cash, quantity adjusted on splits


        # tickers = ["SPY", "TLT", "BND", "VGT"] #yl 
        # tickers = ["SPY","QQQ","DIA","IWM","TLT","GLD"] #JW 156
        # tickers = ["SPY", "VOO", "IVV", "AGG", "BND", "BNDX", 
        #                 "VNQ", "VNQI", "REET", "VEA", "IEFA", "EFA", 
        #                 "DBC", "PDBC", "GSG"] #evelyn
        # tickers = ["SPY", "IVV", "VTI", "VOO", "QQQ", "MLPY", "VEA", "IEFA"] #top 8 by market cap
        # tickers = ["IEF", "TLT", "SPY", "EFA", "EEM", "JPXN", "VGT"] #orig
        # tickers = ["SPY", "IEFA", "EFA", "MLPY", "IAU", "VWO"] #mix1 56
        # tickers = ["SPY", "QQQ", "VTI", "VO", "VB", "CWB" ] #97
        # tickers = ["SPY", "ONEQ", "QQQ", "GLD", "IWM", "DIA"] #139%
        tickers = ["SPY", "QQQ", "TLT", "GLD"] #gives 176%
         
        
        
        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))
        
        
    ## Risk model - additional parameter
        # stopRisk = self.GetParameter("stopRisk")
        # if stopRisk is None:
        #     stopRisk = 0.16 
        # self.SetRiskManagement(TrailingStopRiskManagementModel(float(stopRisk)))


    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.15   #orig is 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