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

from System import *
from QuantConnect import *
from QuantConnect.Algorithm import *
from QuantConnect.Indicators import *
from QuantConnect.Python import PythonQuandl
from QuantConnect.Securities.Equity import EquityExchange

from datetime import datetime, timedelta
import numpy as np
import pandas as pd
from scipy.stats import norm

#from QuantConnect.Data.Custom.Tiingo import *
from QuantConnect.Python import PythonQuandl # quandl data not CLOSE
from QuantConnect.Python import PythonData # custom data
from QuantConnect.Data import SubscriptionDataSource

# Risk Premia RW algorithm
class RPRWAlgorithm(QCAlgorithm):

    def Initialize(self):

        # Initial settings
        self.SetStartDate(2018, 12, 18)
        self.SetEndDate(2019, 8, 30)
        self.SetCash(10000)
        self.MarketAsset = "SPY"
        self.WarmupTime = 310
        self.Window = 300
        
        #parameters
        self.vol_lookback = 90
        self.corr_lookback = 120
        self.formation_periods = np.array([3, 6, 9, 12])*22
        self.z_score_cutoff = 0
        self.momo_multiplier = 0.1        
        
        # these are the growth symbols we'll rotate through
        self.GrowthSymbols = ["VTI",  # Vanguard Total Stock Market ETF
                              "VEA",  # VEA - Vanguard FTSE Developed Markets 
                              "PUTW", # WisdomTree CBOE S&P 500 PutWrite Strategy Fund
                              "TLT",  # iShares 20+ Year Treasury Bond ETF
                              "UST",  # ProShares Ultra 7-10 Year Treasury
                              "VWO",  # iShares MSCI Emerging Markets Indx
                              "VNQI", # VANGUARD INTL E/GLB EX-US RL EST IX
                              "GBTC", #BTC
                              "EMB"]  # iShares J.P. Morgan USD Emerging Markets Bond ETF"
        # these are the safety symbols we go to when things are looking bad for growth
        # this part is not supposed to work
        # I don't know how to open these assets
        #self.SafetySymbols =  "PUTW", # WisdomTree CBOE S&P 500 PutWrite Strategy Fund
        #                      "EMB"]  # iShares J.P. Morgan USD Emerging Markets Bond ETF"
        
        
        #Tiingo.SetAuthCode("49e4ce130ce49c90d9e45193f70723b16c5e4ff7")
        #self.AddData(TiingoDailyData, "PUTW", Resolution.Daily)
        #self.AddData(TiingoDailyData, "EMB", Resolution.Daily)
        
        #self.ticker = "PUTW"
        #self.symbol = self.AddData(TiingoDailyData, self.ticker, Resolution.Daily).Symbol
        self.AddEquity("SPY", Resolution.Daily)
        
        if self.LiveMode:
            self.Debug("Trading Live!")
        
        self.SafetySymbols = []

        
        # all symbols set
        self.AllSymbols = list(set(self.GrowthSymbols) | set(self.SafetySymbols))
        
        # open equity symbols
        for symbol in self.GrowthSymbols:
            self.AddEquity(symbol, Resolution.Daily)

        # this doesn't do anything at the moment. We need to work out how to properly handles these assets
        for symbol in self.SafetySymbols:
            self.AddOption(symbol, Resolution.Daily)

        # wait for warming up
        self.SetWarmUp(self.WarmupTime)
        # schedule the trading function
        #self.Schedule.On(self.DateRules.MonthStart(self.MarketAsset), self.TimeRules.AfterMarketOpen(self.MarketAsset, 10), Action(self.RebalanceAndTrade))
        self.Schedule.On(self.DateRules.EveryDay(), self.TimeRules.AfterMarketOpen(self.MarketAsset, 10), Action(self.RebalanceAndTrade))
        

    def OnData(self, slice):
        if self.LiveMode: self.Debug("Running algorithm!!")
        
        # Make sure all the data we need is in place
        if self.IsWarmingUp: return 

        if not slice.ContainsKey("PUTW"): 
            self.Debug("PUTW not found!!")
            return
        if not slice.ContainsKey("EMB"): 
            self.Debug("EMB not found!!")
            return       
        
        if self.LiveMode: self.Debug("Warm Up Complete Deciding..")

    # trading function
    def RebalanceAndTrade(self):
        
        self.Notify.Email("danielqing@yahoo.com.au", "Current Time" + str(self.Time) , "Portfolio" + str(self.Portfolio))
        # skipping if it is warming up
        if self.IsWarmingUp: return 
        if self.Time.day != 5: return
        self.Liquidate("AUDUSD")

        # creating the pandas DataFrame
        slices = self.History(self.AllSymbols, self.Window, Resolution.Daily)
        slices_df = pd.pivot_table(slices, values = 'close', index='time', columns = 'symbol').reset_index()
        slices_df = slices_df.drop(columns=['time'])
        returns = slices_df.pct_change()
        
        # for debugging
        #self.Debug(self.Time)
        #self.Debug(returns.shape)
        
        # weights calculation
        vol_weights = self.get_srp_weights(returns, self.vol_lookback)
        cor_adjust = self.get_cor_adjustments(returns, self.corr_lookback)
        cor_adjust_weights = self.adjust_weights(vol_weights, cor_adjust, shrinkage=1)
        momo_adjusted_weights = self.get_momo_adjusted_weights(returns, cor_adjust_weights, self.formation_periods, self.z_score_cutoff, self.momo_multiplier)
        # the following should contain asset EMB instead of EEM
        capped_weights = self.cap_allocation_and_rescale(momo_adjusted_weights, ticker="EMB", cap=0.15)
        # the following should VTI and PUTW but I don't know how to handle yet
        final_weights = self.split_allocation(capped_weights, "VTI", "PUTW", ratio=0.5)
        self.Debug(final_weights.shape)
        self.Debug(self.Time)
        self.Debug(final_weights)
        
        # allocating assets
        for i in range(len(final_weights)):
            self.Log("{} : asset {}, allocating {}".format(self.Time, slices_df.columns[i], final_weights[i]))
            self.SetHoldings(slices_df.columns[i], final_weights[i])
        self.Notify.Email("danielqing@yahoo.com.au", "Current Time" + str(self.Time) , "Weight" + str(final_weights))
                    

    def get_srp_weights(self, returns, vol_lookback):
        """
        returns current srp werights given a pandas DataFrame of returns and a vol_lookback period
        """
        n_assets = len(returns.columns)
        vols = returns.iloc[-vol_lookback:, :].apply(lambda x: np.std(x)*np.sqrt(252), axis=0)
        raw_weights = 1/vols
        weights = raw_weights/np.sum(raw_weights)
    
        return weights

    def get_cor_adjustments(self, returns, corr_lookback):
        """
        returns current correlation adjustments given a pandas DataFrame of returns and a corr_lookback period
        """
        cor = returns.iloc[-corr_lookback:, :].corr()
        pairwise_ave_cor = cor.mean(axis=1)
        zscore_pairwise_ave_cor = (pairwise_ave_cor - pairwise_ave_cor.mean())/pairwise_ave_cor.std()
        gauss_scale = 1 - norm.cdf(zscore_pairwise_ave_cor, 0, 1)
        raw_adjustments = gauss_scale/gauss_scale.sum()
        norm_adjustments = raw_adjustments - 1./len(returns.columns)
    
        return norm_adjustments
        
    def adjust_weights(self, vol_weights, corr_adjustments, shrinkage):
        raw_weights = vol_weights * (1 +corr_adjustments * shrinkage)
        adj_weights = raw_weights/raw_weights.sum()
    
        return adj_weights
        

    def get_momo_adjustments(self, returns, formation_period):
        """
        returns current cross-sectional zscore of total return momentum 
        given a pandas DataFrame of returns and formation_period
        """
        synth_prices = (returns+1).cumprod()
        roc = (synth_prices.iloc[-1,:]/synth_prices.iloc[-formation_period-1,:]-1)
        momo_adjustments = (roc - roc.mean())/roc.std()
    
        return momo_adjustments

    def get_sma_slope_adjustments(self, returns, formation_period):
        """
        returns current cross-sectional zscore of slope of moving average 
        given a pandes DataFrame of returns and a formation_period
        """
        synth_prices = (returns+1).cumprod()
        sma = synth_prices.iloc[-formation_period-1:,:].rolling(formation_period).mean()
        sma_slope = (sma.iloc[-1,:]/sma.iloc[-2,:])-1
        momo_adjustments = (sma_slope - sma_slope.mean())/sma_slope.std()
    
        return momo_adjustments

    def adjust_momo_weights(self, base_weights, momo_adjustments, z_score_cutoff, multiplier):
        raw_weights = base_weights * (1 + ((momo_adjustments >= z_score_cutoff) * multiplier))
        adj_weights = raw_weights/raw_weights.sum()
    
        return adj_weights

    def get_momo_adjusted_weights(self, returns, base_weights, formation_periods, z_score_cutoff, multiplier):
        """
        returns current momentum-adjusted weights given a pandes DataFrame of returns and a formation_period
        """
        momo_weights = base_weights

        for period in formation_periods :
            momo_adjustments = self.get_momo_adjustments(returns, period)
            momo_weights = self.adjust_momo_weights(momo_weights, momo_adjustments, z_score_cutoff, multiplier)

        for period in formation_periods :
            momo_adjustments = self.get_sma_slope_adjustments(returns, period)
            momo_weights = self.adjust_momo_weights(momo_weights, momo_adjustments, z_score_cutoff, multiplier)        
        
        return momo_weights
        
    def cap_allocation_and_rescale(self, weights, ticker, cap=0.15):
        """
        cap the allocation into ticker and rescale remaining weights
        """
        if weights[ticker] > cap:
            weights = (1-cap)*weights.drop(ticker)/weights.drop(ticker).sum()
            weights[ticker] = cap
    
        return weights

    def split_allocation(self, weights, ticker, split_ticker, ratio=0.5):
        """
        split the allocation into ticker into ticker and split_ticker according to ratio
        """
        weights[split_ticker] = (1-ratio)*weights[ticker]
        weights[ticker] = ratio*weights[ticker]
    
        #global tradeable_universe
        #if split_ticker not in tradeable_universe:
        #    tradeable_universe.append(split_ticker)
    
        return weights
