from AlgorithmImports import *
import numpy as np
import math


class EndOfMonthSP500Reversal(QCAlgorithm):
    def initialize(self):
        self.set_start_date(2023, 1, 1)
        self.set_cash(1_000_000)

        # Add ES futures continuous contract (price-based, daily)
        self._es = self.add_future("ES", Resolution.DAILY, extended_market_hours=False)
        self._es.set_leverage(1.0)

        # Realistic costs
        self.set_brokerage_model(BrokerageName.INTERACTIVE_BROKERS_BROKERAGE, AccountType.MARGIN)

        # Monthly rebalance: first trading day of each month at 10:00 ET
        self.schedule.on(
            self.date_rules.month_start('SPY'),
            self.time_rules.at(10, 0),
            self._rebalance,
        )

        # Persistent data structures
        self._pairs = []               # list of (r^EoM_s, r_{s+1})
        self._eom_returns = {}         # month_key -> r^EoM
        self._month_end_prices = {}    # month_key -> log_price of last trading day
        self._processed_months = set() # months already added to regression

        # Seed initial ~102 months from history before backtest start
        self._seed_initial_window()

        self._current_weight = 0.0
        self._first_rebalance = True

    def _seed_initial_window(self):
        # ~2500 trading days gives roughly 102 months
        history = self.history(self._es.symbol, 2500, Resolution.DAILY)
        if history.empty:
            self.debug("No history available for initial window")
            return
        self._process_history(history)
        self.debug(f"Seeded initial window: {len(self._pairs)} pairs from {len(self._processed_months)} months")

    def _process_history(self, history):
        history = history.sort_index()
        closes = history['close'].values
        if len(closes) == 0:
            return

        log_prices = np.log(closes)

        # Extract dates from index: futures MultiIndex has (expiry, symbol, time)
        if isinstance(history.index, pd.MultiIndex):
            if 'time' in history.index.names:
                dates = history.index.get_level_values('time')
            else:
                dates = history.index.get_level_values(-1)
        else:
            dates = history.index

        # Group trading days by calendar month
        month_days = {}
        for i, d in enumerate(dates):
            d_date = pd.Timestamp(d).date()
            month_key = (d_date.year, d_date.month)
            if month_key not in month_days:
                month_days[month_key] = []
            month_days[month_key].append((i, d_date, log_prices[i]))

        # Process each month in chronological order
        for month_key in sorted(month_days.keys()):
            if month_key in self._processed_months:
                continue

            days = month_days[month_key]
            if len(days) < 5:
                self._processed_months.add(month_key)
                continue

            last_idx, last_date, last_log_price = days[-1]
            fourth_idx, fourth_date, fourth_log_price = days[-5]

            # r^EoM_t = log(p_last) - log(p_{4th-to-last})
            eom_return = last_log_price - fourth_log_price
            self._eom_returns[month_key] = eom_return
            self._month_end_prices[month_key] = last_log_price
            self._processed_months.add(month_key)

            # Monthly return: r_t = log(p_last_t) - log(p_last_{t-1})
            prev_month = self._get_prev_month(month_key)
            if prev_month in self._month_end_prices:
                monthly_return = last_log_price - self._month_end_prices[prev_month]

                # Add pair (r^EoM_{t-1}, r_t) if we have the EoM of the month before previous
                prev_eom_month = self._get_prev_month(prev_month)
                if prev_eom_month in self._eom_returns:
                    pair = (self._eom_returns[prev_eom_month], monthly_return)
                    self._pairs.append(pair)

    def _get_prev_month(self, month_key):
        year, month = month_key
        if month == 1:
            return (year - 1, 12)
        return (year, month - 1)

    def _rebalance(self):
        # Update with most recent history to catch any newly completed months
        history = self.history(self._es.symbol, 60, Resolution.DAILY)
        if not history.empty:
            self._process_history(history)

        # Latest completed month before the current month
        today = self.time.date()
        current_month = (today.year, today.month)

        completed_months = [m for m in sorted(self._eom_returns.keys()) if m < current_month]
        if len(completed_months) == 0:
            return

        latest_month = completed_months[-1]
        latest_eom = self._eom_returns[latest_month]

        # Need at least a reasonable number of pairs for regression
        if len(self._pairs) < 24:
            return

        # Expanding-window OLS
        X = np.array([[1.0, p[0]] for p in self._pairs])
        y = np.array([p[1] for p in self._pairs])

        try:
            beta = np.linalg.lstsq(X, y, rcond=None)[0]
            alpha_hat = beta[0]
            gamma_hat = beta[1]
        except Exception as e:
            self.debug(f"OLS error: {e}")
            return

        # Forecast next month's return
        r_hat = alpha_hat + gamma_hat * latest_eom

        # Variance of trailing 3 monthly returns (sample variance)
        monthly_returns = [p[1] for p in self._pairs]
        if len(monthly_returns) < 3:
            return

        trailing = np.array(monthly_returns[-3:])
        sigma2 = np.var(trailing, ddof=1)
        if sigma2 <= 0 or np.isnan(sigma2) or np.isinf(sigma2):
            sigma2 = 1e-10

        # Campbell-Thompson optimal allocation, risk aversion = 1
        w = r_hat / sigma2
        w = max(-1.0, min(1.5, w))
        self._current_weight = w

        if self._first_rebalance:
            self.debug(f"First rebalance: month={latest_month}, pairs={len(self._pairs)}, alpha={alpha_hat:.6f}, gamma={gamma_hat:.6f}, eom={latest_eom:.6f}, r_hat={r_hat:.6f}, sigma2={sigma2:.6f}, w={w:.4f}")
            self._first_rebalance = False

        # Trade
        self._trade_futures(w)

    def _trade_futures(self, target_weight):
        mapped = self._es.mapped
        if mapped is None or mapped == self._es.symbol:
            self.debug("No mapped contract available")
            return

        if mapped not in self.securities:
            self.debug("Mapped contract not in securities")
            return

        price = self.securities[mapped].price
        if price <= 0 or np.isnan(price):
            return

        symbol_props = self.securities[mapped].symbol_properties
        multiplier = symbol_props.contract_multiplier if symbol_props is not None else 50.0
        if multiplier is None or multiplier <= 0:
            multiplier = 50.0

        portfolio_value = self.portfolio.total_portfolio_value
        target_notional = portfolio_value * target_weight

        target_contracts = round(target_notional / (price * multiplier))
        current_contracts = self.portfolio[mapped].quantity

        delta = target_contracts - current_contracts
        if delta != 0:
            self.market_order(mapped, delta)
            self.debug(f"Rebalance: target_weight={target_weight:.4f}, contracts={target_contracts}, delta={delta}")

    def on_data(self, data):
        pass