# region imports
from AlgorithmImports import *
# endregion

PRINT_YTD = False
PRINT_TRADES = True
PRINT_BENCHMARK = True

#region imports
from AlgorithmImports import *
from System.Drawing import Color
from pandas.core.frame import DataFrame
from enum import Enum
import charting_setup
from vix_filter import VixFilterType, subscribe_assets, get_VIX_rank, move_to_safe_asset
#endregion

class BBZone(Enum):
    NONE = 0
    ZONE_1 = 1
    ZONE_2 = 2
    ZONE_3 = 3
    ZONE_4 = 4

class DirectionAsset(Enum):
    MARKET = 1
    VIX = 2
    MARKET_AND_VIX = 3

class BollingerBands(QCAlgorithm):

    # directional filter
    _direction_asset: DirectionAsset = DirectionAsset.VIX

    # VIX Filter
    _vix_filter_type: VixFilterType = VixFilterType.VIX_RATIO
    _vix_filter_value_threshold: float = 1.
    
    _safe_asset_ticker: str = 'IEF'

    def initialize(self):
        self.set_start_date(2010, 1, 1)

        self._init_cash: int = 20_000
        self._start_of_year_portfolio_value = self._init_cash
        self.set_cash('USD', self._init_cash)

        tickers: List[str] = [
            # symbol 1, symbol 2, symbol 3, symbol 4
            # 'SPY', 'SPY', 'SPY', 'SPY', 
            'SPY', 'IEF', 'EEM', 'FXC', 

            # symbol 5, symbol 6, symbol 7, symbol 8
            # 'SPY', 'SPY', 'SPY', 'SPY', 
            'BIL', 'IWM', 'QQQ', 'UUP'
        ]
        self._symbols: List[Symbol] = [
            self.add_equity(t, Resolution.DAILY).symbol for t in tickers
        ]

        if self._direction_asset == DirectionAsset.MARKET_AND_VIX:
            self._traded_symbol_by_zone: Dict[Dict[int, Symbol]] = {
                BBZone.ZONE_1 : { (-1, 1): self._symbols[0], (1, -1): self._symbols[4] },
                BBZone.ZONE_2 : { (-1, 1): self._symbols[1], (1, -1): self._symbols[5] },
                BBZone.ZONE_3 : { (-1, 1): self._symbols[2], (1, -1): self._symbols[6] },
                BBZone.ZONE_4 : { (-1, 1): self._symbols[3], (1, -1): self._symbols[7] },
            }
        else:
            self._traded_symbol_by_zone: Dict[Dict[int, Symbol]] = {
                BBZone.ZONE_1 : { -1: self._symbols[0], 1: self._symbols[4] },
                BBZone.ZONE_2 : { -1: self._symbols[1], 1: self._symbols[5] },
                BBZone.ZONE_3 : { -1: self._symbols[2], 1: self._symbols[6] },
                BBZone.ZONE_4 : { -1: self._symbols[3], 1: self._symbols[7] },
            }

        period: int = 2
        self._iv: Symbol = self.add_data(CBOE, 'VIX', Resolution.DAILY).symbol
        self._market: Symbol = self.add_equity('SPY', Resolution.DAILY).symbol
        self._safe_asset: Symbol = self.add_equity(self._safe_asset_ticker, Resolution.DAILY).symbol

        self._signal_assets: List[Symbol] = subscribe_assets(self, self._vix_filter_type)

        # bollinger bands setup
        bb_period: int = 20
        bb_std: int = 2
        self._bb: BollingerBands = self.BB(self._iv, bb_period, bb_std)

        self._historical_prices: Dict[Symbol, RollingWindow] = {
            self._iv : RollingWindow[float](period),
            self._market : RollingWindow[float](period),
        }

        # benchmark
        self._benchmark_values: List[float] = []

        # charting
        # BB chart
        self._bb_chart_name: str = "Bollinger Bands Signals"
        stock_plot: Chart = Chart(self._bb_chart_name)

        # NOTE there's a limited number of chart series we can print
        if charting_setup.PRINT_TRADES:
            # [:4] is temporary workaround to display only trades for the first 4 original symbols
            for symbol in list(set(self._symbols))[:4]:
                stock_plot.add_series(Series(f'Buy {symbol.value}', SeriesType.SCATTER, '', 
                                    Color.GREEN, ScatterMarkerSymbol.TRIANGLE))
                stock_plot.add_series(Series(f'Sell {symbol.value}', SeriesType.SCATTER, '', 
                                    Color.RED, ScatterMarkerSymbol.TRIANGLE_DOWN))
                stock_plot.add_series(Series(f'Liquidate {symbol.value}', SeriesType.SCATTER, '', 
                                    Color.BLUE, ScatterMarkerSymbol.DIAMOND))
                self.add_chart(stock_plot)

        # YTD chart
        if charting_setup.PRINT_YTD:
            self._recent_year: int = -1
            self._ytd_chart_name: str = "YTD Performance [%]"
            self._ytd_chart: Chart = Chart(self._ytd_chart_name)
            self._ytd_chart.add_series(Series("YTD", SeriesType.LINE, ''))
            self.add_chart(self._ytd_chart)

        self.set_warm_up(bb_period, Resolution.DAILY)
        self.settings.minimum_order_margin_portfolio_percentage = 0.
        self.settings.daily_precise_end_time = False

    def on_securities_changed(self, changes: SecurityChanges) -> None:
        for security in changes.added_securities:
            security.set_leverage(3)

    def on_end_of_day(self) -> None:
        # print benchmark in main equity plot
        if charting_setup.PRINT_BENCHMARK:
            mkt_price_df: DataFrame = self.History(self._market, 2, Resolution.DAILY)
            if not mkt_price_df.empty:
                benchmark_price: float = mkt_price_df['close'].unstack(level= 0).iloc[-1]
                if len(self._benchmark_values) == 2:
                    self._benchmark_values[-1] = benchmark_price
                    benchmark_perf: float = self._init_cash * (self._benchmark_values[-1] / self._benchmark_values[0])
                    self.plot('Strategy Equity', self._market, benchmark_perf)
                else:
                    self._benchmark_values.append(benchmark_price)

        # print YTD performance
        if charting_setup.PRINT_YTD:
            if self.time.year != self._recent_year:
                self._recent_year = self.time.year
                self._start_of_year_portfolio_value = self.portfolio.total_portfolio_value
            
            # calculate the YTD performance
            ytd_performance: float = ((self.portfolio.total_portfolio_value / self._start_of_year_portfolio_value) - 1) * 100
            self.Plot(self._ytd_chart_name, "YTD [%]", ytd_performance)

    def on_data(self, slice: Slice) -> None:
        if all(slice.contains_key(s) and slice[s] for s in list(self._historical_prices.keys())) and \
            all(slice.contains_key(s) and slice[s] for s in self._signal_assets) and \
            all(slice.contains_key(s) and slice[s] for s in self._symbols):

            # store prices
            for s in list(self._historical_prices.keys()):
                self._historical_prices[s].add(slice[s].value)
            
            if self.is_warming_up: 
                return
            if not self._bb.is_ready:
                return
            
            iv_price: float = slice[self._iv].value
            market_price: float = slice[self._market].value

            if charting_setup.PRINT_TRADES:
                self.plot(self._bb_chart_name, "Market Price", market_price)
                self.plot(self._bb_chart_name, "IV Price", iv_price)
                self.plot(self._bb_chart_name, "MiddleBand", self._bb.middle_band.current.value)
                self.plot(self._bb_chart_name, "UpperBand", self._bb.upper_band.current.value)
                self.plot(self._bb_chart_name, "LowerBand", self._bb.lower_band.current.value)
            
            if move_to_safe_asset(
                self, 
                self._vix_filter_type, 
                self._vix_filter_value_threshold, 
                self._signal_assets
            ):
                symbol_to_invest: Symbol|None = self._safe_asset
            else:
                bb_zone: BBZone = self.identify_bb_zone(iv_price, self._bb)
                
                if self._direction_asset == DirectionAsset.VIX:
                    direction: int = 1 if self._historical_prices[self._iv][0] > self._historical_prices[self._iv][1] else -1
                    symbol_to_invest: Symbol|None = self._traded_symbol_by_zone[bb_zone][direction] if direction in self._traded_symbol_by_zone[bb_zone] else None
                elif self._direction_asset == DirectionAsset.MARKET:
                    direction: int = 1 if self._historical_prices[self._market][0] > self._historical_prices[self._market][1] else -1
                    symbol_to_invest: Symbol|None = self._traded_symbol_by_zone[bb_zone][direction] if direction in self._traded_symbol_by_zone[bb_zone] else None
                elif self._direction_asset == DirectionAsset.MARKET_AND_VIX:
                    iv_direction: int = 1 if self._historical_prices[self._iv][0] > self._historical_prices[self._iv][1] else -1
                    market_direction: int = 1 if self._historical_prices[self._market][0] > self._historical_prices[self._market][1] else -1
                    symbol_to_invest: Symbol|None = self._traded_symbol_by_zone[bb_zone][(iv_direction, market_direction)] if (iv_direction, market_direction) in self._traded_symbol_by_zone[bb_zone] else None

            # rebalance
            held_symbols: List[Symbol] = [
                x.key for x in self.portfolio if self.portfolio[x.key].invested
            ]

            if symbol_to_invest is None:
                # hold current holdings
                pass
                # for symbol in held_symbols:
                #     self.trade_and_markup_chart(symbol, self._bb_chart_name, 0., slice[self._iv].value)
            else:
                if symbol_to_invest is not None and symbol_to_invest not in held_symbols:
                    for s in held_symbols:
                        self.trade_and_markup_chart(s, self._bb_chart_name, 0., iv_price)

                    self.trade_and_markup_chart(symbol_to_invest, self._bb_chart_name, 1., iv_price)
        
    def trade_and_markup_chart(
        self,
        symbol: Symbol,
        chart_name: str,
        weight: float, 
        chart_value: float) -> None:
        
        self.set_holdings(symbol, weight)

        if charting_setup.PRINT_TRADES:
            self.plot(
                chart_name, 
                f"Buy {symbol.value}" if weight == 1. else (f"Sell {symbol.value}" if weight == -1. else (f"Liquidate {symbol.value}" if weight == 0. else f"Rebalance {symbol.value}")), 
                chart_value
            )

    def identify_bb_zone(self, price: float, bb: BollingerBands) -> BBZone:
        if price >= bb.upper_band.current.value:
            return BBZone.ZONE_1
        elif price < bb.upper_band.current.value and price >= bb.middle_band.current.value:
            return BBZone.ZONE_2
        elif price < bb.middle_band.current.value and price >= bb.lower_band.current.value:
            return BBZone.ZONE_3
        elif price < bb.lower_band.current.value:
            return BBZone.ZONE_4

