# region imports
from AlgorithmImports import *
# endregion

class PensiveYellowGreenHorse(QCAlgorithm):

    def initialize(self):
        self.set_start_date(2023, 1, 1)
        self.set_cash(1000000)
        self.settings.seed_initial_prices = True
        
        self._target_dte = 45
        self._close_dte = 15
        self._min_delta = 0.03
        self._max_delta = 0.07
        self._contracts_per_asset = 2
        self._num_holdings = 15
        
        self._open_positions = {}
        self._selected_symbols = []
        
        self._spy = self.add_equity("SPY", Resolution.DAILY).symbol
        self._universe = self.add_universe(self._select_constituents)
        
        self.schedule.on(self.date_rules.month_start(self._spy), 
                        self.time_rules.after_market_open(self._spy, 30), 
                        self._rebalance)
        
        self.schedule.on(self.date_rules.every_day(self._spy),
                        self.time_rules.after_market_open(self._spy, 30),
                        self._manage_positions)
    
    def _select_constituents(self, fundamentals: List[Fundamental]) -> List[Symbol]:
        spy_constituents = [f for f in fundamentals 
                           if f.has_fundamental_data 
                           and f.asset_classification.morningstar_sector_code != 0
                           and f.market_cap > 0]
        
        if not spy_constituents:
            return []
        
        sorted_by_volume = sorted(spy_constituents, 
                                 key=lambda f: f.dollar_volume, 
                                 reverse=True)
        
        self._selected_symbols = [f.symbol for f in sorted_by_volume[:self._num_holdings]]
        
        return self._selected_symbols
    
    def on_securities_changed(self, changes: SecurityChanges) -> None:
        for security in changes.added_securities:
            if security.symbol in self._selected_symbols:
                security.set_data_normalization_mode(DataNormalizationMode.RAW)
    
    def _rebalance(self) -> None:
        if not self._selected_symbols:
            return
        
        for symbol in self._selected_symbols:
            self._open_put_position(symbol)
    

    
    def _open_put_position(self, underlying: Symbol) -> None:
        chain = self.option_chain_provider.get_option_contract_list(underlying, self.time)
        if not chain:
            return
        
        puts = [x for x in chain if x.id.option_right == OptionRight.PUT]
        if not puts:
            return
        
        target_expiry = self.time + timedelta(days=self._target_dte)
        puts = sorted(puts, key=lambda x: abs((x.id.date - target_expiry).days))
        
        if not puts:
            return
        
        expiry_group = [x for x in puts if x.id.date == puts[0].id.date]
        
        underlying_price = self.securities[underlying].price
        if underlying_price == 0:
            return
        
        valid_contracts = []
        for contract in expiry_group:
            strike = contract.id.strike_price
            if strike < underlying_price:
                valid_contracts.append(contract)
        
        if not valid_contracts:
            return
        
        valid_contracts = sorted(valid_contracts, key=lambda x: x.id.strike_price, reverse=True)
        
        target_strike_index = int(len(valid_contracts) * 0.10)
        if target_strike_index >= len(valid_contracts):
            target_strike_index = len(valid_contracts) - 1
        
        selected_contract = valid_contracts[target_strike_index]
        
        if selected_contract in [pos['symbol'] for pos in self._open_positions.values()]:
            return
        
        self.add_option_contract(selected_contract, Resolution.DAILY)
        
        ticket = self.market_order(selected_contract, -self._contracts_per_asset)
        
        self._open_positions[selected_contract] = {
            'symbol': selected_contract,
            'expiry': selected_contract.id.date,
            'underlying': underlying,
            'ticket': ticket
        }
        
        self.debug(f"Sold {self._contracts_per_asset} puts on {underlying.value}: {selected_contract.value}, Strike: {selected_contract.id.strike_price}, Expiry: {selected_contract.id.date}")
    
    def _manage_positions(self) -> None:
        positions_to_close = []
        
        for option_symbol, position_info in self._open_positions.items():
            days_to_expiry = (position_info['expiry'].date() - self.time.date()).days
            
            if days_to_expiry <= self._close_dte:
                positions_to_close.append(option_symbol)
        
        for option_symbol in positions_to_close:
            if self.portfolio[option_symbol].invested:
                self.liquidate(option_symbol)
                self.debug(f"Closed position: {option_symbol.value}")
            
            del self._open_positions[option_symbol]