#region imports
from AlgorithmImports import *

from utils import get_position_size
from futures import categories
#endregion

class CarryAndTrendAlphaModel(AlphaModel):

    _futures = []
    _BUSINESS_DAYS_IN_YEAR = 256
    _TREND_FORECAST_SCALAR_BY_SPAN = {64: 1.91, 32: 2.79, 16: 4.1, 8: 5.95, 4: 8.53, 2: 12.1} # Table 29 on page 177
    _CARRY_FORECAST_SCALAR = 30 # Provided on p.216
    _FDM_BY_RULE_COUNT = { # Table 52 on page 234
        1: 1.0,
        2: 1.02,
        3: 1.03,
        4: 1.23,
        5: 1.25,
        6: 1.27,
        7: 1.29,
        8: 1.32,
        9: 1.34,
    }

    def __init__(self, algorithm, emac_filters, abs_forecast_cap, sigma_span, target_risk, blend_years):
        self._algorithm = algorithm
        
        self._emac_spans = [2**x for x in range(4, emac_filters+1)]
        self._fast_ema_spans = self._emac_spans
        # "Any ratio between the two moving average lengths of two and six gives statistically indistinguishable results." (p.165)
        self._slow_ema_spans = [fast_span * 4 for fast_span in self._emac_spans] 
        self._all_ema_spans = sorted(list(set(self._fast_ema_spans + self._slow_ema_spans)))

        self._carry_spans = [5, 20, 60, 120]

        self._annualization_factor = self._BUSINESS_DAYS_IN_YEAR ** 0.5

        self._abs_forecast_cap = abs_forecast_cap
        
        self._sigma_span = sigma_span
        self._target_risk = target_risk
        self._blend_years = blend_years

        # Instrument Diversification Multiplier. Hardcoded in https://gitfront.io/r/user-4000052/iTvUZwEUN2Ta/AFTS-CODE/blob/chapter8.py
        self._idm = 1.5                                                    

        self._categories = categories
        self._total_lookback = timedelta(sigma_span*(7/5) + blend_years*365)

        self._day = -1

    def update(self, algorithm: QCAlgorithm, data: Slice) -> List[Insight]:
        if data.quote_bars.count:
            for future in self._futures:
                future.latest_mapped = future.mapped

        # Rebalance daily
        if self._day == data.time.day or data.quote_bars.count == 0:
            return []

        # Update annualized carry data
        for future in self._futures:
            # Get the near and far contracts
            contracts = self._get_near_and_further_contracts(algorithm.securities, future.mapped)
            if contracts is None:
                continue
            near_contract, further_contract = contracts[0], contracts[1]
            
            # Save near and further contract for later
            future.near_contract = near_contract
            future.further_contract = further_contract

            # Check if the daily consolidator has provided a bar for these contracts yet
            if not hasattr(near_contract, "raw_history") or not hasattr(further_contract, "raw_history") or near_contract.raw_history.empty or further_contract.raw_history.empty:
                continue
            # Update annualized raw carry history
            raw_carry = near_contract.raw_history.iloc[0] - further_contract.raw_history.iloc[0]
            months_between_contracts = round((further_contract.expiry - near_contract.expiry).days / 30)
            expiry_difference_in_years = abs(months_between_contracts) / 12
            annualized_raw_carry = raw_carry / expiry_difference_in_years
            future.annualized_raw_carry_history.loc[near_contract.raw_history.index[0]] = annualized_raw_carry

        # If warming up and still > 7 days before start date, don't do anything
        # We use a 7-day buffer so that the algorithm has active insights when warm-up ends
        if algorithm.start_date - algorithm.time > timedelta(7):
            self._day = data.time.day
            return []
        
        # Estimate the standard deviation of % daily returns for each future
        sigma_pcts_by_future = {}
        for future in self._futures:
            sigma_pcts = self._estimate_std_of_pct_returns(future.raw_history, future.adjusted_history)
            # Check if there is sufficient history
            if sigma_pcts is None:
                continue
            sigma_pcts_by_future[future] = sigma_pcts

        # Create insights
        insights = []
        weight_by_symbol = get_position_size({future.symbol: self._categories[future.symbol].classification for future in sigma_pcts_by_future.keys()})
        for symbol, instrument_weight in weight_by_symbol.items():
            future = algorithm.securities[symbol]
            target_contract = [future.near_contract, future.further_contract][self._categories[future.symbol].contract_offset]
            sigma_pct = sigma_pcts_by_future[future]
            # "The price should be for the expiry date we currently hold (not the back-adjusted price)" (p.55)
            daily_risk_price_terms = sigma_pct / (self._annualization_factor) * target_contract.price 

            # Calculate target position
            position = (
                (algorithm.portfolio.total_portfolio_value * self._idm * instrument_weight * self._target_risk)
                / (future.symbol_properties.contract_multiplier * daily_risk_price_terms * (self._annualization_factor))
            )

            # Calculate forecast type 1: EMAC
            trend_forecasts = self._calculate_emac_forecasts(future.ewmac_by_span, daily_risk_price_terms)
            if not trend_forecasts:
                continue
            emac_combined_forecasts = sum(trend_forecasts) / len(trend_forecasts) # Aggregate EMAC factors -- equal-weight

            # Calculate factor type 2: Carry
            carry_forecasts = self._calculate_carry_forecasts(future.annualized_raw_carry_history, daily_risk_price_terms)
            if not carry_forecasts:
                continue
            carry_combined_forecasts = sum(carry_forecasts) / len(carry_forecasts) # Aggregate Carry factors -- equal-weight
            
            # Aggregate factors -- 60% for trend, 40% for carry
            raw_combined_forecast = 0.6 * emac_combined_forecasts + 0.4 * carry_combined_forecasts
            # Apply a forecast diversification multiplier to keep the average forecast at 10 (p 193-194)
            scaled_combined_forecast = raw_combined_forecast * self._FDM_BY_RULE_COUNT[len(trend_forecasts) + len(carry_forecasts)] 
            capped_combined_forecast = max(min(scaled_combined_forecast, self._abs_forecast_cap), -self._abs_forecast_cap)

            if capped_combined_forecast * position == 0:
                continue
            target_contract.forecast = capped_combined_forecast
            target_contract.position = position
            
            local_time = Extensions.convert_to(algorithm.time, algorithm.time_zone, future.exchange.time_zone)
            expiry = future.exchange.hours.get_next_market_open(local_time, False) - timedelta(seconds=1)
            insights.append(Insight.price(target_contract.symbol, expiry, InsightDirection.UP if capped_combined_forecast * position > 0 else InsightDirection.DOWN))
        
        if insights:
            self._day = data.time.day

        return insights

    def _align_history(self, a, b):
        idx = sorted(list(set(a.index).intersection(set(b.index)))) 
        return a.loc[idx], b.loc[idx]

    def _calculate_emac_forecasts(self, ewmac_by_span, daily_risk_price_terms):
        forecasts = []
        for span in self._emac_spans:
            risk_adjusted_ewmac = ewmac_by_span[span].current.value / daily_risk_price_terms
            scaled_forecast_for_ewmac = risk_adjusted_ewmac * self._TREND_FORECAST_SCALAR_BY_SPAN[span]
            capped_forecast_for_ewmac = max(min(scaled_forecast_for_ewmac, self._abs_forecast_cap), -self._abs_forecast_cap)
            forecasts.append(capped_forecast_for_ewmac)
        return forecasts

    def _calculate_carry_forecasts(self, annualized_raw_carry, daily_risk_price_terms):
        carry_forecast = annualized_raw_carry / daily_risk_price_terms

        forecasts = []
        for span in self._carry_spans:
            ## Smooth out carry forecast
            smoothed_carry_forecast = carry_forecast.ewm(span=span, min_periods=span).mean().dropna()
            if smoothed_carry_forecast.empty:
                continue
            smoothed_carry_forecast = smoothed_carry_forecast.iloc[-1]
            ## Apply forecast scalar (p. 264)
            scaled_carry_forecast = smoothed_carry_forecast * self._CARRY_FORECAST_SCALAR
            ## Cap forecast
            capped_carry_forecast = max(min(scaled_carry_forecast, self._abs_forecast_cap), -self._abs_forecast_cap)
            forecasts.append(capped_carry_forecast)

        return forecasts

    def _get_near_and_further_contracts(self, securities, mapped_symbol):
        ## Gather and align history of near/further contracts
        contracts_sorted_by_expiry = sorted(
            [
                kvp.Value for kvp in securities 
                if not kvp.key.is_canonical() and kvp.key.canonical == mapped_symbol.canonical and kvp.Value.Expiry >= securities[mapped_symbol].Expiry
            ], 
            key=lambda contract: contract.expiry
        )
        if len(contracts_sorted_by_expiry) < 2:
            return None
        near_contract = contracts_sorted_by_expiry[0]
        further_contract = contracts_sorted_by_expiry[1]
        return near_contract, further_contract

    def _estimate_std_of_pct_returns(self, raw_history, adjusted_history):
        # Align history of raw and adjusted prices
        raw_history_aligned, adjusted_history_aligned = self._align_history(raw_history, adjusted_history)

        # Calculate exponentially weighted standard deviation of returns
        returns = adjusted_history_aligned.diff().dropna() / raw_history_aligned.shift(1).dropna() 
        rolling_ewmstd_pct_returns = returns.ewm(span=self._sigma_span, min_periods=self._sigma_span).std().dropna()
        if rolling_ewmstd_pct_returns.empty: # Not enough history
            return None
        # Annualize sigma estimate
        annulized_rolling_ewmstd_pct_returns = rolling_ewmstd_pct_returns * (self._annualization_factor)
        # Blend the sigma estimate (p.80)
        blended_estimate = 0.3*annulized_rolling_ewmstd_pct_returns.mean() + 0.7*annulized_rolling_ewmstd_pct_returns.iloc[-1]
        return blended_estimate

    def _consolidation_handler(self, sender: object, consolidated_bar: TradeBar) -> None:
        security = self._algorithm.securities[consolidated_bar.symbol]
        end_date = consolidated_bar.end_time.date()
        if security.symbol.is_canonical():
            # Update adjusted history
            security.adjusted_history.loc[end_date] = consolidated_bar.close
            security.adjusted_history = security.adjusted_history[security.adjusted_history.index >= end_date - self._total_lookback]
        else:
            # Update raw history
            continuous_contract = self._algorithm.securities[security.symbol.canonical]
            if hasattr(continuous_contract, "latest_mapped") and consolidated_bar.symbol == continuous_contract.latest_mapped:
                continuous_contract.raw_history.loc[end_date] = consolidated_bar.close
                continuous_contract.raw_history = continuous_contract.raw_history[continuous_contract.raw_history.index >= end_date - self._total_lookback]
            
            # Update raw carry history
            security.raw_history.loc[end_date] = consolidated_bar.close
            security.raw_history = security.raw_history.iloc[-1:]

    def on_securities_changed(self, algorithm: QCAlgorithm, changes: SecurityChanges) -> None:
        for security in changes.added_securities:
            symbol = security.symbol

            # Create a consolidator to update the history
            security.consolidator = TradeBarConsolidator(timedelta(1))
            security.consolidator.data_consolidated += self._consolidation_handler
            algorithm.subscription_manager.add_consolidator(symbol, security.consolidator)

            # Get raw and adjusted history
            security.raw_history = pd.Series()

            if symbol.is_canonical():
                security.adjusted_history = pd.Series()
                security.annualized_raw_carry_history = pd.Series()

                # Create indicators for the continuous contract
                ema_by_span = {span: algorithm.ema(symbol, span, Resolution.DAILY) for span in self._all_ema_spans}
                security.ewmac_by_span = {}
                for i, fast_span in enumerate(self._emac_spans):
                    security.ewmac_by_span[fast_span] = IndicatorExtensions.minus(ema_by_span[fast_span], ema_by_span[self._slow_ema_spans[i]])

                security.automatic_indicators = ema_by_span.values()

                self._futures.append(security)

        for security in changes.removed_securities:
            # Remove consolidator + indicators
            algorithm.subscription_manager.remove_consolidator(security.symbol, security.consolidator)
            if security.symbol.is_canonical():
                for indicator in security.automatic_indicators:
                    algorithm.deregister_indicator(indicator)

