# region imports
from AlgorithmImports import *
from io import StringIO
from typing import List, Dict
from pandas.core.frame import DataFrame
from pandas.core.series import Series
# endregion

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

class SystematicInnovationFactorinStocks(QCAlgorithm):

    _momentum_period: int = 12 * 21
    _momentum_count: int = 5

    _vix_filter_type: VixFilterType = VixFilterType.VIX_RANK
    _vix_filter_value_threshold: float = 0.5

    _safe_asset_ticker: str = 'IEF'

    _trade_safe_asset: bool = True
    _trade_stocks: bool = False

    def initialize(self) -> None:
        self.set_start_date(2013, 1, 1)
        self.set_cash(100_000)

        self.set_brokerage_model(BrokerageName.INTERACTIVE_BROKERS_BROKERAGE, AccountType.MARGIN)
        self.settings.minimum_order_margin_portfolio_percentage = 0
        self.settings.daily_precise_end_time = False

        self._price_data: Dict[Symbol, float] = {}
        self._long: List[Symbol] = []
        self._selected_stock_universe: List[Symbol] = []

        market: Symbol = self.add_equity('SPY', Resolution.DAILY).symbol
        self._safe_asset: Symbol = self.add_equity(self._safe_asset_ticker, Resolution.DAILY).symbol

        file: str = self.Download(f'data.quantpedia.com/backtesting_data/economic/most_innovative_companies.csv')
        self._innovative_companies_df: DataFrame = pd.read_csv(StringIO(file), delimiter=';')
        self._innovative_companies_df.set_index('date', inplace=True)

        tickers: np.ndarray = np.array([list(self._innovative_companies_df[col].values) for col in list(self._innovative_companies_df.columns)])
        self._unique_tickers: List[str] = list(set(tickers.reshape(1, tickers.size)[0]))

        # subscribe to VIX filter asset/s
        if self._vix_filter_type == VixFilterType.VVIX:
            iv: Symbol = self.add_data(CBOE, "VVIX", Resolution.DAILY).symbol
            self._signal_assets = [iv]
        elif self._vix_filter_type == VixFilterType.VIX_RANK:
            iv: Symbol = self.add_data(CBOE, "VIX", Resolution.DAILY).symbol
            self._signal_assets = [iv]
        elif self._vix_filter_type == VixFilterType.VIX_RATIO:
            iv: Symbol = self.add_data(CBOE, "VIX", Resolution.DAILY).symbol
            iv_3m: Symbol = self.add_data(CBOE, "VIX3M", Resolution.DAILY).symbol
            self._signal_assets = [iv, iv_3m]

        self._selection_month: int = 6
        self._selection_flag: bool = False
        self._rebalance_flag: bool = False

        self.universe_settings.resolution = Resolution.DAILY
        self.add_universe(self._fundamental_selection_function)
        self.schedule.on(self.date_rules.month_end(market), self.time_rules.after_market_close(market), self._selection)
        self.schedule.on(self.date_rules.every_day(), self.time_rules.after_market_close(market), self._daily_selection)

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

    def _fundamental_selection_function(self, fundamental: List[Fundamental]) -> List[Symbol]:
        # update the rolling window every day
        for stock in fundamental:
            symbol: Symbol = stock.symbol

            if symbol in self._price_data:
                self._price_data[symbol].add(stock.adjusted_price)
        
        if not self._selection_flag:
            return Universe.UNCHANGED
        self._selection_flag = False
        
        # select new stock universe
        self._long.clear()

        # select only dataset related stocks
        selected: List[Fundamental] = [
            f.symbol for f in fundamental if f.symbol.value in self._unique_tickers
        ]

        # warmup price rolling windows
        for symbol in selected:
            if symbol in self._price_data:
                continue
            
            self._price_data[symbol] = RollingWindow[float](self._momentum_period)

            history: DataFrame = self.history(TradeBar, symbol, self._momentum_period, Resolution.DAILY)
            if history.empty:
                self.log(f"Not enough data for {symbol} yet.")
                continue
            
            closes: Series = history.loc[symbol].close
            for time, close in closes.items():
                self._price_data[symbol].add(close)

        # pick top momentum stocks each year
        if self.time.year in self._innovative_companies_df.index:
            momentum_by_symbol: Dict[Symbol, float] = {
                s : self._price_data[s][0] / self._price_data[s][self._momentum_period - 1] - 1 
                for s in selected 
                if s.value in self._innovative_companies_df.loc[self.time.year].values
                and s in self._price_data
                and self._price_data[s].is_ready
            }
            
            if len(momentum_by_symbol) >= self._momentum_count:
                sorted_by_momentum: List[Symbol] = sorted(momentum_by_symbol, key=momentum_by_symbol.get, reverse=True)
                self._long = sorted_by_momentum[:self._momentum_count]
        
        return self._long

    def on_data(self, slice: Slice) -> None:
        if not self._rebalance_flag:
            return
        self._rebalance_flag = False

        # trade safe asset
        move_to_safe_asset = False

        if self._vix_filter_type == VixFilterType.VVIX:
            vvix: float = self.securities[self._signal_assets[0]].price
            if vvix > self._vix_filter_value_threshold:
                move_to_safe_asset = True
        elif self._vix_filter_type == VixFilterType.VIX_RANK:
            vix_rank: float = self._get_VIX_rank(self._signal_assets[0])
            if vix_rank > self._vix_filter_value_threshold:
                move_to_safe_asset = True
        elif self._vix_filter_type == VixFilterType.VIX_RATIO:
            vix_ratio: float = self.securities[self._signal_assets[0]].price / self.securities[self._signal_assets[1]].price
            if vix_ratio > self._vix_filter_value_threshold:
                move_to_safe_asset = True

        rebalance: bool = False
        long: List[Symbol] = []

        if move_to_safe_asset:
            if self._trade_safe_asset:
                long = [self._safe_asset]
                if not self.portfolio[self._safe_asset].invested:
                    rebalance = True
            else:
                long = []
                if self.portfolio.invested:
                    rebalance = True
        else:
            if self._trade_stocks:
                long = self._long
                if not self.portfolio.invested:
                    rebalance = True
            else:
                long = []
                if self.portfolio.invested:
                    rebalance = True

        # trade execution
        if rebalance:
            targets: List[PortfolioTarget] = []
            for symbol in long:
                if slice.contains_key(symbol) and slice[symbol]:
                    targets.append(PortfolioTarget(symbol, 1. / len(long)))
            
            self.set_holdings(targets, True)
    
    def _daily_selection(self) -> None:
        self._rebalance_flag = True

    def _selection(self) -> None:
        if self.time.month == self._selection_month:
            # safe asset in subscribed and has data
            if self.securities.contains_key(self._safe_asset) and self.securities[self._safe_asset].get_last_data():
                self._selection_flag = True
        
        # self._rebalance_flag = True

    def _get_VIX_rank(self, vix: Symbol, lookback: int = 150) -> float:
        history: DataFrame = self.history(CBOE, vix, lookback, Resolution.DAILY)
        rank: float = ((self.securities[vix].price - min(history["low"])) / (max(history["high"]) - min(history["low"])))
        return rank