# region imports
from AlgorithmImports import *

from itertools import combinations
import statsmodels.api as sm
from scipy.stats import norm
# endregion


class FuturesCollection:

    _minimum_observations = 24 # months

    def __init__(self, algorithm, tickers):
        self._algorithm = algorithm
        self._futures = []
        for ticker in tickers:
            # Add the Future contracts and define the continuous contract settings.
            future = algorithm.add_future(
                ticker,
                data_mapping_mode=DataMappingMode.FIRST_DAY_MONTH,
                data_normalization_mode=DataNormalizationMode.BACKWARDS_PANAMA_CANAL,
                contract_depth_offset=0
            )
            # Select the 2 front contracts.
            future.set_filter(lambda universe: universe.contracts(lambda symbols: sorted(symbols, key=lambda s: s.id.date)[:2]))
            
            # Create a Series to hold the monthly returns.
            future.returns = pd.Series()
            roc = RateOfChange(1)
            algorithm.consolidate(future.symbol, Calendar.MONTHLY, lambda bar, roc=roc: roc.update(bar.end_time, bar.close))
            roc.updated += lambda indicator, indicator_data_point, future=future: self._update_returns_history(indicator, indicator_data_point, future)

            # Create a DataFrame to hold the factor history.
            future.factors = pd.DataFrame()

            # Define the momentum factor: Average trailing 12-month return (using monthly returns)
            momentum = IndicatorExtensions.sma(roc, 12)
            momentum.updated += lambda indicator, indicator_data_point, future=future: self._update_factor_history(indicator, indicator_data_point, future)

            # Create an indicator we'll need to calculate the value factor (5-year mean price).
            future.mean_price = algorithm.sma(future.symbol, 5*252, Resolution.DAILY)

            self._futures.append(future)
    
    def _update_returns_history(self, indicator, indicator_data_point, future):
        if indicator.is_ready: 
            future.returns.loc[indicator_data_point.end_time] = indicator_data_point.value

    def _update_factor_history(self, indicator, indicator_data_point, future):
        t = indicator_data_point.end_time
        # Momentum
        if indicator.is_ready:
            future.factors.loc[t, 'momentum'] = indicator_data_point.value
        # Value
        if future.mean_price.is_ready:
            future.factors.loc[t, 'value'] = 1 - future.price / future.mean_price.current.value
        # Carry
        if future.symbol not in self._algorithm.current_slice.future_chains:
            return
        chain = self._algorithm.current_slice.future_chains[future.symbol]
        if len(list(chain.contracts)) < 2:
            return
        front_contract, next_contract = sorted(chain, key=lambda c: c.expiry)
        carry = front_contract.last_price - next_contract.last_price
        months_between_contracts = round((next_contract.expiry - front_contract.expiry).days / 30)
        expiry_difference_in_years = abs(months_between_contracts) / 12
        annualized_carry = carry / expiry_difference_in_years
        future.factors.loc[t, 'carry'] = annualized_carry

    def tradable(self):
        return all([self._algorithm.is_market_open(future.symbol) and future.symbol in self._algorithm.current_slice.future_chains for future in self._futures])

    def create_weights(self, selectivity_level, maximum_observations):
        # Calculate the conditional expected return of the base pairs for each factor.
        thetas_by_factor = {}
        for future_i, future_j in combinations(self._futures, 2):
            # Analyze each factor independently.
            for factor in set(future_i.factors.columns) & set(future_j.factors.columns):
                # Fit the base pairs models.
                x_i = future_i.factors[factor].dropna()
                x_j = future_j.factors[factor].dropna()
                if len(x_i) < self._minimum_observations or len(x_j) < self._minimum_observations:
                    continue
                r_i = future_i.returns.shift(-1)[:-1]  # `shift` so that it represents forward return
                r_j = future_j.returns.shift(-1)[:-1]
                idx = sorted(list(set(r_i.index) & set(r_j.index) & set(x_i.index) & set(x_j.index)))[-maximum_observations:]
                X = sm.add_constant(pd.concat([x_i[idx], x_j[idx]], axis=1))
                mu_i, b_ii, b_ji = sm.OLS(r_i[idx], X).fit().params
                mu_j, b_ij, b_jj = sm.OLS(r_j[idx], X).fit().params
                # Decompose theta into the three components.
                if x_i[-1] == x_j[-1]:
                    continue
                m_i, m_j = x_i.mean(), x_j.mean()
                s_i, s_j = x_i.std(), x_j.std()
                p_ij = x_i.corr(x_j)
                s_ij = (s_i**2 + s_j**2 - (2*s_i*s_j*p_ij))**(1/2)
                y_ij = -((m_i - m_j)/s_ij)
                cdf = norm.cdf(y_ij) 
                pdf = norm.pdf(y_ij)
                oa_factor = ((b_ii*s_i**2 + b_jj*s_j**2 - (b_ii+b_jj)*s_i*s_j*p_ij) / s_ij)
                ca_factor = ((b_ij*s_i**2 + b_ji*s_j**2 - (b_ij+b_ji)*s_i*s_j*p_ij) / s_ij)
                if x_i[-1] > x_j[-1]:
                    ue = mu_i - mu_j  # Unexplained effect of asset means
                    oa = b_ii*m_i - b_jj*m_j + (pdf / (1-cdf)) * oa_factor  # Own-asset effect
                    ca = -(b_ij*m_i - b_ji*m_j + (pdf / (1-cdf)) * ca_factor)  # Cross-asset effect
                    pair = (future_i, future_j)
                else:
                    ue = mu_j - mu_i  # Unexplained effect of asset means 
                    oa = b_jj*m_j - b_ii*m_i + (pdf / cdf) * oa_factor  # Own-asset effect
                    ca = -(b_ji*m_j - b_ij*m_i + (pdf / cdf) * ca_factor)  # Cross-asset effect
                    pair = (future_j, future_i)
                if factor not in thetas_by_factor:
                    thetas_by_factor[factor] = {}
                thetas_by_factor[factor][pair] = ue + oa + ca
        
        # Create a dictionary to track the weights.
        weight_by_symbol = {future.mapped: 0 for future in self._futures}
        for theta_by_pair in thetas_by_factor.values():
            # Select the top s% of the base pairs.
            selected_pairs = [kvp[0] for kvp in sorted(theta_by_pair.items(), key=lambda kvp: kvp[1], reverse=True)[:int(selectivity_level*len(theta_by_pair))]]
            # Calculate portfolio weights.
            n = 2*len(selected_pairs)
            denominator = n**2 - int(n % 2 != 0)  # `n % 2 != 0` checks if n is odd
            if not denominator:
                continue
            a = (2*(n-1))/denominator
            h = 4/denominator
            for i, pair in enumerate(selected_pairs):
                weight = (a - i*h) / len(thetas_by_factor)  # Divide to give equal weight to each factor.
                weight_by_symbol[pair[0].mapped] += weight  # Long the Future with the larger signal.
                weight_by_symbol[pair[1].mapped] -= weight  # Short the Future with the smaller signal.
        return weight_by_symbol