# region imports
from AlgorithmImports import *
# endregion


class FutureData:
    def __init__(self, classification, contract_offset):
        self.classification = classification
        self.contract_offset = contract_offset

categories = {
    pair[0]: FutureData(pair[1], pair[2]) for pair in [
        (Symbol.create(Futures.Indices.SP_500_E_MINI, SecurityType.FUTURE, Market.CME), ("Equity", "US"), 0),
        (Symbol.create(Futures.Indices.NASDAQ_100_E_MINI, SecurityType.FUTURE, Market.CME), ("Equity", "US"), 0),
        (Symbol.create(Futures.Indices.RUSSELL_2000_E_MINI, SecurityType.FUTURE, Market.CME), ("Equity", "US"), 0),
        (Symbol.create(Futures.Indices.VIX, SecurityType.FUTURE, Market.CFE), ("Volatility", "US"), 0),
        (Symbol.create(Futures.Energies.NATURAL_GAS, SecurityType.FUTURE, Market.NYMEX), ("Energies", "Gas"), 1),
        (Symbol.create(Futures.Energies.CRUDE_OIL_WTI, SecurityType.FUTURE, Market.NYMEX), ("Energies", "Oil"), 0),
        (Symbol.create(Futures.Grains.CORN, SecurityType.FUTURE, Market.CBOT), ("Agricultural", "Grain"), 0),
        (Symbol.create(Futures.Metals.COPPER, SecurityType.FUTURE, Market.COMEX), ("Metals", "Industrial"), 0),
        (Symbol.create(Futures.Metals.GOLD, SecurityType.FUTURE, Market.COMEX), ("Metals", "Precious"), 0),
        (Symbol.create(Futures.Metals.SILVER, SecurityType.FUTURE, Market.COMEX), ("Metals", "Precious"), 0)
    ]
}

