from AlgorithmImports import *
import datetime
import math

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 = {}
        self.new_contract_flags = {}
        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)

        self.sell_window_open = datetime.time(15, 1)
        self.sell_window_close = datetime.time(15, 58)

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

    def OnSecuritiesChanged(self, changes):
        for added_security in changes.AddedSecurities:
            symbol = added_security.Symbol
            # Add earnings data for the new symbol if needed
            self.earnings_data = self.AddData(EODHDUpcomingEarnings, "earnings").Symbol

    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):
        # Retrieve benchmark (SPY) data
        benchmark_history = self.History(self.spy, 30, Resolution.Daily)
        if benchmark_history.empty:
            # If no benchmark data, just return the top fine fundamentals
            return [f.Symbol for f in fine[:self.high_iv_stock_count]]

        benchmark_closes = benchmark_history['close']
        benchmark_returns = benchmark_closes.pct_change().dropna()
        benchmark_avg_return = benchmark_returns.mean()
        benchmark_cum_return = (benchmark_closes[-1] / benchmark_closes[0]) - 1 if len(benchmark_closes) > 1 else 0

        symbols_metrics = {}
        for f_obj in fine:
            history = self.History(f_obj.Symbol, 30, Resolution.Daily)
            if history.empty:
                continue

            closes = history['close']
            daily_returns = closes.pct_change().dropna()
            if daily_returns.empty:
                continue

            # Historical Volatility (annualized)
            hist_volatility = daily_returns.std() * math.sqrt(252) if len(daily_returns) > 1 else 0

            # Sharpe Ratio
            excess_returns = daily_returns.mean() - benchmark_avg_return
            sharpe_ratio = excess_returns / hist_volatility if hist_volatility > 0 else 0

            # Relative Strength: Compare stock returns to benchmark returns
            stock_cum_return = (closes[-1] / closes[0]) - 1 if len(closes) > 1 else 0
            # One simple RS metric: difference between stock and benchmark cumulative returns
            relative_strength = stock_cum_return - benchmark_cum_return

            # Momentum: Use cumulative return over the last 30 days
            momentum = stock_cum_return

            # Composite Score:
            # Adjust these weights based on preference:
            # For example:
            # Sharpe Ratio: 30%
            # Inverse Volatility: We actually want lower volatility, so we might consider 1/hist_vol, but let's keep it simple
            # Relative Strength: 25%
            # Momentum: 25%
            # Hist Vol can be included as a negative factor if we want lower volatility assets
            # For simplicity, let's just do a weighted combo:
            # We'll treat lower volatility as better by taking negative vol.
            w_sharpe = 0.40
            w_rel_str = 0.20
            w_momentum = 0.10
            w_vol = 0.30  # Negative weighting for volatility still applies

            # Lower volatility is better, so we use negative hist_volatility
            composite_score = (w_sharpe * sharpe_ratio) + (w_rel_str * relative_strength) + (w_momentum * momentum) + (w_vol * (-hist_volatility))

            symbols_metrics[f_obj.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)

    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

            chain = self.OptionChain(symbol, flatten=True).DataFrame
            if chain.empty:
                continue

            filtered_chain = chain[
                (chain.expiry.dt.date == next_friday) &
                (chain.right == OptionRight.CALL) &
                (chain.delta < 0.35) &
                (chain.delta > 0.25)
            ]

            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)

    def MonitorDeltas(self, data: Slice):
        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:
                continue
            if self.delta_flags.get(contract, False):
                continue
            if contract in self.contract_indicators and data.ContainsKey(contract):
                delta_value = self.contract_indicators[contract].Current.Value
                if delta_value > 0.40:
                    underlying_symbol = contract.Underlying
                    # 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)

                    chain = self.OptionChain(underlying_symbol).DataFrame
                    if not chain.empty:
                        filtered_chain = chain[
                            (chain.expiry.dt.date == current_friday) &
                            (chain.right == OptionRight.CALL)
                        ]
                        if not filtered_chain.empty:
                            filtered_chain["delta_diff"] = abs(filtered_chain.delta - 0.30)
                            closest_to_03 = filtered_chain.sort_values("delta_diff").iloc[0]
                            new_contract = closest_to_03.name
                            self.AddOptionContract(new_contract, Resolution.Minute)
                            self.contract_indicators[new_contract] = self.d(new_contract)
                            self.delta_flags[contract] = True
                            self.new_contract_flags[contract] = new_contract
                            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:
                    self.MarketOrder(new_contract, -1)
                    gtc_price = round(premium * 0.1, 2)
                    self.LimitOrder(new_contract, 1, gtc_price, tag="GTC Buy Order")
                    self.selected_contracts[new_contract] = True

                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()

        #if self.Time.weekday() in [0, 1, 2]:
        #    self.MonitorDeltas(data)
        #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

        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 = []
            for symbol in self.selected_symbols:
                if symbol in self.earnings_dict.keys() and self.earnings_dict[symbol] is not None:
                    continue
                selected_contract = self.SelectContractClosestToDelta(symbol)
                if selected_contract and not self.selected_contracts[selected_contract]:
                    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
                            if premium >= 0.50:
                                delta_value = self.contract_indicators[selected_contract].Current.Value
                                self.MarketOrder(selected_contract, -1*quantity_multiplier)
                                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)

            for contract in processed_contracts:
                self.selected_contracts[contract] = True
        else:
            self.processed_symbols.clear()

        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):
        symbol_contracts = {
            contract: self.contract_indicators[contract].Current.Value
            for contract in self.selected_contracts
            if contract.Underlying == symbol
        }

        if not symbol_contracts:
            return None

        closest_contract = min(symbol_contracts, key=lambda c: abs(symbol_contracts[c] - target_delta))
        return closest_contract