class EquityIndexFutures(FuturesCollection):

    def __init__(self, algorithm):
        tickers = [
            Futures.Indices.EURO_STOXX_50,
            Futures.Indices.NIKKEI_225_DOLLAR,
            Futures.Indices.NASDAQ_100_E_MINI,
            Futures.Indices.RUSSELL_2000_E_MINI,
            Futures.Indices.SP_500_E_MINI
        ]
        super().__init__(algorithm, tickers)
    
    def create_weights(self, selectivity_level):
        return super().create_weights(selectivity_level, 10*12)

    #def value(self):  # earnings/price ratio (3-year z-score)
    #    pass

    #def carry(self):  # futures roll (spot minus first)
    #    pass


class BondFutures(FuturesCollection):

    def __init__(self, algorithm):
        tickers = [
            Futures.Financials.Y_2_TREASURY_NOTE,
            Futures.Financials.Y_5_TREASURY_NOTE,
            Futures.Financials.Y_10_TREASURY_NOTE
        ]
        super().__init__(algorithm, tickers)

    def create_weights(self, selectivity_level):
        return super().create_weights(selectivity_level, 15*12)

    #def value(self):  # real yield
    #    pass

    #def carry(self):  # excess yield plus term-structure roll
    #    pass


class CurrencyFutures(FuturesCollection):

    def __init__(self, algorithm):
        tickers = [
            #Futures.Currencies.AUD,
            Futures.Currencies.CAD,
            Futures.Currencies.CHF,
            Futures.Currencies.EUR,
            Futures.Currencies.GBP,
            #Futures.Currencies.JPY,  Data issue
            Futures.Currencies.NZD
        ]
        super().__init__(algorithm, tickers)

    def create_weights(self, selectivity_level):
        return super().create_weights(selectivity_level, 20*12)

    #def value(self):  # 5-year real exchange rate reversal (3-year z-score)
    #    pass

    #def carry(self):  # cash rate differential
    #    pass