# region imports
from AlgorithmImports import *
from universe import AdvancedFuturesUniverseSelectionModel
from alpha import CarryAndTrendAlphaModel
from portfolio import BufferedPortfolioConstructionModel
from utils import get_position_size
# endregion


class FuturesCombinedCarryAndTrendAlgorithm(QCAlgorithm):
    
    def initialize(self):
        self.set_start_date(self.end_date - timedelta(5*365))
        self.set_cash(1_000_000)

        self.set_brokerage_model(BrokerageName.INTERACTIVE_BROKERS_BROKERAGE, AccountType.MARGIN)
        self.settings.seed_initial_prices = True
        self.settings.minimum_order_margin_portfolio_percentage = 0

        self.universe_settings.data_normalization_mode = DataNormalizationMode.BACKWARDS_PANAMA_CANAL
        self.universe_settings.data_mapping_mode = DataMappingMode.LAST_TRADING_DAY
        self.add_universe_selection(AdvancedFuturesUniverseSelectionModel())

        self.add_alpha(CarryAndTrendAlphaModel(
            self,
            self.get_parameter("emac_filters", 6), 
            self.get_parameter("abs_forecast_cap", 20),           # Hardcoded on p.173
            self.get_parameter("sigma_span", 32),
            self.get_parameter("target_risk", 0.2),               # Recommend value is 0.2 on p.75
            self.get_parameter("blend_years", 3)                  # Number of years to use when blending sigma estimates
        ))

        self.settings.rebalance_portfolio_on_security_changes = False
        self.settings.rebalance_portfolio_on_insight_changes = False
        self.total_count = 0
        self.day = -1
        self.set_portfolio_construction(BufferedPortfolioConstructionModel(
            self.rebalance_func,
            self.get_parameter("buffer_scaler", 0.1)              # Hardcoded on p.167 & p.173
        ))

        self.add_risk_management(NullRiskManagementModel())

        self.set_execution(ImmediateExecutionModel())

        # We need several years of data to warm-up. Data before 2014 can have issues.
        self.set_warm_up(self.start_date - datetime(2015, 1, 1)) 

    def rebalance_func(self, time):
        if (self.total_count != self.insights.total_count or self.day != self.time.day) and not self.is_warming_up and self.current_slice.quote_bars.count > 0:
            self.total_count = self.insights.total_count
            self.day = self.time.day
            return time
        return None

