#region imports
from AlgorithmImports import *
from collections import deque
#endregion

class ClassicMomentum(PythonIndicator):
    # Momentum measured as p11 / p0 - 1 (excluding last month, or in this case generalized to exclude last 1/12 of period)
    def __init__(self, name, period):
        self.Name = name
        self.WarmUpPeriod = period
        self.Time = datetime.min
        self.Value = 0
        self.queue = deque(maxlen=period)

    def Update(self, input: BaseData) -> bool:
        self.queue.appendleft(input.Value)
        self.Time = input.Time
        self.Value =  (self.queue[-int((1/12)*len(self.queue))] / self.queue[-1]) - 1
        return len(self.queue) == self.queue.maxlen

class SimpleMomentum(PythonIndicator):
    # Momentum measured as p11 / p0 - 1
    def __init__(self, name, period):
        self.Name = name
        self.WarmUpPeriod = period
        self.Time = datetime.min
        self.Value = 0
        self.queue = deque(maxlen=period)

    def Update(self, input: BaseData) -> bool:
        self.queue.appendleft(input.Value)
        self.Time = input.Time
        self.Value =  (self.queue[0] / self.queue[-1]) - 1
        return len(self.queue) == self.queue.maxlen

class ClassicVolatility(PythonIndicator):
    # Simple standard deviation of log returns
    def __init__(self, name, period):
        self.Name = name
        self.WarmUpPeriod = period
        self.Time = datetime.min
        self.Value = 0
        self.window = RollingWindow[float](period)

    def Update(self, input: BaseData) -> bool:
        self.window.Add(input.Value)
        self.Time = input.Time
        self.Value =  np.std([x for x in self.window])
        return self.window.IsReady

#region imports
from AlgorithmImports import *
import riskfolio as rp
#endregion

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

from indicators import SimpleMomentum, ClassicVolatility

class AdaptiveAssetAllocation(QCAlgorithm):

    def Initialize(self):
        self.SetCash(100000)
        self.SetStartDate(2016, 1, 1) # RWX only start end of 12/006
        self.SetEndDate(2023, 10, 1)
        self.Settings.MinimumOrderMarginPortfolioPercentage = 0.  

        # Hyper-parameters #
        ####################
        self.symbols = ['ITOT', 'EWJ', 'VNQ', 'RWX', 'IEF', 'DBC', 'VGK', 'EEM', 'TLT', 'GLD']
        self.allocation_method = 'pos_sizing_momentum'  # change to any of the allocation methods
        self.volatility_lookback = 3*21
        self.momentum_lookback = 6*21
        self.top_n_momentum = 5
        ####################

        self.volatility_ind = {}
        self.momentum_ind = {}

        for i, asset in enumerate(self.symbols):
            self.AddEquity(asset, Resolution.Daily)
            
            mom_indicator = SimpleMomentum(f"mom_{asset}", period=self.momentum_lookback)
            self.RegisterIndicator(asset, 
                mom_indicator, 
                TradeBarConsolidator(1))
            self.momentum_ind[asset] = mom_indicator

            vol_indicator = ClassicVolatility(f"vol_{asset}", period=self.volatility_lookback)
            self.RegisterIndicator(asset, 
                vol_indicator, 
                TradeBarConsolidator(1))
            self.volatility_ind[asset] = vol_indicator

        self.Schedule.On(self.DateRules.MonthStart(self.symbols[0], daysOffset=0),
                         self.TimeRules.BeforeMarketClose(self.symbols[0], minutesBeforeClose=0),
                         self.rebalance)

        self.SetWarmUp(max(self.volatility_lookback, self.momentum_lookback))
        self.weights = None

    def OnData(self, slice: Slice) -> None:

        if self.IsWarmingUp: return

        # Plot the current portfolio weights
        if self.weights is not None:
            for sym in self.weights:
                self.Plot("Portfolio Weights", sym, self.weights[sym])

        # when trading live, reset indicators on splits & dividends
        # if data.Splits.ContainsKey(self.symbol) or data.Dividends.ContainsKey(self.symbol):

    def rebalance(self):
        if self.IsWarmingUp: return
        # get all weights
        self.weights = getattr(self, self.allocation_method)()
        self.Log(self.weights)
        self.SetHoldings([PortfolioTarget(asset, self.weights[asset]) for asset in self.weights], True)
    

    ########################
    ## Allocation methods ##
    ########################

    def pos_sizing_sixtyfortybench(self):
        # Track US 6040 benchmark 
        return {'ITOT': 0.6,
                'TLT': 0.4}

    def pos_sizing_equal_weight(self):
        # Buy all assets, equal weighted
        return {k: ((1)/len(self.symbols))  for k in self.symbols}
    
    def pos_sizing_inverse_volatility(self):
        # Buy all assets, weight them inversely proportional to their volatility
        total_inv_vol = sum([1/self.volatility_ind[sym].Value for sym in self.symbols if self.volatility_ind[sym].IsReady])
        return {sym: math.floor(100* (1/self.volatility_ind[sym].Value) / (total_inv_vol))/100 for sym in self.symbols if self.volatility_ind[sym].IsReady}

    def pos_sizing_momentum(self):
        # Buy the top n momentum assets, equal weighted
        top = self._get_top_momentum(self.top_n_momentum)
        return {sym: 1/self.top_n_momentum if sym in top else 0 for sym in self.symbols }

    def pos_sizing_inverse_volatility_momentum(self):
        # Buy the top n momentum assets, weight them inversely proportional to their volatility
        top = self._get_top_momentum(self.top_n_momentum)

        total_inv_vol = sum([1/self.volatility_ind[sym].Value for sym in top.index])
        return {sym: math.floor(100* (1/self.volatility_ind[sym].Value) / (total_inv_vol))/100 if sym in top else 0 for sym in self.symbols }

    def pos_sizing_minimum_variance(self):
        # Buy all assets, weight them via minvar optimization
        symbols = self.symbols
        ret = self.History(symbols, self.volatility_lookback, Resolution.Daily).unstack(level=0).close.pct_change()[1:]
        port = rp.Portfolio(returns=ret)

        port.assets_stats(method_mu='hist', method_cov='hist', d=0.94)
        w = port.optimization(model='Classic', rm='MV', obj='MinRisk', rf=0, l=0, hist=True)
        allocation = pd.Series([math.floor(x*100)/100 for x in w.values.flatten()], index=ret.columns).to_dict()
        return allocation

    def pos_sizing_momentum_minimum_variance(self):
        # Buy the top n momentum assets, weight them via minvar optimization
        symbols = self._get_top_momentum(self.top_n_momentum)

        ret = self.History(symbols, self.volatility_lookback, Resolution.Daily).unstack(level=0).close.pct_change()[1:]
        port = rp.Portfolio(returns=ret)
        port.assets_stats(method_mu='hist', method_cov='hist', d=0.94)
        w = port.optimization(model='Classic', rm='MV', obj='MinRisk', rf=0, l=0, hist=True)
        allocation = pd.Series([math.floor(x*100)/100 for x in w.values.flatten()], index=ret.columns).to_dict()
        return allocation

    def _get_top_momentum(self, n):
        # returns list of n symbols with biggest momentum
        mom_dict = {k: v for k, v in self.momentum_ind.items() if v.IsReady}
        sorted_signals = pd.Series(mom_dict).sort_values(ascending = False)
        symbols = sorted_signals[:self.top_n_momentum].index.tolist()
        return symbols