# region imports
from AlgorithmImports import *
# endregion

class TradedStrategy(Enum):
    MOMENTUM = 1
    FRONT_MONTH = 2
    COMBO = 3

# 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 MetatronCommoditySeasonality(QCAlgorithm):

    _notional_value: int = 160_000
    _long_period: int = 250
    _short_period: int = 20
    _asset_count: int = 2
    _offset_months: int = 10
    _trade_exec_minute_offset: int = 15

    _traded_strategy: TradedStrategy = TradedStrategy.COMBO

    def initialize(self) -> None:
        self.set_start_date(2008, 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.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

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

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

        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_period and len(prices.columns) == len(self._traded_assets):
            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_period:].corr()
            short_corr: DataFrame = returns.iloc[-self._short_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_flag: bool = False if short_corr_mean > long_corr_mean else True
            sorted_assets: List[str] = list((prices.iloc[-1] / prices.iloc[0] - 1).sort_values(ascending=trade_flag if self._traded_strategy == TradedStrategy.MOMENTUM else False).index)

            if self._traded_strategy == TradedStrategy.MOMENTUM:
                long = sorted_assets[:self._asset_count]
                short = sorted_assets[-self._asset_count:]
            else:
                if trade_flag:
                    if self._traded_strategy in [TradedStrategy.FRONT_MONTH, TradedStrategy.COMBO]:
                        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:
                    if self._traded_strategy == TradedStrategy.COMBO:
                        long = sorted_assets[:self._asset_count]
                        short = sorted_assets[-self._asset_count:]
        else:
            self.log('Not enough data for correlation signal.')

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

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