# from datetime import timedelta
# from QuantConnect import *
# from QuantConnect.Algorithm.Framework.Selection import FutureUniverseSelectionModel
# from QuantConnect.Data.UniverseSelection import *
# from QuantConnect.Securities.Future import FutureFilterUniverse
# from QuantConnect.Securities import MarketHoursDatabase

from AlgorithmImports import * 

class OpenInterestFutureUniverseSelectionModel(FutureUniverseSelectionModel):
    """
    This model selects futures contracts based on open interest. It extends the FutureUniverseSelectionModel
    and applies an additional filter to sort the contracts by open interest.

    Attributes:
        algorithm (QCAlgorithm): The algorithm instance using this selection model.
        future_chain_symbol_selector (Callable[[DateTime], List[Symbol]]): Function to select futures symbols.
        _chain_contracts_lookup_limit (int): The maximum number of contracts to inspect for open interest.
        _results_limit (int): The maximum number of contracts to return in the universe.
        _market_hours_database (MarketHoursDatabase): Provides exchange hours needed for history requests.
    """

    def __init__(self, algorithm, future_chain_symbol_selector, chain_contracts_lookup_limit=6, results_limit=1):
        """
        Creates a new instance of the model.

        Args:
            algorithm (QCAlgorithm): The algorithm instance.
            future_chain_symbol_selector (Callable[[DateTime], List[Symbol]]): Selector function for futures symbols.
            chain_contracts_lookup_limit (int): Maximum number of contracts to query for open interest.
            results_limit (int): Maximum number of contracts that will be part of the universe.
        """
        super().__init__(timedelta(days=1), future_chain_symbol_selector)
        self.algorithm = algorithm
        self._chain_contracts_lookup_limit = chain_contracts_lookup_limit
        self._results_limit = results_limit
        self._market_hours_database = MarketHoursDatabase.FromDataFolder()

    def Filter(self, filter: FutureFilterUniverse) -> FutureFilterUniverse:
        """
        Overrides the base method to filter futures contracts based on open interest.

        Args:
            filter (FutureFilterUniverse): The future filter universe to apply the filtering.

        Returns:
            FutureFilterUniverse: The filtered universe of futures contracts.
        """
        contracts = filter.ToDictionary(lambda x: x.Symbol, lambda x: self._market_hours_database.GetEntry(x.Symbol.ID.Market, x.Symbol, x.Symbol.ID.SecurityType))
        return filter.Contracts(self.FilterByOpenInterest(contracts))

    def FilterByOpenInterest(self, contracts):
        """
        Filters the provided contracts by open interest.

        Args:
            contracts (Dictionary[Symbol, MarketHoursDatabase.Entry]): Contracts to filter.

        Returns:
            List[Symbol]: The list of symbols sorted by open interest and date.
        """
        sorted_contracts = sorted(contracts.keys(), key=lambda x: x.ID.Date)
        if self._chain_contracts_lookup_limit is not None:
            sorted_contracts = sorted_contracts[:self._chain_contracts_lookup_limit]
        
        open_interest_data = self.GetOpenInterest(contracts, sorted_contracts)
        filtered_contracts = sorted(open_interest_data.items(), key=lambda x: (-x[1], x[0].ID.Date))

        if self._results_limit is not None:
            filtered_contracts = filtered_contracts[:self._results_limit]
        
        return [symbol for symbol, _ in filtered_contracts]

    def GetOpenInterest(self, contracts, symbols):
        """
        Retrieves and processes historical open interest data for the given symbols.

        Args:
            contracts (Dictionary[Symbol, MarketHoursDatabase.Entry]): Contracts for which to retrieve open interest.
            symbols (List[Symbol]): The list of symbols to get historical data for.

        Returns:
            Dictionary[Symbol, float]: A dictionary mapping each symbol to its latest open interest value.
        """
        current_utc_time = self.algorithm.UtcTime
        start_time = current_utc_time - timedelta(days=1)
        
        # Request historical open interest data
        history = self.algorithm.History(symbols, start_time, current_utc_time, Resolution.Daily)
        open_interest_data = {}

        for symbol in symbols:
            try:
                # Access the open interest data for the symbol
                symbol_data = history.get(symbol)
                if symbol_data.empty:
                    continue
                open_interest_data[symbol] = symbol_data['openinterest'][-1]
            except KeyError:
                # If data is missing, log a message and skip the symbol
                self.algorithm.Debug(f'No open interest data for {symbol.Value}')
        
        return open_interest_data

