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(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.Value] = future
            self.logr_indicators[future.Symbol.Value] = 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.Log(
            f"Testing - lookback period: {lookback_multiplier} months, holding period: {holding_multiplier} months")
        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)

                self.Log(rollover_tag)
            # else:
            #     # raise error
            #     self.Log(rollover_tag + ' --- SKIPPED: old contract not invested')

        # if first day of month
        if self.Time < (self.nextRebalance):
            return

        self.Log('---------------------')
        self.Log(f"Rebalancing at {self.Time}")
        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.Log(
            f"Portfolio value: {portfolio_value}, profit: {portfolio_pnl}, cash: {portfolio_cash}")
        self.Log(f"Current holdings: {invested}")
        self.Log("Liquidating all positions...")
        self.Liquidate(tag=f'Rebalance - Liquidate ALL - {self.test_tag}')

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

        # highest_logr_mapped = self.continuous_contracts[highest_logr].Mapped
        # lowest_logr_mapped = self.continuous_contracts[lowest_logr].Mapped
        # highest_logr_mapped_debug = self.Symbol(
        #     highest_logr_mapped).Value
        # lowest_logr_mapped_debug = self.Symbol(
        #     lowest_logr_mapped).Value

        for continuous_contract_symbol, continuous_contract in self.continuous_contracts.items():
            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.Log(mapping_debug + f' --> LONG {mapped_contract_debug}')
                self.MarketOrder(mapped_contract, self.quantity,
                                 tag=f'LONG {mapped_contract_debug} - {self.test_tag}')
            elif lowest_logr == continuous_contract_symbol:
                self.Log(mapping_debug + f' --> SHORT {mapped_contract_debug}')
                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
        self.Log(f"Next rebalance at {self.nextRebalance}")

    def Momentum_Signal(self, mom_indicator_dict):
        self.Log('Calculating momentum signal:')
        # 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)
            else:
                self.Log(
                    f'{i} indicator not ready, {mom_indicator_dict[i].Samples} samples')

        # sort logr values
        sorted_logr_values = dict(
            sorted(logr_values.items(), key=lambda x: x[1]))
        self.Log(f"Sorted logr values: {sorted_logr_values}")

        # get highest and lowest logr values
        if len(sorted_logr_values) < 2:
            self.Log('Not enough data to calculate momentum signal')
            return None, None

        highest_logr = list(sorted_logr_values.keys())[-1]
        lowest_logr = list(sorted_logr_values.keys())[0]
        self.Log(
            f"Highest logr: {highest_logr} with value {sorted_logr_values[highest_logr]}")
        self.Log(
            f"Lowest logr: {lowest_logr} with value {sorted_logr_values[lowest_logr]}")

        return highest_logr, lowest_logr