# region imports
from AlgorithmImports import *
import pandas as pd
from pandas.core.frame import DataFrame
from typing import List
from traded_strategy import TradedStrategy
# endregion

class MetatronCommodityETFSeasonalityPlusMomentum(QCAlgorithm):

    _notional_value: int = 160_000
    _long_corr_period: int = 250
    _short_corr_period: int = 20
    _asset_count: int = 2
    _offset_months: int = 10
    _trade_exec_minute_offset: int = 15

    _traded_strategy: TradedStrategy = TradedStrategy.FRONT_RUN_SEASONALITY

    def initialize(self) -> None:
        self.set_start_date(2007, 1, 1)
        self.set_cash(self._notional_value)

        leverage: int = 3

        self.set_brokerage_model(BrokerageName.INTERACTIVE_BROKERS_BROKERAGE, AccountType.MARGIN)

        tickers: List[str] = ['DBA', 'DBB', 'DBE', 'DBP']
        self._traded_assets: List[Symbol] = [
            self.add_equity(ticker, Resolution.MINUTE, leverage=leverage).symbol for ticker in tickers
        ]

        self._trade_flag: bool = False
        self.settings.minimum_order_margin_portfolio_percentage = 0.
        self.settings.daily_precise_end_time = False

        self.schedule.on(
            self.date_rules.month_end(self._traded_assets[0]), 
            self.time_rules.before_market_close(self._traded_assets[0], self._trade_exec_minute_offset), 
            self.selection
        )

    def on_data(self, slice: Slice) -> None:
        # monthly rebalance
        if not self._trade_flag:
            return
        self._trade_flag = False
        self.log('New monthly rebalance')

        long: List[str] = []
        short: List[str] = []

        # correlation signal
        history: DataFrame = self.history(
            TradeBar, self._traded_assets, self._long_corr_period, resolution=Resolution.DAILY
        )

        if 'close' not in history.columns:
            self.log('Close column is not present in history dataframe')
            return

        prices: DataFrame = history.close.unstack(level=0)
        monthly_returns: DataFrame = prices.groupby(pd.Grouper(freq='M')).last().pct_change().dropna()

        if len(prices) == self._long_corr_period and len(prices.columns) == len(self._traded_assets):
            self.log("Calculating correlation signal")
            
            observed_performance: DataFrame = monthly_returns[monthly_returns.index.month == (self.time - pd.DateOffset(months=self._offset_months)).month].iloc[0]
            returns: DataFrame = prices.pct_change().dropna()

            long_corr: DataFrame = returns.iloc[-self._long_corr_period:].corr()
            short_corr: DataFrame = returns.iloc[-self._short_corr_period:].corr()

            long_corr_mean: float = long_corr.values[np.triu_indices_from(long_corr.values, 1)].mean()
            short_corr_mean: float = short_corr.values[np.triu_indices_from(short_corr.values, 1)].mean()

            trade_front_run_seasonality: bool = True if short_corr_mean < long_corr_mean else False
            sorted_assets: List[str] = list(
                (prices.iloc[-1] / prices.iloc[0] - 1).sort_values(ascending=trade_front_run_seasonality if self._traded_strategy == TradedStrategy.MOMENTUM else False).index
            )

            momentum_flag: bool = False

            if trade_front_run_seasonality:
                # trade front run seasonality
                if self._traded_strategy in [TradedStrategy.FRONT_RUN_SEASONALITY, TradedStrategy.MOMENTUM_AND_FRONT_RUN_SEASONALITY]:
                    self.log(f"Selecting traded symbols for front run seasonality strategy")

                    above_median_bool: Series = observed_performance > monthly_returns.median()
                    for symbol, is_above in zip(observed_performance.index, above_median_bool):
                        long.append(symbol) if is_above else short.append(symbol)
            else:
                # otherwise, trade momentum
                if self._traded_strategy in [TradedStrategy.MOMENTUM, TradedStrategy.MOMENTUM_AND_FRONT_RUN_SEASONALITY]:
                    self.log(f"Selecting traded symbols for momentum strategy")

                    long = sorted_assets[:self._asset_count]
                    short = sorted_assets[-self._asset_count:]
                    momentum_flag = True
        else:
            self.log('Not enough data for correlation signal')

        # order execution
        self.log('Rebalancing portfolio...')

        for i, portfolio in enumerate([long, short]):
            for symbol in portfolio:
                portfolio_length: int = len(portfolio) if momentum_flag else len(self._traded_assets)
                if slice.contains_key(str(symbol)) and slice[str(symbol)]:
                    q: int = self._notional_value / portfolio_length // slice[symbol].price
                    self.market_order(
                        symbol, 
                        ((-1) ** i) * q,
                    )

    def selection(self) -> None:
        self._trade_flag = True
        self.liquidate()

# region imports
from AlgorithmImports import *
# endregion

class TradedStrategy(Enum):
    MOMENTUM = 1
    FRONT_RUN_SEASONALITY = 2
    MOMENTUM_AND_FRONT_RUN_SEASONALITY = 3