from AlgorithmImports import *

# Define a custom universe selection model that inherits from OpenInterestFutureUniverseSelectionModel.
# This model will select futures contracts based on open interest.
class CustomFutureUniverseSelectionModel(OpenInterestFutureUniverseSelectionModel):
    """
    A universe selection model that selects futures contracts based on open interest
    and inherits the behavior of the OpenInterestFutureUniverseSelectionModel.
    """
    def __init__(self, algorithm, chainContractsLookupLimit=6, resultsLimit=1):
        """
        Initialize the custom universe selection model with the algorithm instance and limits for contracts lookup.

        :param algorithm: The algorithm instance.
        :param chainContractsLookupLimit: The maximum number of contracts to consider based on open interest.
        :param resultsLimit: The maximum number of contracts to include in the universe.
        """
        super().__init__(algorithm, self.select_future_chain_symbols, chainContractsLookupLimit, resultsLimit)

    def select_future_chain_symbols(self, utcTime: datetime) -> List[Symbol]:
        """
        Selects the future symbols to be included in the universe.

        :param utcTime: The current UTC time as a datetime object.
        :return: A list of Symbols representing the selected futures contracts.
        """
        return [
            Symbol.Create(Futures.Indices.SP500EMini, SecurityType.Future, Market.CME),
            Symbol.Create(Futures.Metals.GOLD, SecurityType.Future, Market.COMEX)
        ]

from AlgorithmImports import * 
from Universe import CustomFutureUniverseSelectionModel

# Main algorithm class that uses the CustomFutureUniverseSelectionModel for universe selection.
class OpenInterestFutureUniverseSelectionAlgorithm(QCAlgorithm):
    """
    This algorithm uses a custom future universe selection model to select futures contracts based on open interest.
    It logs the contract data and rollover events for analysis.
    """
    def Initialize(self):
        """
        Initialize the algorithm settings, including start date, cash, universe settings, and universe selection model.
        """
        self.SetStartDate(2023, 1, 1)
        self.SetEndDate(2023,2,1)  # Start date for the backtest
        self.SetCash(100000)           # Starting cash for the algorithm

        # Set properties of the universe settings
        self.UniverseSettings.Resolution = Resolution.MINUTE
        self.UniverseSettings.Leverage = 1.0
        self.UniverseSettings.FillForward = True
        self.UniverseSettings.ExtendedMarketHours = True
        self.UniverseSettings.DataNormalizationMode = DataNormalizationMode.Raw

        # Add the custom universe selection model to the algorithm
        self.AddUniverseSelection(CustomFutureUniverseSelectionModel(self))

        # Keep track of active contracts to detect rollovers
        self.previous_active_contracts = set()

    def OnData(self, slice: Slice):
        """
        Event handler for new data.
        
        :param slice: Slice object containing the data for the current time step.
        """
        # Log the current time slice
        self.Log(f"OnData for time slice: {self.Time}")

        # Process data for each futures chain in the slice
        if slice.FutureChains:
            for chain in slice.FutureChains:
                # Log the underlying Symbol for the current futures chain
                self.Log(f"Processing future chain for underlying: {chain.Key.Value}")

                # Iterate through contracts in the current chain
                for contract in chain.Value:
                    # Log contract details including price and open interest
                    self.Log(f"Contract: {contract.Symbol.Value} Price: {contract.LastPrice} OpenInterest: {contract.OpenInterest}")

                    # Detect and log new active contracts based on open interest
                    if contract.Symbol not in self.previous_active_contracts and contract.OpenInterest > 0:
                        self.Log(f"New active contract detected based on open interest: {contract.Symbol.Value}")
                        # Update the active contracts list
                        self.previous_active_contracts.clear()
                        self.previous_active_contracts.add(contract.Symbol)

    def OnSecuritiesChanged(self, changes: SecurityChanges):
        """
        Event handler for changes in the securities in the universe, such as additions and removals.
        
        :param changes: SecurityChanges object containing lists of added and removed securities.
        """
        # Log securities that have been added to the universe
        for security in changes.AddedSecurities:
            self.Log(f"Added security: {security.Symbol.Value}")

        # Log securities that have been removed from the universe
        for security in changes.RemovedSecurities:
            self.Log(f"Removed security: {security.Symbol.Value}")

