#region imports
from AlgorithmImports import *

from math import floor
#endregion
# https://quantpedia.com/Screener/Details/23


class CommodityMomentumCombinedWithTermStructureAlgorithm(QCAlgorithm):

    def initialize(self):
        self.set_start_date(self.end_date - timedelta(5*365))
        self.set_cash(10_000_000)
        self.settings.seed_initial_prices = True
        self._lookback = int(7.5*60*21) # 21 trading days
        # Add some Futures to the algorithm.
        tickers = [
            Futures.Softs.COCOA,
            Futures.Softs.COFFEE,
            Futures.Grains.CORN,
            Futures.Softs.COTTON_2,
            Futures.Grains.OATS,
            Futures.Softs.ORANGE_JUICE,
            Futures.Grains.SOYBEAN_MEAL,
            Futures.Grains.SOYBEAN_OIL,
            Futures.Grains.SOYBEANS,
            Futures.Softs.SUGAR_11,
            Futures.Grains.WHEAT,
            Futures.Meats.FEEDER_CATTLE,
            Futures.Meats.LEAN_HOGS,
            Futures.Meats.LIVE_CATTLE,
            Futures.Energies.CRUDE_OIL_WTI,
            Futures.Energies.HEATING_OIL,
            Futures.Energies.NATURAL_GAS,
            Futures.Energies.GASOLINE,
            Futures.Metals.MICRO_GOLD,
            Futures.Metals.PALLADIUM,
            Futures.Metals.PLATINUM,
            Futures.Metals.SILVER
        ]
        for ticker in tickers:
            future = self.add_future(ticker)
            future.set_filter(timedelta(0), timedelta(90))
        # Add a warm-up period so the algorithm trades on deployment.
        self.set_warmup(timedelta(45))

    def on_securities_changed(self, changes):
        for security in changes.added_securities:
            roc = self.roc(security, 1)
            security.mean_roc = IndicatorExtensions.sma(roc, self._lookback)
            for bar in self.history[TradeBar](security, self._lookback, Resolution.MINUTE):
                roc.update(bar)

    def on_warmup_finished(self):
        # Add a Scheduled event to rebalance the portfolio monthly.
        time_rule = self.time_rules.after_market_open("SPY", 30)
        self.schedule.on(self.date_rules.month_start("SPY"), time_rule, self._rebalance)
        # Rebalance the portfolio today too.
        if self.live_mode:
            self._rebalance()
        else:
            self.schedule.on(self.date_rules.today, time_rule, self._rebalance)

    def _rebalance(self):
        # Calculate the roll return of each contract.
        roll_return = {}
        for symbol, chain in self.current_slice.future_chains.items():
            if len(chain) < 2:
                continue
            contracts = sorted(chain, key=lambda x: x.expiry)
            expiry_nearest = contracts[0].expiry
            price_nearest = float(contracts[0].last_price) if contracts[0].last_price>0 else 0.5*float(contracts[0].ask_price+contracts[0].bid_price)
            for x in contracts[1:]:
                roll_return[x] = (price_nearest-float(x.last_price))*365 / (x.expiry-expiry_nearest).days
        
        # Split the contracts into high and low roll return buckets.
        sorted_by_roll_return = sorted(roll_return, key=lambda x: roll_return[x], reverse=True)
        tertile = floor(1/3*len(sorted_by_roll_return))
        high = sorted_by_roll_return[:tertile]
        low = sorted_by_roll_return[-tertile:]
        
        # For the contracts with the highest roll returns, select the subset with the greatest mean return.
        mean_return_high = {i: self.securities[i.symbol].mean_roc.current.value for i in high}            
        high_winners = sorted(mean_return_high, key=lambda x: mean_return_high[x], reverse=True)[:int(len(high)*0.5)]
        # For the contracts with the lowest roll returns, select the subset with the lowest mean return.
        mean_return_low = {i: self.securities[i.symbol].mean_roc.current.value for i in high}        
        low_losers = sorted(mean_return_low, key=lambda x: mean_return_low[x], reverse=True)[-int(len(low)*0.5):]

        # Form a long-short portfolio.
        short_weight = -0.05/len(low_losers)
        long_weight = 0.05/len(high_winners)
        targets = [PortfolioTarget(short.symbol, short_weight) for short in low_losers]
        targets += [PortfolioTarget(long_.symbol, long_weight) for long_ in high_winners]
        self.set_holdings(targets, True)