'''
# Duplicate code?
from System import *
from QuantConnect import *
from QuantConnect.Algorithm import *
from QuantConnect.Indicators import *
from QuantConnect.Python import PythonQuandl
from QuantConnect.Securities.Equity import EquityExchange

from datetime import datetime, timedelta
import numpy as np
import pandas as pd
from scipy.stats import norm

#from QuantConnect.Data.Custom.Tiingo import *
from QuantConnect.Python import PythonQuandl # quandl data not CLOSE
from QuantConnect.Python import PythonData # custom data
from QuantConnect.Data import SubscriptionDataSource

# Risk Premia RW algorithm
class RPRWAlgorithm(QCAlgorithm):

    def Initialize(self):

        # Initial settings
        self.SetStartDate(2019, 1, 1)
        self.SetEndDate(2019, 8, 1)
        self.SetCash(10000)
        self.MarketAsset = "SPY"
        self.WarmupTime = 310
        self.Window = 300
        
        #parameters
        self.vol_lookback = 90
        self.corr_lookback = 120
        self.formation_periods = np.array([3, 6, 9, 12])*22
        self.z_score_cutoff = 0
        self.momo_multiplier = 0.1        
        
        # these are the growth symbols we'll rotate through
        self.GrowthSymbols = ["VTI",  # Vanguard Total Stock Market ETF
                              "SPY",  # SPDR S&P 500 ETF Trust
                              "TLT",  # iShares 20+ Year Treasury Bond ETF
                              "EFA",  # iShares MSCI EAFE ETF
                              "EEM",  # iShares J.P. Morgan USD Emerging Markets Bond ETF 
                              "PUTW", # WisdomTree CBOE S&P 500 PutWrite Strategy Fund
                              "EMB"]  # iShares J.P. Morgan USD Emerging Markets Bond ETF"
        # these are the safety symbols we go to when things are looking bad for growth
        # this part is not supposed to work
        # I don't know how to open these assets
        #self.SafetySymbols =  "PUTW", # WisdomTree CBOE S&P 500 PutWrite Strategy Fund
        #                      "EMB"]  # iShares J.P. Morgan USD Emerging Markets Bond ETF"
        
        
        #Tiingo.SetAuthCode("49e4ce130ce49c90d9e45193f70723b16c5e4ff7")
        #self.AddData(TiingoDailyData, "PUTW", Resolution.Daily)
        #self.AddData(TiingoDailyData, "EMB", Resolution.Daily)
        
        #self.ticker = "PUTW"
        #self.symbol = self.AddData(TiingoDailyData, self.ticker, Resolution.Daily).Symbol
        #self.AddEquity("EMB", Resolution.Daily)
        
        
        
        self.SafetySymbols = []

        
        # all symbols set
        self.AllSymbols = list(set(self.GrowthSymbols) | set(self.SafetySymbols))
        
        # open equity symbols
        for symbol in self.GrowthSymbols:
            self.AddEquity(symbol, Resolution.Daily)

        # this doesn't do anything at the moment. We need to work out how to properly handles these assets
        for symbol in self.SafetySymbols:
            self.AddOption(symbol, Resolution.Daily)

        # wait for warming up
        self.SetWarmUp(self.WarmupTime)
        # schedule the trading function
        self.Schedule.On(self.DateRules.MonthStart(self.MarketAsset), self.TimeRules.AfterMarketOpen(self.MarketAsset, 10), Action(self.RebalanceAndTrade))
        

    def OnData(self, slice):
        if self.LiveMode: self.Debug("Running algorithm!!")
        
        # Make sure all the data we need is in place
        if self.IsWarmingUp: return 

        if not slice.ContainsKey("PUTW"): 
            self.Debug("PUTW not found!!")
            return
        if not slice.ContainsKey("EMB"): 
            self.Debug("EMB not found!!")
            return       
        
        if self.LiveMode: self.Debug("Warm Up Complete Deciding..")

    # trading function
    def RebalanceAndTrade(self):

        # skipping if it is warming up
        if self.IsWarmingUp: return 
    

        # creating the pandas DataFrame
        slices = self.History(self.AllSymbols, self.Window, Resolution.Daily)
        slices_df = pd.pivot_table(slices, values = 'close', index='time', columns = 'symbol').reset_index()
        slices_df = slices_df.drop(columns=['time'])
        returns = slices_df.pct_change()
        
        # for debugging
        #self.Debug(self.Time)
        #self.Debug(returns.shape)
        
        # weights calculation
        vol_weights = self.get_srp_weights(returns, self.vol_lookback)
        cor_adjust = self.get_cor_adjustments(returns, self.corr_lookback)
        cor_adjust_weights = self.adjust_weights(vol_weights, cor_adjust, shrinkage=1)
        momo_adjusted_weights = self.get_momo_adjusted_weights(returns, cor_adjust_weights, self.formation_periods, self.z_score_cutoff, self.momo_multiplier)
        # the following should contain asset EMB instead of EEM
        capped_weights = self.cap_allocation_and_rescale(momo_adjusted_weights, ticker="EMB", cap=0.15)
        # the following should VTI and PUTW but I don't know how to handle yet
        final_weights = self.split_allocation(capped_weights, "VTI", "PUTW", ratio=0.5)
        self.Debug(final_weights.shape)

        # allocating assets
        for i in range(len(final_weights)):
            self.Log("{} : asset {}, allocating {}".format(self.Time, slices_df.columns[i], final_weights[i]))
            self.SetHoldings(slices_df.columns[i], final_weights[i])
                    

    def get_srp_weights(self, returns, vol_lookback):
        """
        returns current srp werights given a pandas DataFrame of returns and a vol_lookback period
        """
        n_assets = len(returns.columns)
        vols = returns.iloc[-vol_lookback:, :].apply(lambda x: np.std(x)*np.sqrt(252), axis=0)
        raw_weights = 1/vols
        weights = raw_weights/np.sum(raw_weights)
    
        return weights

    def get_cor_adjustments(self, returns, corr_lookback):
        """
        returns current correlation adjustments given a pandas DataFrame of returns and a corr_lookback period
        """
        cor = returns.iloc[-corr_lookback:, :].corr()
        pairwise_ave_cor = cor.mean(axis=1)
        zscore_pairwise_ave_cor = (pairwise_ave_cor - pairwise_ave_cor.mean())/pairwise_ave_cor.std()
        gauss_scale = 1 - norm.cdf(zscore_pairwise_ave_cor, 0, 1)
        raw_adjustments = gauss_scale/gauss_scale.sum()
        norm_adjustments = raw_adjustments - 1./len(returns.columns)
    
        return norm_adjustments
        
    def adjust_weights(self, vol_weights, corr_adjustments, shrinkage):
        raw_weights = vol_weights * (1 +corr_adjustments * shrinkage)
        adj_weights = raw_weights/raw_weights.sum()
    
        return adj_weights
        

    def get_momo_adjustments(self, returns, formation_period):
        """
        returns current cross-sectional zscore of total return momentum 
        given a pandas DataFrame of returns and formation_period
        """
        synth_prices = (returns+1).cumprod()
        roc = (synth_prices.iloc[-1,:]/synth_prices.iloc[-formation_period-1,:]-1)
        momo_adjustments = (roc - roc.mean())/roc.std()
    
        return momo_adjustments

    def get_sma_slope_adjustments(self, returns, formation_period):
        """
        returns current cross-sectional zscore of slope of moving average 
        given a pandes DataFrame of returns and a formation_period
        """
        synth_prices = (returns+1).cumprod()
        sma = synth_prices.iloc[-formation_period-1:,:].rolling(formation_period).mean()
        sma_slope = (sma.iloc[-1,:]/sma.iloc[-2,:])-1
        momo_adjustments = (sma_slope - sma_slope.mean())/sma_slope.std()
    
        return momo_adjustments

    def adjust_momo_weights(self, base_weights, momo_adjustments, z_score_cutoff, multiplier):
        raw_weights = base_weights * (1 + ((momo_adjustments >= z_score_cutoff) * multiplier))
        adj_weights = raw_weights/raw_weights.sum()
    
        return adj_weights

    def get_momo_adjusted_weights(self, returns, base_weights, formation_periods, z_score_cutoff, multiplier):
        """
        returns current momentum-adjusted weights given a pandes DataFrame of returns and a formation_period
        """
        momo_weights = base_weights

        for period in formation_periods :
            momo_adjustments = self.get_momo_adjustments(returns, period)
            momo_weights = self.adjust_momo_weights(momo_weights, momo_adjustments, z_score_cutoff, multiplier)

        for period in formation_periods :
            momo_adjustments = self.get_sma_slope_adjustments(returns, period)
            momo_weights = self.adjust_momo_weights(momo_weights, momo_adjustments, z_score_cutoff, multiplier)        
        
        return momo_weights
        
    def cap_allocation_and_rescale(self, weights, ticker, cap=0.15):
        """
        cap the allocation into ticker and rescale remaining weights
        """
        if weights[ticker] > cap:
            weights = (1-cap)*weights.drop(ticker)/weights.drop(ticker).sum()
            weights[ticker] = cap
    
        return weights

    def split_allocation(self, weights, ticker, split_ticker, ratio=0.5):
        """
        split the allocation into ticker into ticker and split_ticker according to ratio
        """
        weights[split_ticker] = (1-ratio)*weights[ticker]
        weights[ticker] = ratio*weights[ticker]
    
        #global tradeable_universe
        #if split_ticker not in tradeable_universe:
        #    tradeable_universe.append(split_ticker)
    
        return weights
        
        
'''