from AlgorithmImports import *

class OpenInterestFutureUniverseSelectionAlgorithm(QCAlgorithm):
    """
    Algorithm using the OpenInterestFutureUniverseSelectionModel to select futures contracts based on open interest
    and prints out data and rollover events.
    """
    def Initialize(self):
        self.SetStartDate(2023, 1, 1)  # Set the start date
        # self.SetEndDate(2023, 6, 1)    # Set the end date
        self.SetCash(100000)           # Set the starting cash

        # Set UniverseSettings properties as required
        self.UniverseSettings.Resolution = Resolution.Daily 
        self.UniverseSettings.Leverage = 1.0
        self.UniverseSettings.FillForward = True
        self.UniverseSettings.ExtendedMarketHours = True
        self.UniverseSettings.DataNormalizationMode = DataNormalizationMode.Raw

        # Add OpenInterestFutureUniverseSelectionModel with the selector function for the futures chains
        self.AddUniverseSelection(
            OpenInterestFutureUniverseSelectionModel(self, self.FutureChainSymbolSelector, chainContractsLookupLimit=6, resultsLimit=1)
        )

        # Keep track of the active contracts and rollover
        self.previous_active_contracts = set()

    def FutureChainSymbolSelector(self, utc_time: datetime) -> List[Symbol]:
        # This function returns the future symbols to be included in the universe
        return [
            Symbol.Create(Futures.Indices.SP500EMini, SecurityType.Future, Market.CME),
            Symbol.Create(Futures.Metals.GOLD, SecurityType.Future, Market.COMEX)
        ]

    def OnData(self, slice: Slice):
        # Log the time of the data slice
        self.Log(f"OnData for time slice: {self.Time}")

        if slice.FutureChains:
            for chain in slice.FutureChains:
                # Use 'chain.Key' to get the underlying Symbol for the futures chain
                self.Log(f"Processing future chain for {chain.Key}")

                for contract in chain.Value:
                    # Log the contract Symbol, Price, and OpenInterest for each contract in the chain
                    self.Log(f"Contract: {contract.Symbol.Value} Price: {contract.LastPrice} OpenInterest: {contract.OpenInterest}")

                    # Check if this is a newly active contract based on open interest
                    if contract.Symbol not in self.previous_active_contracts and contract.OpenInterest > 0:
                        # Log the new active contract
                        self.Log(f"New active contract detected based on open interest: {contract.Symbol.Value}")
                        # Clear the previous contracts and add this contract as the currently active one
                        self.previous_active_contracts.clear()
                        self.previous_active_contracts.add(contract.Symbol)

    def OnSecuritiesChanged(self, changes: SecurityChanges):
        # This method will be called when the active future contract changes, i.e., a rollover event
        for security in changes.AddedSecurities:
            self.Log(f"Added security: {security.Symbol.Value}")

        for security in changes.RemovedSecurities:
            self.Log(f"Removed security: {security.Symbol.Value}")