#region imports
from AlgorithmImports import *
#endregion


class BufferedPortfolioConstructionModel(EqualWeightingPortfolioConstructionModel):

    def __init__(self, rebalance, buffer_scaler):
        super().__init__(rebalance)
        self._buffer_scaler = buffer_scaler

    def create_targets(self, algorithm: QCAlgorithm, insights: List[Insight]) -> List[PortfolioTarget]:
        targets = super().create_targets(algorithm, insights)
        adj_targets = []
        for insight in insights:
            future_contract = algorithm.securities[insight.symbol]
            optimal_position = future_contract.forecast * future_contract.position / 10

            # Create buffer zone to reduce churn
            buffer_width = self._buffer_scaler * abs(future_contract.position)
            upper_buffer = round(optimal_position + buffer_width)
            lower_buffer = round(optimal_position - buffer_width)
            
            # Determine quantity to put holdings into buffer zone
            current_holdings = future_contract.holdings.quantity
            if lower_buffer <= current_holdings <= upper_buffer:
                continue
            quantity = lower_buffer if current_holdings < lower_buffer else upper_buffer

            # Place trades
            adj_targets.append(PortfolioTarget(insight.symbol, quantity))
        
        # Liquidate contracts that have an expired insight
        for target in targets:
            if target.quantity == 0:
                adj_targets.append(target)

        return adj_targets

# region imports
from AlgorithmImports import *
from futures import categories
# endregion


class AdvancedFuturesUniverseSelectionModel(FutureUniverseSelectionModel):
    
    def __init__(self) -> None:
        super().__init__(timedelta(1), self._select_future_chain_symbols)
        self._symbols = list(categories.keys())

    def _select_future_chain_symbols(self, utc_time: datetime) -> List[Symbol]:
        return self._symbols

    def filter(self, filter: FutureFilterUniverse) -> FutureFilterUniverse:
        return filter.expiration(0, 365)

#region imports
from AlgorithmImports import *
#endregion


def get_position_size(group):
    subcategories = {}
    for category, subcategory in group.values():
        if category not in subcategories:
            subcategories[category] = {subcategory: 0}
        elif subcategory not in subcategories[category]:
            subcategories[category][subcategory] = 0
        subcategories[category][subcategory] += 1

    category_count = len(subcategories.keys())
    subcategory_count = {
        category: len(subcategory.keys()) 
        for category, subcategory in subcategories.items()
    }
    
    weights = {}
    for symbol in group:
        category, subcategory = group[symbol]
        weights[symbol] = (
            1 
            / category_count 
            / subcategory_count[category] 
            / subcategories[category][subcategory]
        )
    return weights