# region imports
from AlgorithmImports import *
from enum import Enum
# endregion

class VixFilterType(Enum):
    VVIX = 1
    VIX_RANK = 2
    VIX_RATIO = 3

def subscribe_assets(
    algo: QCAlgorithm, 
    vix_filter_type: VixFilterType
) -> List[Symbol]:

    if vix_filter_type == VixFilterType.VVIX:
        iv: Symbol = algo.add_data(CBOE, "VVIX", Resolution.DAILY).symbol
        signal_assets = [iv]
    elif vix_filter_type == VixFilterType.VIX_RANK:
        iv: Symbol = algo.add_data(CBOE, "VIX", Resolution.DAILY).symbol
        signal_assets = [iv]
    elif vix_filter_type == VixFilterType.VIX_RATIO:
        iv: Symbol = algo.add_data(CBOE, "VIX", Resolution.DAILY).symbol
        iv_3m: Symbol = algo.add_data(CBOE, "VIX3M", Resolution.DAILY).symbol
        signal_assets = [iv, iv_3m]
    
    return signal_assets

def move_to_safe_asset(
    algo: QCAlgorithm,
    vix_filter_type: VixFilterType,
    vix_filter_value_threshold: float,
    signal_assets: List[Symbol]
) -> bool:

    result: bool = False

    if vix_filter_type == VixFilterType.VVIX:
        vvix: float = algo.securities[signal_assets[0]].price
        if vvix > vix_filter_value_threshold:
            result = True
    elif vix_filter_type == VixFilterType.VIX_RANK:
        vix_rank: float = get_VIX_rank(algo, signal_assets[0])
        if vix_rank > vix_filter_value_threshold:
            result = True
    elif vix_filter_type == VixFilterType.VIX_RATIO:
        vix_ratio: float = algo.securities[signal_assets[0]].price / algo.securities[signal_assets[1]].price
        if vix_ratio > vix_filter_value_threshold:
            result = True

    return result

def get_VIX_rank(
    algo: QCAlgorithm,
    vix: Symbol, 
    lookback: int = 150
) -> float:

    history: DataFrame = algo.history(
        CBOE, vix, lookback, Resolution.DAILY
    )
    
    rank: float = ((algo.securities[vix].price - min(history["low"])) / (max(history["high"]) - min(history["low"])))
    return rank