# https://github.com/QuantConnect/Lean/blob/master/Algorithm.Framework/Risk/MaximumDrawdownPercentPortfolio.py
# QUANTCONNECT.COM - Democratizing Finance, Empowering Individuals.
# Lean Algorithmic Trading Engine v2.0. Copyright 2014 QuantConnect Corporation.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

from clr import AddReference
AddReference("System")
AddReference("QuantConnect.Common")
AddReference("QuantConnect.Algorithm")
AddReference("QuantConnect.Algorithm.Framework")

from QuantConnect import *
from QuantConnect.Algorithm import *
from QuantConnect.Algorithm.Framework import *
from QuantConnect.Algorithm.Framework.Portfolio import PortfolioTarget
from QuantConnect.Algorithm.Framework.Risk import RiskManagementModel

class MaximumDrawdownPercentPortfolio(RiskManagementModel):
    '''Provides an implementation of IRiskManagementModel that limits the drawdown of the portfolio to the specified percentage.'''

    def __init__(self, maximumDrawdownPercent = 0.2, isTrailing = False):
        '''Initializes a new instance of the MaximumDrawdownPercentPortfolio class
        Args:
            maximumDrawdownPercent: The maximum percentage drawdown allowed for algorithm portfolio compared with starting value, defaults to 5% drawdown</param>
            isTrailing: If "false", the drawdown will be relative to the starting value of the portfolio.
                        If "true", the drawdown will be relative the last maximum portfolio value'''
        self.maximumDrawdownPercent = -abs(maximumDrawdownPercent)
        self.isTrailing = isTrailing
        self.initialised = False
        self.portfolioHigh = 0;

    def ManageRisk(self, algorithm, targets):
        '''Manages the algorithm's risk at each time step
        Args:
            algorithm: The algorithm instance
            targets: The current portfolio targets to be assessed for risk'''
        currentValue = algorithm.Portfolio.TotalPortfolioValue

        if not self.initialised:
            self.portfolioHigh = currentValue   # Set initial portfolio value
            self.initialised = True

        # Update trailing high value if in trailing mode
        if self.isTrailing and self.portfolioHigh < currentValue:
            self.portfolioHigh = currentValue
            return []   # return if new high reached

        pnl = self.GetTotalDrawdownPercent(currentValue)
        if pnl < self.maximumDrawdownPercent and len(targets) != 0:
            self.initialised = False # reset the trailing high value for restart investing on next rebalcing period
            return [ PortfolioTarget(target.Symbol, 0) for target in targets ]

        return []

    def GetTotalDrawdownPercent(self, currentValue):
        return (float(currentValue) / float(self.portfolioHigh)) - 1.0

from clr import AddReference
AddReference("System")
AddReference("QuantConnect.Algorithm")
AddReference("QuantConnect.Common")

from System import *
from QuantConnect import *
from QuantConnect.Algorithm import *
from QuantConnect.Orders import *
from QuantConnect.Orders.Fees import *
from QuantConnect.Securities import *
from QuantConnect.Orders.Fills import *
from QuantConnect.Brokerages import *
# from QuantConnect.Algorithm.Framework.Risk import *
from QuantConnect.Data.Custom.USTreasury import *

from sklearn.linear_model import LinearRegression
import numpy as np
import pandas as pd
import math 
import statsmodels.api as sm
from statsmodels.tsa.stattools import coint, adfuller
from datetime import datetime

# from VolatilityModel import CustomVolatilityModel
from RiskManagement import MaximumDrawdownPercentPortfolio

