from AlgorithmImports import *
from datetime import timedelta


class FutureCLGCInflationHedge(QCAlgorithm):

    def initialize(self):
        self.set_start_date(self.end_date - timedelta(5*365))
        self.set_cash(100000)
        self.settings.seed_initial_prices = True
        
        # Add continuous futures
        self._gc = self.add_future("GC", Resolution.DAILY)
        self._cl = self.add_future("CL", Resolution.DAILY)
        
        # Rolling windows for 30-day history
        self._gc_prices = RollingWindow[float](30)
        self._cl_prices = RollingWindow[float](30)
        self._ratio_window = RollingWindow[float](30)
        
        self._hedge_ratio = 1.0
        self._state = 0  # 0=flat, 1=long ratio, -1=short ratio
        
        # Schedule daily signal check before market open
        self.schedule.on(
            self.date_rules.every_day("GC"),
            self.time_rules.at(7, 0),
            self._check_signals
        )
        
        # Monthly hedge ratio re-fit before market open
        self.schedule.on(
            self.date_rules.month_start("GC"),
            self.time_rules.at(7, 15),
            self._refit_hedge_ratio
        )

    def _refit_hedge_ratio(self):
        if self._gc_prices.count < 30 or self._cl_prices.count < 30:
            return
        
        gc_list = list(self._gc_prices)[::-1]
        cl_list = list(self._cl_prices)[::-1]
        n = len(gc_list)
        
        mean_gc = sum(gc_list) / n
        mean_cl = sum(cl_list) / n
        
        cov = sum((gc_list[i] - mean_gc) * (cl_list[i] - mean_cl) for i in range(n))
        var_cl = sum((cl_list[i] - mean_cl) ** 2 for i in range(n))
        
        if var_cl != 0:
            self._hedge_ratio = cov / var_cl

    def _check_signals(self):
        gc_price = self.securities[self._gc.symbol].close
        cl_price = self.securities[self._cl.symbol].close
        
        if gc_price == 0 or cl_price == 0:
            return
            
        ratio = gc_price / cl_price
        
        self._gc_prices.add(gc_price)
        self._cl_prices.add(cl_price)
        self._ratio_window.add(ratio)
        
        if self._ratio_window.count < 30:
            return
            
        ratios = list(self._ratio_window)[::-1]
        mean_ratio = sum(ratios) / len(ratios)
        variance = sum((r - mean_ratio) ** 2 for r in ratios) / len(ratios)
        std_ratio = variance ** 0.5
        
        if std_ratio == 0:
            return
            
        z_score = (ratio - mean_ratio) / std_ratio
        
        # Entry logic
        if self._state == 0:
            if z_score > 2:
                # Ratio is high, short GC long CL (short the ratio)
                self._state = -1
                self._enter_positions()
            elif z_score < -2:
                # Ratio is low, long GC short CL (long the ratio)
                self._state = 1
                self._enter_positions()
        else:
            # Exit when z crosses 0
            if (self._state == 1 and z_score >= 0) or (self._state == -1 and z_score <= 0):
                self._liquidate_all()
                self._state = 0

    def _liquidate_all(self):
        for symbol in [self._gc.mapped, self._cl.mapped]:
            sec = self.securities[symbol]
            if sec.is_tradable and sec.price > 0 and self.portfolio[symbol].holdings.quantity != 0:
                self.liquidate(symbol)

    def _enter_positions(self):
        gc_symbol = self._gc.mapped
        cl_symbol = self._cl.mapped
        
        gc_sec = self.securities[gc_symbol]
        cl_sec = self.securities[cl_symbol]
        
        if not gc_sec.is_tradable or gc_sec.price <= 0 or not cl_sec.is_tradable or cl_sec.price <= 0:
            return
        
        gc_target = 0.5 * self._state
        cl_target = -0.5 * self._state * self._hedge_ratio
        
        gross = abs(gc_target) + abs(cl_target)
        if gross > 0:
            gc_target = gc_target / gross
            cl_target = cl_target / gross
        
        gc_qty = self.calculate_order_quantity(gc_symbol, gc_target)
        cl_qty = self.calculate_order_quantity(cl_symbol, cl_target)
        
        if gc_qty != 0:
            self.market_order(gc_symbol, gc_qty)
        if cl_qty != 0:
            self.market_order(cl_symbol, cl_qty)

    def on_data(self, data: Slice):
        pass