from AlgorithmImports import *
import datetime

class HighIVCoveredCallComposite(QCAlgorithm):
    def Initialize(self):
        self.SetStartDate(2024, 1, 1)
        self.SetEndDate(2024, 12, 31)
        self.SetCash(400000)
        self.high_iv_stock_count = 10
        self.selected_symbols = []
        self.selected_contracts = {}
        self.contract_indicators = {}
        self.spy = self.AddEquity("SPY", Resolution.Daily, dataNormalizationMode=DataNormalizationMode.Raw).Symbol
        self.delta_flags = {}  # Dictionary to track delta flags per contract
        self.new_contract_flags = {}  # Track new contracts per symbol
        self.processed_symbols = []
        self.earnings_dict = {}
        self.collect_earnings_info_window = False
        self.earnings_data = self.AddData(EODHDUpcomingEarnings, "earnings").Symbol

                # Schedule earnings collection window
        self.Schedule.On(self.DateRules.WeekEnd(self.spy, -1),
                         self.TimeRules.At(23, 59),
                         self.ActivateEarningsCollection)
        self.Schedule.On(self.DateRules.WeekEnd(self.spy),
                         self.TimeRules.At(0, 1),
                         self.DeactivateEarningsCollection)

        # Universe selection settings
        self.UniverseSettings.Resolution = Resolution.Daily
        self.UniverseSettings.DataNormalizationMode = DataNormalizationMode.Raw

        # Seeder and fill forward setting
        seeder = FuncSecuritySeeder(self.get_last_known_prices)
        self.SetSecurityInitializer(BrokerageModelSecurityInitializer(self.brokerage_model, seeder))
        self.fill_forward = True

        # Schedule quarterly equity selection on the first trading day of each quarter
        self.Schedule.On(self.DateRules.MonthStart(self.spy, 1),
                         self.TimeRules.AfterMarketOpen(self.spy, 1),
                         self.SelectSymbols)
        
        # Schedule for preparing covered calls on Fridays at 3 pm
        self.Schedule.On(self.DateRules.WeekEnd(self.spy),
                         self.TimeRules.At(15, 0),
                         self.SelectCoveredCalls)

        # Define sell window time range (15:01 - 15:58)
        self.sell_window_open = datetime.time(15, 1)
        self.sell_window_close = datetime.time(15, 58)

        # Universe selection functions
        self.AddUniverse(self.CoarseSelectionFunction, self.FineSelectionFunction)

    def OnSecuritiesChanged(self, changes):
        # Iterate over newly added securities
        for added_security in changes.AddedSecurities:
            symbol = added_security.Symbol
            # Add earnings data for the new symbol
            self.earnings_data = self.AddData(EODHDUpcomingEarnings, "earnings").Symbol
            # Store the earnings data symbol for future reference
     
    def ActivateEarningsCollection(self):
        self.collect_earnings_info_window = True

    def DeactivateEarningsCollection(self):
        self.collect_earnings_info_window = False


    def CoarseSelectionFunction(self, coarse):
        filtered = [x for x in coarse if x.HasFundamentalData and x.Price > 5 and x.DollarVolume > 1e7]
        sortedByDollarVolume = sorted(filtered, key=lambda x: x.DollarVolume, reverse=True)
        return [x.Symbol for x in sortedByDollarVolume[:20]]

    def FineSelectionFunction(self, fine):
        symbols_metrics = {}
        benchmark_history = self.History(self.spy, 30, Resolution.Daily)
        benchmark_returns = benchmark_history['close'].pct_change().dropna() if not benchmark_history.empty else None

        for f in fine:
            history = self.History(f.Symbol, 30, Resolution.Daily)
            if history.empty or benchmark_returns is None:
                continue

            # Calculate daily returns for the symbol
            daily_returns = history['close'].pct_change().dropna()
            if daily_returns.empty:
                continue

            # Calculate historical volatility (standard deviation of returns)
            hist_volatility = daily_returns.std() * (252 ** 0.5)  # Annualized volatility

            # Calculate Sharpe Ratio
            excess_returns = daily_returns.mean() - (benchmark_returns.mean() if not benchmark_returns.empty else 0)
            sharpe_ratio = excess_returns / hist_volatility if hist_volatility > 0 else 0

            # Combine historical volatility and Sharpe Ratio for a composite score
            composite_score = hist_volatility * 0.7 + sharpe_ratio * 0.3  # Adjust weights as needed
            symbols_metrics[f.Symbol] = composite_score

        # Sort symbols by composite score and select the top N
        sorted_symbols = sorted(symbols_metrics.items(), key=lambda x: x[1], reverse=True)
        self.selected_symbols = [kv[0] for kv in sorted_symbols[:self.high_iv_stock_count]]

        return self.selected_symbols


    def SelectSymbols(self):
        self.selected_contracts.clear()
        for symbol in self.selected_symbols:
            price = self.Securities[symbol].Price
            required_cash = price * (100 - self.Portfolio[symbol].Quantity)
            if self.Portfolio.Cash >= required_cash:
                if self.Portfolio[symbol].Quantity < 100:
                    self.MarketOrder(symbol, 100 - self.Portfolio[symbol].Quantity)
           # else:
                #self.Debug(f"Not enough cash to buy 100 shares of {symbol}. Required: {required_cash}, Available: {self.Portfolio.Cash}")

    def SelectCoveredCalls(self):
        days_until_friday = (4 - self.Time.weekday()) % 7
        if days_until_friday == 0:
            days_until_friday = 7
        next_friday = (self.Time + timedelta(days=days_until_friday)).date()

        for symbol in self.selected_symbols:




            # Check if any contract for this symbol is already invested
            if any(self.Portfolio[contract].Invested for contract in self.selected_contracts if contract.Underlying == symbol):
                continue

            # Get the option chain for the symbol
            chain = self.OptionChain(symbol).DataFrame
            if chain.empty:
                #self.Debug(f"No options in chain for {symbol} at {self.Time}")
                continue

            # Filter contracts expiring next Friday with delta between 0.25 and 0.35
            filtered_chain = chain[
                (chain.expiry.dt.date == next_friday) & 
                (chain.right == OptionRight.CALL) & 
                (chain.delta < 0.35) & 
                (chain.delta > 0.25)
            ]

            # Add each contract that meets the criteria
            if not filtered_chain.empty:
                for contract in filtered_chain.index:
                    self.AddOptionContract(contract, Resolution.Minute)
                    self.selected_contracts[contract] = False
                    self.contract_indicators[contract] = self.d(contract)
                    delta_value = self.contract_indicators[contract].Current.Value

                    #self.Debug(f"Added contract with delta {delta_value} for {symbol} expiring on {next_friday}")
           # else:
                #self.Debug(f"No suitable contracts expiring on {next_friday} for {symbol}")
    def MonitorDeltas(self, data: Slice):
        # Calculate the current week's Friday
        days_until_current_friday = (4 - self.Time.weekday()) % 7
        current_friday = (self.Time + timedelta(days=days_until_current_friday)).date()

        for contract, sold in list(self.selected_contracts.items()):
            if not sold:  # Skip unsold contracts
                continue

            # Skip if delta flag is raised for this contract
            if self.delta_flags.get(contract, False):
                continue

            # Ensure the contract data and indicators are available
            if contract in self.contract_indicators and data.ContainsKey(contract):
                delta_value = self.contract_indicators[contract].Current.Value

                # Check if delta exceeds the threshold
                if delta_value > 0.40:
                    underlying_symbol = contract.Underlying
                    #self.Debug(f"Delta {delta_value:.2f} exceeds threshold for {underlying_symbol}")

                    if self.Portfolio[underlying_symbol].Quantity >= 100:
                        # Liquidate the current contract
                        open_orders = self.Transactions.GetOpenOrders(contract)
                        for order in open_orders:
                            if order.Type == OrderType.Limit:
                                self.Transactions.CancelOrder(order.Id)
                        self.Liquidate(contract)

                        # Get the option chain for the current week's Friday
                        chain = self.OptionChain(underlying_symbol).DataFrame
                        if not chain.empty:
                            # Filter chain for closest delta to 0.3
                            filtered_chain = chain[
                                (chain.expiry.dt.date == current_friday) & 
                                (chain.right == OptionRight.CALL)
                            ]
                            filtered_chain["delta_diff"] = abs(filtered_chain.delta - 0.30)
                            closest_to_03 = filtered_chain.sort_values("delta_diff").iloc[0]

                            # Select a new contract
                            new_contract = closest_to_03.name
                            self.AddOptionContract(new_contract, Resolution.Minute)

                            # Add to contract indicators and set delta flag
                            self.contract_indicators[new_contract] = self.d(new_contract)
                            self.delta_flags[contract] = True  # Raise flag for this contract
                            self.new_contract_flags[contract] = new_contract
                            #self.Debug(f"New contract added for {contract}: {new_contract}. Waiting for data.")
                            return


    def ProcessNewContract(self, data: Slice):
        for original_contract, new_contract in list(self.new_contract_flags.items()):
            if data.ContainsKey(new_contract):
                new_delta = self.contract_indicators[new_contract].Current.Value
                premium = self.Securities[new_contract].Price

                if premium >= 0.50 and new_delta < 0.40 and new_delta > 0:
                    # Place orders for the new contract
                    self.MarketOrder(new_contract, -1)
                    gtc_price = round(premium * 0.1, 2)
                    self.LimitOrder(new_contract, 1, gtc_price, tag="GTC Buy Order")

                    # Update selected contracts
                    self.selected_contracts[new_contract] = True
                    #self.Debug(f"Rolled contract to {new_contract} with delta {new_delta:.2f}")
                #else:
                    #self.Debug(f"Skipped rolling: Premium {premium:.2f} or delta {new_delta:.2f} not suitable.")

                # Reset flags for this contract
                self.delta_flags[original_contract] = False
                del self.new_contract_flags[original_contract]


    def on_data(self, data: Slice) -> None:
        current_time = self.Time.time()

        # Perform delta monitoring on Monday, Tuesday, and Wednesday
        #if self.Time.weekday() in [0, 1, 2]:  # Monday = 0, Tuesday = 1, Wednesday = 2
            #self.MonitorDeltas(data)

        # Process new contracts once data is available
        #self.ProcessNewContract(data)

        if self.collect_earnings_info_window:
            for symbol in self.selected_symbols:
                upcomings_earnings_for_symbol = data.get(EODHDUpcomingEarnings).get(symbol)
                if upcomings_earnings_for_symbol and upcomings_earnings_for_symbol.report_date <= self.Time + timedelta(days=10):
                    self.earnings_dict[symbol] = upcomings_earnings_for_symbol.report_date

                # Remove earnings info for past dates
        for symbol, report_date in list(self.earnings_dict.items()):
            if report_date < self.Time:
                del self.earnings_dict[symbol]



        if self.Time.weekday() == 4 and self.sell_window_open <= current_time <= self.sell_window_close:
            processed_contracts = []
        
            # Iterate over selected contracts by symbol
            for symbol in self.selected_symbols:
                if symbol in self.earnings_dict.keys() and self.earnings_dict[symbol] is not None:
                    #self.Debug(f"Skipping covered calls for {symbol} due to upcoming earnings on {self.earnings_dict[symbol]}")
                    continue  # Skip this symbol if earnings announcement is within 7 days
                selected_contract = self.SelectContractClosestToDelta(symbol)  # Filter contracts for delta closest to 0.30

                if selected_contract and not self.selected_contracts[selected_contract]:  # Ensure contract hasn't been sold
                    if data.ContainsKey(selected_contract) and data[selected_contract] is not None:
                        underlying_symbol = selected_contract.Underlying

                        if self.Portfolio[underlying_symbol].Quantity >= 100:
                            premium = data[selected_contract].Price
                            quantity_multiplier = 1#self.Portfolio[underlying_symbol].Quantity / 100
                            if premium >= 0.50:
                                delta_value = self.contract_indicators[selected_contract].Current.Value
                                self.MarketOrder(selected_contract, -1*quantity_multiplier, tag=f"{self.contract_indicators[selected_contract].Current.Value} {self.Portfolio[underlying_symbol].Quantity} {self.Portfolio[selected_contract].Quantity}")
                                gtc_price = round(premium * 0.1, 2)
                                self.LimitOrder(selected_contract, 1*quantity_multiplier, gtc_price, tag="GTC Buy Order")
                                processed_contracts.append(selected_contract)
                                self.processed_symbols.append(symbol)
        
            # Mark processed contracts as sold
            for contract in processed_contracts:
                self.selected_contracts[contract] = True
        else:
            self.processed_symbols.clear()
        # Clean up unsold contracts after the sell window closes
        if self.Time.weekday() == 4 and current_time > self.sell_window_close:
            unsold_contracts = [contract for contract, sold in self.selected_contracts.items() if not sold]
            for contract in unsold_contracts:
                self.selected_contracts.pop(contract)

    def SelectContractClosestToDelta(self, symbol, target_delta=0.30):
        # Filter for contracts in selected_contracts that match the specified symbol
        symbol_contracts = {
            contract: self.contract_indicators[contract].Current.Value
            for contract in self.selected_contracts
            if contract.Underlying == symbol
        }

        if not symbol_contracts:
            #self.Debug(f"No contracts available for {symbol} in selected contracts.")
            return None

        # Find the contract with the delta closest to target_delta (0.30)
        closest_contract = min(symbol_contracts, key=lambda c: abs(symbol_contracts[c] - target_delta))
        # self.Debug(f"Selected contract with delta {symbol_contracts[closest_contract]:.2f} closest to target delta {target_delta} for {symbol}")
        return closest_contract