# Compute the Value-at-Risk and Tail VaR for Pairs-Trading Portfolio
class RiskManagementForPairTrading(QCAlgorithm):
    def Initialize(self):
        # Backtest between 7/1/2017 and 7/1/2018
        self.SetStartDate(2017, 7, 1)      
        self.SetEndDate(2018,7,1)
        self.cash = 1000000  
        self.SetCash(self.cash)
        self.enter = 1
        self.exit = 0
        
        self.betarange = 0.25
        self.lookback = 30
        self.resolution = Resolution.Daily
        
        # Pairs ----------------------------------------------------------------
        self.pairs =['MS','GS']  # 0.25 order matters since beta -> 1 is much more profitable
        self.symbols =[]
        self.securities = []
        for ticker in self.pairs:
            security = self.AddEquity(ticker, Resolution.Hour)
            self.symbols.append(security.Symbol)
            self.securities.append(security)
        
        # Cash Buffer ------------------------------------------------------------
        # Adjust the cash buffer from the default 2.5% to 5%
        self.Settings.FreePortfolioValuePercentage = 45
        self.buffer = 0.4  
        
        # Brokerage ------------------------------------------------------------
        self.SetBrokerageModel(BrokerageName.InteractiveBrokersBrokerage, AccountType.Margin)
        
        # Risk Management: Maximum Drawdown ------------------------------------
        # self.AddRiskManagement( MaximumDrawdownPercentPortfolio() ) 
        # stop-loss
        self.stop = 4 
        # daily-pnl-volatility
        self.port_daily_pnl = []
        # sharpe ratio
        self.rate_of_return = []    
        self.daily_performance = []
        # valu-at-risk
        self.var_limit = 5000
        self.ci = 0.95
        
        
        # Execution ------------------------------------------------------------
        self.SetExecution(ImmediateExecutionModel())       
        
        # Risk Free Rate ------------------------------------------------------------
        self.yield_curve = self.AddData(USTreasuryYieldCurveRate, "USTYCR").Symbol
        self.rates = []
        
        # Warm Up --------------------------------------------------------------        
        self.SetWarmUp(timedelta(self.lookback))
        
        #  ------------------------------------------------------------
        self.yesterday_total_profit = 0
        self.yesterday_total_fees = 0
        self.yesterday_total_unrealized_profit = 0       
        self.yesterday_total_portfolio_value = self.cash
    
    def risk_management_sharpe_ratio(self,rates, c):
        if (self.sharpe_ratio_calc(rates) <= c):
                self.Liquidate()
    
    # Calculate annualized sharpe ratio
    def sharpe_ratio_calc(self, rates): 
        sharpe_ratio =( self.annualized_rate_of_return_calc() - rates) /self.annualized_port_std_calc()
        self.Log(f"SR: {sharpe_ratio}")    
        return sharpe_ratio
    
    # Calculate annualized rate of return 
    def annualized_rate_of_return_calc(self):
        annualized_rate_of_return = Math.Pow(np.mean(self.daily_performance) + 1, 252) - 1
        # annualized_rate_of_return = Math.Pow(np.mean(self.rate_of_return[-self.lookback:]) + 1, 252) - 1
        self.Log(f"annualized ror: {annualized_rate_of_return}")
        return annualized_rate_of_return 
    
    # Calculate annualted portfolio standard deviation 
    def annualized_port_std_calc(self):
        annualized_port_std =  np.std(self.daily_performance) * math.sqrt(252)
        self.Log(f"std: {annualized_port_std}")
        return annualized_port_std
    
        
    def OnData(self,data):
        if self.IsWarmingUp: return
        if data.ContainsKey(self.yield_curve):
            rate = data[self.yield_curve]
            if rate.OneYear:
                # self.Log(f"one year {rate.OneYear}")
                self.rates.append(rate.OneYear/100.0)
            
            self.risk_free_rate = np.mean(self.rates[-self.lookback:])
            # self.Log(f" one year {self.risk_free_rate}")
            
        if self.Time.hour == 10:
            if not data.ContainsKey(self.pairs[0]): return
            if not data.ContainsKey(self.pairs[1]): return
            
            # self.rate_of_return.append(self.Portfolio.TotalPortfolioValue/self.cash - 1 )
            [zscore,self.beta,self.adf] = self.port_check(self.symbols[0], self.symbols[1])
        
            
            ticker1_holdings = self.Portfolio[self.pairs[0]]
            self.equity = self.Portfolio.TotalPortfolioValue
    
    
            if ticker1_holdings.Invested:
                if (ticker1_holdings.IsShort and zscore >= self.exit) or \
                    (ticker1_holdings.IsLong and zscore <= self.exit):
                        self.Liquidate(self.pairs[0])
                        self.Liquidate(self.pairs[1])
                
                # Risk Management: Sharpe Ratio --------------------------------
                # Sharpe Ratio > 0.5
                self.risk_management_sharpe_ratio(self.risk_free_rate, 0.5)

                     
                
            elif [self.beta >= 1 - self.betarange and self.beta <= 1+self.betarange]:
                # If portfolio is not invested, then equity = cash = margin = buying power when leverage is 2 
                wt1 = self.beta/(1+self.beta)
                wt2 = 1/(1+self.beta)
                C =  self.equity *2 *(1-self.buffer)
                # self.Debug(str(data[self.symbols[0]].Price) + ", "+ str(self.Portfolio[self.pairs[0]].Price)) 
    
                price1 = data[self.symbols[0]].Price
                price2 = data[self.symbols[1]].Price
                pos1 = round(wt1 * C / price1)
                pos2 = round(wt2 * C / price2)
                
                # Buy ticker1, Sell ticker2
                if zscore > self.enter:
                    self.MarketOrder(self.symbols[0], pos1)
                    self.MarketOrder(self.symbols[1], -pos2)
                    # self.Log(f"adf {self.adf}")
                # Sell ticker1, Buy ticker2
                if zscore < - self.enter:
                    self.MarketOrder(self.symbols[0], -pos1)
                    self.MarketOrder(self.symbols[1], pos2)
                    # self.Log(f"adf {self.adf}")
            else:
                pass
    
    def OnEndOfDay(self, symbol):
        self.daily_performance.append((self.Portfolio.TotalPortfolioValue - self.yesterday_total_portfolio_value)/self.yesterday_total_portfolio_value  )
        self.yesterday_total_portfolio_value = self.Portfolio.TotalPortfolioValue           
    
    def port_check(self, symbol1, symbol2):
        x = np.log(self.History(symbol1, self.lookback, self.resolution).open.values)
        y = np.log(self.History(symbol2, self.lookback,  self.resolution).open.values)
        
        regr = LinearRegression()
        x_constant = np.column_stack([np.ones(len(x)), x])
        regr.fit(x_constant, y)
        beta = regr.coef_[1]  # weight of ticker1: ticker2 = beta : 1
        alpha = regr.intercept_
        res = y - x*beta - alpha
        mean = np.mean(res)
        std = np.std(res)
        zscore = (res[-1] - mean)/std 
        adf = adfuller(res)
    
        return [zscore, beta, adf]