class CommodityFutures(FuturesCollection):

    def __init__(self, algorithm):
        tickers = [
            Futures.Energy.CRUDE_OIL_WTI,
            Futures.Energy.GASOLINE,
            Futures.Energy.HEATING_OIL,
            Futures.Energy.NATURAL_GAS,
            Futures.Grains.CORN,
            Futures.Grains.SOYBEAN_MEAL,
            Futures.Grains.SOYBEAN_OIL,
            Futures.Grains.SOYBEANS,
            Futures.Grains.WHEAT,
            Futures.Meats.FEEDER_CATTLE,
            Futures.Meats.LEAN_HOGS,
            Futures.Meats.LIVE_CATTLE,
            Futures.Metals.COPPER,
            Futures.Metals.GOLD,
            Futures.Metals.PLATINUM,
            Futures.Metals.SILVER,
            Futures.Softs.SUGAR_11
        ]
        super().__init__(algorithm, tickers)

    def create_weights(self, selectivity_level):
        return super().create_weights(selectivity_level, 20*12)

    #def value(self):  # 5-year real exchange rate reversal
    #    pass

    #def carry(self):  # futures roll (spot minus first)
    #    pass

# region imports
from AlgorithmImports import *
from futures_collection import *
# endregion

#TODO: 
# - Add leverage scaling (page 31)

class InvestmentBasePairsAlgorithm(QCAlgorithm):

    def initialize(self):
        self.set_start_date(2019, 1, 1)
        self.set_cash(10_000_000)
        self.settings.automatic_indicator_warm_up = True
        self._month = 0
        self._selectivity_level = self.get_parameter('selectivity_level', 0.5)  # 0.05 = 5%
        self._asset_classes = [EquityIndexFutures(self)]  # [EquityIndexFutures(self), BondFutures(self), CurrencyFutures(self), CommodityFutures(self)]
        self.set_warm_up(timedelta(400 + 24*31))

    def on_data(self, data: Slice):
        # Rebalance at the start of each month, when all of the markets are open.
        if (self.is_warming_up or 
            self._month == self.time.month or 
            not all([futures_collection.tradable() for futures_collection in self._asset_classes])):
            return
        self._month = self.time.month
        # Calculate the weight of each Future.
        weight_by_symbol = {}
        for futures_collection in self._asset_classes:
            for symbol, weight in futures_collection.create_weights(self._selectivity_level).items():
                if symbol not in weight_by_symbol:
                    weight_by_symbol[symbol] = 0
                weight_by_symbol[symbol] += (weight/len(self._asset_classes))  # Divide to give equal weight to each asset class.
        # Re-scale weights to keep unit gross exposure. 
        abs_weight_sum = sum([abs(x) for x in weight_by_symbol.values()])
        if not abs_weight_sum:
            return
        weight_by_symbol = {s: w/abs_weight_sum for s, w in weight_by_symbol.items()}    
        # Rebalance the portfolio.
        self.set_holdings([PortfolioTarget(symbol, weight/2) for symbol, weight in weight_by_symbol.items()], True)  ### NOTE: Dividing weight by 2 here to avoid order errors!