from AlgorithmImports import *

from datetime import timedelta
import numpy as np


class FuturesMomentumEffect(QCAlgorithm):
    '''Basic template algorithm simply initializes the date range and cash'''

    def Initialize(self):
        self.SetStartDate(2010, 1, 1)
        self.SetEndDate(2012, 9, 15)
        # self.SetEndDate(2020, 10, 1)
        self.SetCash(1000000)
        futures_universe = [
            Futures.Grains.Wheat,
            Futures.Grains.HRWWheat,
            Futures.Grains.Corn,
            Futures.Grains.Soybeans,
            Futures.Grains.SoybeanMeal,
            Futures.Grains.SoybeanOil
        ]

        indicator_period = 21
        # lookback_multiplier = 1
        # holding_multiplier = 6
        lookback_multiplier = int(self.GetParameter("lookback-period"))
        holding_multiplier = int(self.GetParameter("holding-period"))

        self.continuous_contracts = {}
        self.logr_indicators = {}
        for i in futures_universe:
            future = self.AddFuture(i, dataNormalizationMode=DataNormalizationMode.BackwardsRatio,
                                    dataMappingMode=DataMappingMode.FirstDayMonth, contractDepthOffset=0,
                                    resolution=Resolution.Daily)
            future.SetFilter(0, 182)
            self.continuous_contracts[future.Symbol] = future
            self.logr_indicators[future.Symbol] = self.LOGR(
                future.Symbol, indicator_period * lookback_multiplier, Resolution.Daily)

        self.nextRebalance = self.Time
        self.Rebalancetime = timedelta(30) * (holding_multiplier-1)

        # set warmup
        self.SetWarmUp(timedelta(indicator_period * lookback_multiplier))

        self.test_tag = f"({lookback_multiplier}-{holding_multiplier})"

        self.most_liquid_contracts = {}
        self.long_holdings = {}
        self.short_holdings = {}

        self.quantity = 10
        self.SetSecurityInitializer(BrokerageModelSecurityInitializer(self.BrokerageModel,
                                                                      FuncSecuritySeeder(self.GetLastKnownPrices)))

    def OnData(self, slice: Slice) -> None:
        if self.IsWarmingUp:
            return

        for changed_event in slice.SymbolChangedEvents.Values:
            old_contract = self.Symbol(changed_event.OldSymbol)
            new_contract = self.Symbol(changed_event.NewSymbol)
            old_month = old_contract.ID.Date.month
            new_month = new_contract.ID.Date.month
            rollover_tag = f"Rollover: {old_contract.Value} -> {new_contract.Value}, ({old_month}) -> ({new_month})"
            liquidate_tag = rollover_tag + \
                f', [Liquidate {old_contract.Value}]'
            new_tag = rollover_tag + f', [Open {new_contract.Value}]'

            if self.Portfolio[old_contract].Invested:
                if self.Portfolio[old_contract].Quantity > 0:
                    self.Liquidate(old_contract, tag=liquidate_tag)
                    self.MarketOrder(new_contract,
                                     self.quantity, tag=new_tag)
                elif self.Portfolio[old_contract].Quantity < 0:
                    self.Liquidate(old_contract, tag=liquidate_tag)
                    self.MarketOrder(new_contract, -
                                     self.quantity, tag=new_tag)


        # if first day of month
        if self.Time < (self.nextRebalance) or slice.FuturesChains.Count == 0:
            return

        invested = [x.Symbol.Value for x in self.Portfolio.Values if x.Invested]
        # get portfolio value
        portfolio_value = self.Portfolio.TotalPortfolioValue
        # get portfolio cash
        portfolio_cash = self.Portfolio.Cash
        # get portfolio pnl
        portfolio_pnl = self.Portfolio.TotalProfit

        self.Liquidate(tag=f'Rebalance - Liquidate ALL - {self.test_tag}')

        highest_logr, lowest_logr = self.Momentum_Signal(self.logr_indicators)



        for continuous_contract_symbol, chain in slice.FuturesChains.items():
            continuous_contract = self.continuous_contracts[continuous_contract_symbol]
            mapped_contract = continuous_contract.Mapped
            mapped_contract_debug = self.Symbol(mapped_contract).Value
            mapping_debug = f'{continuous_contract_symbol} Mapped: {mapped_contract_debug} ({mapped_contract})'

            if highest_logr == continuous_contract_symbol:
                self.MarketOrder(mapped_contract, self.quantity,
                                 tag=f'LONG {mapped_contract_debug} - {self.test_tag}')
            elif lowest_logr == continuous_contract_symbol:
                self.MarketOrder(mapped_contract, -self.quantity,
                                 tag=f'SHORT {mapped_contract_debug} - {self.test_tag}')

        # set next rebalance time
        self.nextRebalance = Expiry.EndOfMonth(self.Time) + self.Rebalancetime

    def Momentum_Signal(self, mom_indicator_dict):
        # get current value for logr indicators
        logr_values = {}
        for i in mom_indicator_dict:
            # check if indicator is ready
            if mom_indicator_dict[i].IsReady:
                logr_values[i] = round(mom_indicator_dict[i].Current.Value, 4)

        # sort logr values
        sorted_logr_values = dict(
            sorted(logr_values.items(), key=lambda x: x[1]))

        # get highest and lowest logr values
        if len(sorted_logr_values) < 2:
            return None, None

        highest_logr = list(sorted_logr_values.keys())[-1]
        lowest_logr = list(sorted_logr_values.keys())[0]

        return highest_logr, lowest_logr