import datetime
import pandas as pd
import numpy as np

def cal_volatility(df): 
    prices = np.array(df)
    returns = (prices[1:]-prices[:-1])/prices[:-1]
    volatility = np.std(returns, axis = 0)
    
    return volatility
    
    
def inverse_volatility(df): 
    ''' 
    caculate weight using inverse volatility
    :param df: (DataFrame) datetime, asset_1, asset_2, ....
            This dataframe has been bounded in a range of time
    :return: (Series) weight of corresponding assets
    '''
    vol = cal_volatility(df)
    inverse_vol = 1/vol
    #weight for each asset 
    weights = pd.Series(inverse_vol/sum(inverse_vol), index = df.columns)
    
    return weights

## import to run locally 
# from AlgorithmImports import *
# from clr import AddReference   
# AddReference("System")   
# AddReference("QuantConnect.Algorithm")   
# AddReference("QuantConnect.Common")       
# from System import *   
# from QuantConnect import *   
# from QuantConnect.Algorithm import *

from risk_parity import risk_parity
from inverse_volatility import inverse_volatility
import constant

import random
random.seed(10)

class BareKnuckle(QCAlgorithm):

    def Initialize(self):
        self.SetStartDate(2010, 5, 1)  # Set Start Date
        
        init_cash = 1000000
        self.SetCash(init_cash)  # Set Strategy Cash
        
        #set maximum leverage level 
        self.leverage_level = 2 if self.GetParameter(constant.LEVERAGE) is None else int(self.GetParameter(constant.LEVERAGE))
        
        if self.leverage_level > 1: 
            self.SetBrokerageModel(BrokerageName.InteractiveBrokersBrokerage, AccountType.Margin) #change Brokerage >> get different fee
        
        #max porfolio value (to calculate drawdown)
        self.max_portfolio_value = init_cash

        #param
        stocks = self.GetParameter(constant.STOCK_NAME)
        self.stocks = [
            "SPY",
            "QQQ",
            "TLT"
            ] if stocks is None else str(stocks).split(constant.SPLIT_CHARACTER)

        self.lookback_period = 42 if self.GetParameter(constant.RISK_LOOKBACK) is None else int(self.GetParameter(constant.RISK_LOOKBACK))

        self.optimization = constant.RISK_PARITY if self.GetParameter(constant.OPTIMIZATION) is None else self.GetParameter(constant.OPTIMIZATION)

        #add stock and set normalization mode (use raw price, dividents added directly to portfolio)
        for stock in self.stocks: 
            equity = self.AddEquity(stock, Resolution.Daily)
            self.Securities[stock].SetLeverage(self.leverage_level)
            equity.SetDataNormalizationMode(DataNormalizationMode.Raw) 
        
        
        self.SetWarmUp(self.lookback_period)  # Skip 30 beginning ticks
        
        #choose random asset for scheduling
        asset_for_scheduling = random.choice(self.stocks)
        self.Schedule.On(  # schedule reallocaton every monday
            self.DateRules.Every(DayOfWeek.Monday),
            self.TimeRules.AfterMarketOpen(asset_for_scheduling, 10),
            self.Reallocate
        )

        # Chart - Master Container for the Chart:
        stockPlot = Chart('Weight Plot')
        # On the Trade Plotter Chart we want 3 series: trades and price:
        [stockPlot.AddSeries(Series(stock, SeriesType.Line, 0)) for stock in self.stocks]

        #schedule to add drawdown (monthly). Note: can choose any common etf/stock, not just SPY
        self.Schedule.On(
                        self.DateRules.MonthStart(asset_for_scheduling), 
                        self.TimeRules.AfterMarketOpen(asset_for_scheduling), 
                        self.CalculateDrawdown
                        )
                         
        drawdownMontlyPlot = Chart('Monthly Drawdown')
        drawdownMontlyPlot.AddSeries(Series('Drawdown', SeriesType.Line, 0))
        
    
    def CalculateDrawdown(self): 
        current_portfolio_value = self.Portfolio.TotalPortfolioValue
        #calculate drawdown
        drawdown = round(min(0, current_portfolio_value - self.max_portfolio_value)/self.max_portfolio_value, 2) * 100
        
        #plot new drawdown
        self.Plot('Monthly Drawdown', 'Drawdown', drawdown)
        
        # update max_portfolio_value
        if self.max_portfolio_value < current_portfolio_value: 
            self.max_portfolio_value = current_portfolio_value
        
        
    def Reallocate(self): 
        '''
        Reallocate after a predefine x days/weeks
        
        '''
        #get historical open price of all stocks and calculate weight using risk parity 
        self.stocks_symbol = [self.Symbol(d) for d in self.stocks]
        df = self.History(self.stocks_symbol, self.lookback_period)
        df = df[constant.QC_OPEN_PRICE].unstack(level=0) 
        
        #calculate weight 
        if self.optimization == constant.INVERSE_VOLATILITY:
            weight = inverse_volatility(df)
        else:
            weight = risk_parity(df)

        weight[constant.VISUALIZATION_DATETIME] = str(self.Time)
        self.Debug("weight after risk parity {}".format(weight.to_dict()))
        
        #ploting weight chart 
        for symbol in self.stocks: 
            self.Plot('Weight Plot', symbol, weight[self.Symbol(symbol)])
            
        #rebalance
        portfolio_target = [PortfolioTarget(d, max(weight[d], 0) * self.leverage_level) for d in self.stocks] 
        self.SetHoldings(portfolio_target)
        

    def OnData(self, data):
        '''OnData event is the primary entry point for your algorithm. Each new data point will be pumped in here.
            Arguments:
                data: Slice object keyed by symbol containing the stock data
        '''
        
        # debug: print the day that do not have data 
        for stock in self.stocks: 
            if not data.Bars.ContainsKey(stock):
                self.Debug(self.Time)
                self.Debug("{}, data doesn\'t contain {}".format(self.Time, stock))
                return 
     

