"""
Adaptive Volatility (position sizing)

credit attribution:
    David Varadi
    https://cssanalytics.wordpress.com/2017/11/15/adaptive-volatility/

Aim: Get a better position sizing than  [target_vol / realized_vol_{t-1}]
     (where the realized_vol is calculated over a fixed lookback period, e.g. past 20 days)
     using a more 'adaptive' volatility that varies its lookback period according to market conditions.

The simplest method is to use the R-squared of the regression of prices vs time:
 1. high R-squared indicates a trending market
            -> use short lookback periods to capture sudden changes in volatilities;
 2. low R-squared instead iimplies a rangebound/mean-reverting market
            -> lengthen lookbacks since vol will revert to historical means.

To translate the R_squared value into the alpha for an exponential moving average,
the following exponential function is used (motivation: returns supposed lognormal):

    raw_alpha =  exp[-10. * (1 - R_squared(price vs. time, period=20)]
    alpha = min(raw_alpha, 0.5)

    the 0.5 lower bound effectively  limits the lookback to 3 days, since alpha := 2 / (1 + lookback).

Such a capped aplha is used in an EMA of the squared returns for the past 20 days.

Finally the (theoretical) daily exposure is:

    target_vol / sqrt( EMA_{t-1}(squared rturns, alpha) * 252)

    and target_vol is an annualised target vol, say 20%.

To limit excessive trading, I only rebalace if theoretical exposure changes above a certain threshold (say 5%).

Application hereby:
 long SPY (or similar) with a daily position sizing

A more interesting use of this position sizing scheme is when using algorithms with
long periodical rebalacings, say monthly or quarterly.
"""
import numpy as np
import pandas as pd
from datetime import datetime, timedelta
from scipy.stats import linregress
import decimal as d

class AdaptiveVolatility(QCAlgorithm):

    def __init__(self):
        self.symbols = ['SPY',
                        'TLT'
        ]

        self.back_period = 21 * 3 + 1     # 3 months

        self.vol_period = 21    # days for calc vol
        self.target_vol = 0.2
        self.lev = 1.5          # max lev from ratio targ_vol / real_vol
        
        self.delta = 0.05       # min rebalancing
        
        self.w = 1. / len(self.symbols)
        self.x = np.asarray(range(self.vol_period))
        
        self.SetBenchmark('SPY')
        

    ###################################### 
    def Initialize(self):

        self.SetCash(100000)
        self.SetStartDate(2019,1,1)  # (2006,1,1) 
    #   self.SetEndDate(datetime.now().date() - timedelta(1))
        self.SetBrokerageModel(BrokerageName.InteractiveBrokersBrokerage,
                               AccountType.Margin)

        # register and replace 'tkr symbol' with 'tkr object'
        for i, tkr in enumerate(self.symbols):
            self.symbols[i] = self.AddEquity(tkr, Resolution.Minute).Symbol     # was: .Daily

        # schedule: rebalance
        self.Schedule.On(self.DateRules.EveryDay(self.symbols[0]), 
                         self.TimeRules.AfterMarketOpen(self.symbols[0], -90),
                         Action(self.rebalance))
                         
        # schedule: email for recap at around close
        self.Schedule.On(self.DateRules.EveryDay(self.symbols[0]), 
                         self.TimeRules.BeforeMarketClose(self.symbols[0], 0), Action(self.JustBeforeMarketClose))

        # DEBUG: testing email once
        string_test = "This is a test. \nNo need to to anything" 
        self.Notify.Email("alex.muci@gmail.com", "IB Algo: algo test", string_test)


    ###################################### 
    def rebalance(self):
        
        # get all weights
        try:
            weight, close = self.pos_sizing() 
        except Exception as e:
                self.Notify.Email("alex.muci@gmail.com", "ERROR", "Exception: "  + str(e) )

        
        tot_port = self.Portfolio.TotalPortfolioValue
        
        for tkr in self.symbols:
            
            price = self.Securities[tkr.Value].Price  # in case we move to trade during session
            if price == 0: price = close[tkr.Value]
            
            curr_no_shares = self.Portfolio[tkr.Value].Quantity
            
            # gauge if needs to trade (new weight vs. current one > self.delta)
            curr_weight = curr_no_shares * price / tot_port
            new_weight = weight[tkr.Value]  
            shall_trade = abs(float(new_weight) - float(curr_weight)) > self.delta
            
            if shall_trade: 
                
                # self.SetHoldings(tkr, new_weight)
                
                delta_shares = int(new_weight * tot_port/ price) - curr_no_shares
                self.MarketOnOpenOrder(tkr, delta_shares)
                
                # DEBUG: testing email once             
                _string_trades = "TRADE: tkr: %s -- weight: %.2f (old weight was: %.2f) -- last price: %.2f"  \
                                 %(str(tkr), float(new_weight), float(curr_weight), price)
                self.Log(_string_trades)
                self.Notify.Email("alex.muci@gmail.com", "IB Algo Execution: short vol", _string_trades)


    def pos_sizing(self):

        # get daily returns for period = self.back_period
        prices = self.History(self.symbols, self.back_period, Resolution.Daily)["close"].unstack(level=0)     # .dropna(axis=1)
        daily_rtrn = prices.pct_change().dropna() # or: np.log(price / price.shift(1)).dropna()
        
        pos = {}
        yest_close = {}

        # calculate alpha for EWM
        for tkr in self.symbols:
            
            _rsq = self.rsquared(self.x, np.asarray(prices[tkr.Value])[-self.vol_period:])
                
            alpha_raw = np.exp(-10. * (1. - _rsq))
            alpha_ = min(alpha_raw, 0.5)
               
            vol = daily_rtrn[tkr.Value].ewm(alpha=alpha_).std() # alpha = 2/(span+1) = 1-exp(log(0.5)/halflife)
            ann_vol = vol.tail(1) * np.sqrt(252)
            
            # self.Log("rsqr: %s, alpha_raw: %s, ann_vol = %s" %(str(_rsq), str(alpha_raw), str(ann_vol)) )
            
            pos[tkr.Value] = (self.target_vol / ann_vol).clip(0.0, self.lev)  * self.w  # NB: self.w = 1/no_assets  
            yest_close[tkr.Value] = prices[tkr.Value].values[-1]

        return pos, yest_close

   
    ###################################### 
    def rsquared(self, x, y):
        # slope, intercept, r_value, p_value, std_err
        _, _, r_value, _, _ = linregress(x, y)
        return r_value**2
    
    
    ###################################### ###################################### 
    def OnMarginCallWarning(self):
        margin_msg = "check warning margin call! Fast"
        self.Log(margin_msg)
        self.Notify.Email("alex.muci@gmail.com", "IB Algo: WARNING", margin_msg)

    ######################################    
    def JustBeforeMarketClose(self):
        my_msg = "End of day: %s \nPortfolio value is %.2f and Margin Remaining is: %.2f  (Total Holdings Value: %.2f)" \
                 %( str(self.Time), self.Portfolio.TotalPortfolioValue, self.Portfolio.MarginRemaining, self.Portfolio.TotalHoldingsValue)
        self.Log(my_msg)
        self.Notify.Email("alex.muci@gmail.com", "IB: portfolio and margins at end of day", my_msg)

    ######################################            
    def OnOrderEvent(self, orderEvent):
        order = self.Transactions.GetOrderById(orderEvent.OrderId)
        self.Log("{0}: {1}: {2}".format(self.Time, order.Type, orderEvent))
        
    ###################################### 
    def TimeIs(self, day, hour, minute):
        return self.Time.day == day and self.Time.hour == hour and self.Time.minute == minute