from AlgorithmImports import *
from datetime import datetime

class FedAnnouncementIronCondorStrategy(QCAlgorithm):

    def Initialize(self):
        self.SetStartDate(2022, 5, 13)
        self.SetEndDate(2024, 12, 1)
        self.SetCash(2000000)
        self.SetTimeZone(TimeZones.NewYork)

        # Add SPX index and options
        self.index = self.AddIndex("SPX")
        
        # Add SPXW options with iron condor filter
        option = self.AddIndexOption(self.index.Symbol, "SPXW")
        option.SetFilter(lambda x: x.IncludeWeeklys().IronCondor(0, 20, 40))
        self._symbol = option.Symbol

        # Risk management parameters
        self.max_portfolio_risk = 0.20  # 20%! max risk per trade
        self.profit_target = 1.5 # 200% of initial credit
        self.stop_loss = 0.75  # 100% of max potential loss
        self.trade_open = False
        self.initial_credit = 0
        self.max_potential_loss = 0
        self.target_delta = 0.2  # Target delta for option selection

        self.announcement_dates = [
            datetime(2022, 5, 16), datetime(2022, 5, 18), datetime(2022, 5, 19), datetime(2022, 5, 20),
            datetime(2022, 5, 23), datetime(2022, 5, 25), datetime(2022, 5, 26), datetime(2022, 5, 31),
            datetime(2022, 6, 3), datetime(2022, 6, 7), datetime(2022, 6, 8), datetime(2022, 6, 9),
            datetime(2022, 6, 14), datetime(2022, 6, 15), datetime(2022, 6, 16), datetime(2022, 6, 17),
            datetime(2022, 6, 27), datetime(2022, 6, 28), datetime(2022, 6, 29), datetime(2022, 6, 30),
            datetime(2022, 7, 6), datetime(2022, 7, 20), datetime(2022, 7, 21), datetime(2022, 7, 28),
            datetime(2022, 8, 1), datetime(2022, 8, 11), datetime(2022, 8, 15), datetime(2022, 8, 16),
            datetime(2022, 9, 13), datetime(2022, 9, 14), datetime(2022, 9, 16), datetime(2022, 9, 21),
            datetime(2022, 9, 28), datetime(2022, 9, 29), datetime(2022, 10, 5), datetime(2022, 10, 7),
            datetime(2022, 10, 12), datetime(2022, 10, 14), datetime(2022, 10, 19), datetime(2022, 10, 24),
            datetime(2022, 10, 25), datetime(2022, 11, 2), datetime(2022, 11, 9), datetime(2022, 11, 23),
            datetime(2022, 11, 28), datetime(2022, 12, 1), datetime(2022, 12, 5), datetime(2022, 12, 7),
            datetime(2022, 12, 14), datetime(2023, 1, 12), datetime(2023, 1, 25), datetime(2023, 1, 30),
            datetime(2023, 2, 1), datetime(2023, 2, 14), datetime(2023, 3, 13), datetime(2023, 3, 15),
            datetime(2023, 3, 16), datetime(2023, 3, 22), datetime(2023, 3, 30), datetime(2023, 4, 3),
            datetime(2023, 4, 4), datetime(2023, 4, 19), datetime(2023, 5, 1), datetime(2023, 6, 5),
            datetime(2023, 6, 12), datetime(2023, 6, 14), datetime(2023, 7, 19), datetime(2023, 10, 23),
        ]
        
        # self.announcement_dates = [
        #     datetime(2022, 5, 19), datetime(2022, 7, 11), datetime(2022, 7, 20), datetime(2022, 8, 2),
        #     datetime(2022, 8, 23), datetime(2022, 9, 5), datetime(2022, 9, 9), datetime(2022, 9, 20),
        #     datetime(2022, 10, 6), datetime(2022, 10, 11), datetime(2022, 10, 17), datetime(2022, 10, 19),
        #     datetime(2022, 10, 25), datetime(2022, 11, 24), datetime(2022, 12, 1), datetime(2022, 12, 19),
        #     datetime(2023, 1, 2), datetime(2023, 1, 4), datetime(2023, 1, 9), datetime(2023, 1, 10),
        #     datetime(2023, 1, 16), datetime(2023, 2, 27), datetime(2023, 2, 28), datetime(2023, 5, 11),
        #     datetime(2023, 5, 18), datetime(2023, 5, 19), datetime(2023, 5, 26), datetime(2023, 6, 7),
        #     datetime(2023, 6, 9), datetime(2023, 7, 4), datetime(2023, 8, 3), datetime(2023, 8, 8),
        #     datetime(2023, 8, 21), datetime(2023, 9, 8), datetime(2023, 9, 11), datetime(2023, 9, 27),
        #     datetime(2023, 9, 29), datetime(2023, 10, 16), datetime(2023, 10, 19), datetime(2023, 12, 4),
        #     datetime(2023, 12, 13), datetime(2023, 12, 28), datetime(2024, 1, 2), datetime(2024, 1, 9),
        #     datetime(2024, 2, 8), datetime(2024, 2, 16), datetime(2024, 3, 5), datetime(2024, 3, 21),
        #     datetime(2024, 3, 25), datetime(2024, 3, 26), datetime(2024, 4, 11), datetime(2024, 4, 30),
        #     datetime(2024, 5, 14), datetime(2024, 5, 31), datetime(2024, 6, 17), datetime(2024, 6, 24),
        #     datetime(2024, 7, 4), datetime(2024, 7, 16), datetime(2024, 7, 18), datetime(2024, 11, 11)
        # ]


    def OnData(self, slice):
        # Check positions for management
        if self.trade_open:
            self.CheckPositionManagement()
        
        # Open new positions only on announcement dates
        if self.Portfolio.Invested or not self.Time.date() in [date.date() for date in self.announcement_dates]:
            return

        # Option chain handling and trade opening logic
        chain = slice.OptionChains.get(self._symbol)
        if not chain:
            return

        expiry = max([x.Expiry for x in chain])
        chain = sorted([x for x in chain if x.Expiry == expiry], key=lambda x: x.Strike)

        put_contracts = [x for x in chain if x.Right == OptionRight.PUT and abs(x.Greeks.Delta) <= self.target_delta]
        call_contracts = [x for x in chain if x.Right == OptionRight.CALL and abs(x.Greeks.Delta) <= self.target_delta]

        if len(call_contracts) < 2 or len(put_contracts) < 2:
            return

        near_call = min(call_contracts, key=lambda x: abs(x.Greeks.Delta - self.target_delta))
        far_call = min([x for x in call_contracts if x.Strike > near_call.Strike], key=lambda x: abs(x.Greeks.Delta - self.target_delta))
        
        near_put = min(put_contracts, key=lambda x: abs(x.Greeks.Delta + self.target_delta))
        far_put = min([x for x in put_contracts if x.Strike < near_put.Strike], key=lambda x: abs(x.Greeks.Delta + self.target_delta))

        credit = (near_call.BidPrice - far_call.AskPrice) + (near_put.BidPrice - far_put.AskPrice)
        spread_width = max(far_call.Strike - near_call.Strike, near_put.Strike - far_put.Strike)
        max_potential_loss = spread_width * 100 - credit * 100

        total_portfolio_value = self.Portfolio.TotalPortfolioValue
        max_trade_risk = total_portfolio_value * self.max_portfolio_risk
        contracts = int(max_trade_risk / max_potential_loss)

        if contracts == 0:
            return

        iron_condor = OptionStrategies.IronCondor(
            self._symbol, 
            far_put.Strike,
            near_put.Strike,
            near_call.Strike,
            far_call.Strike,
            expiry)

        self.Buy(iron_condor, contracts)
        self.initial_credit = credit * 100 * contracts
        self.max_potential_loss = max_potential_loss * contracts
        self.trade_open = True
        self.Debug(f"Opened iron condor at {self.Time}, Contracts: {contracts}, Credit: ${self.initial_credit:.2f}")

    def CheckPositionManagement(self):
        total_pnl = sum([holding.UnrealizedProfit for holding in self.Portfolio.Values if holding.Invested])

        if total_pnl >= self.initial_credit * self.profit_target:
            self.Liquidate()
            self.Debug(f"Closed position at profit target on {self.Time}")
            self.trade_open = False
        elif total_pnl <= -self.max_potential_loss * self.stop_loss:
            self.Liquidate()
            self.Debug(f"Closed position at stop loss on {self.Time}")
            self.trade_open = False