# ---- Our constant
SPLIT_CHARACTER=" "
STOCK_NAME="stock"
BOND_NAME="bond"
RISK_LOOKBACK="rlb"
MOMENTUM_LOOKBACK="mlb"
THRES_LOW="tl"
THRES_HIGH="th"
ALPHA_MEDIUM="am"
ALPHA_HIGH="ah"
LEVERAGE="leverage"

OPTIMIZATION="optimization"
RISK_PARITY="rp"
INVERSE_VOLATILITY="iv"

# ---- QuantConnect Pararms
QC_CLOSE_PRICE="close"
QC_LOW_PRICE="low"
QC_HIGH_PRICE="high"
QC_OPEN_PRICE="open"

# ---- Visualization
VISUALIZATION_DATETIME="DateTime"

import datetime
import pandas as pd
import numpy as np
from scipy.optimize import minimize


TOLERANCE = 1e-10

#### EQUALLY WEIGHTED
def equal_weight(df, 
                portfolio_value): 

    weights = [1 / df.shape[1]] * df.shape[1]

    # Convert the weights to a pandas Series
    weights = pd.Series(weights, index=df.columns, name='weight')

    return weights


#### RISK PARITY
def _allocation_risk(weights, covariances):
    portfolio_risk = np.sqrt((weights * covariances * weights.T))[0, 0]

    return portfolio_risk


def _assets_risk_contribution_to_allocation_risk(weights, covariances):
    portfolio_risk = _allocation_risk(weights, covariances)
    assets_risk_contribution = np.multiply(weights.T, covariances * weights.T) / portfolio_risk

    return assets_risk_contribution


def _risk_budget_objective_error(weights, args):
    covariances = args[0]
    assets_risk_budget = args[1]
    weights = np.matrix(weights)
    portfolio_risk = _allocation_risk(weights, covariances)
    assets_risk_contribution = _assets_risk_contribution_to_allocation_risk(weights, covariances)
    assets_risk_target = np.asmatrix(np.multiply(portfolio_risk, assets_risk_budget))
    error = sum(np.square(assets_risk_contribution - assets_risk_target.T))[0, 0]

    return error


def _get_risk_parity_weights(covariances, assets_risk_budget, initial_weights):
    constraints = ({'type': 'eq', 'fun': lambda x: np.sum(x) - 1.0},
                   {'type': 'ineq', 'fun': lambda x: x})

    optimize_result = minimize(fun=_risk_budget_objective_error,
                               x0=initial_weights,
                               args=[covariances, assets_risk_budget],
                               method='SLSQP',
                               constraints=constraints,
                               tol=TOLERANCE,
                               options={'disp': False})
    weights = optimize_result.x

    return weights


def risk_parity(df):
    """
    Execute risk parity measurement
    :param df: (DataFrame) datetime, asset_1, asset_2, ....
        This dataframe has been bounded in a range of time
    :return: (Series) weight of corresponding assets
    """

    covariances = 52.0 * df.asfreq('W-FRI').pct_change().iloc[1:, :].cov().values
    assets_risk_budget = [1 / df.shape[1]] * df.shape[1]
    init_weights = [1 / df.shape[1]] * df.shape[1]
    weights = _get_risk_parity_weights(covariances, assets_risk_budget, init_weights)
    weights = pd.Series(weights, index=df.columns, name='weight')

    return weights