from AlgorithmImports import *
from datetime import datetime
import math
from scipy.stats import kurtosis

class TwoUniverseKurtosisIronCondorStrategy(QCAlgorithm):

    def Initialize(self):
        self.SetStartDate(2022, 5, 13)
        self.SetEndDate(2022, 6, 20)
        self.SetCash(2000000)
        self.SetTimeZone(TimeZones.NewYork)

        # Add SPX index
        self.index = self.AddIndex("SPX")

        # Universe 1 (option1): Wide filter for kurtosis calculations
        self.option1 = self.AddIndexOption(self.index.Symbol, "SPXW")
        self.option1.SetFilter(lambda universe: universe.IncludeWeeklys().Strikes(-30,30).Expiration(0, 0))
        self._symbol1 = self.option1.Symbol

        # Universe 2 (option2): Iron Condor filter for placing trades
        self.option2 = self.AddIndexOption(self.index.Symbol, "SPXW")
        self.option2.SetFilter(lambda x: x.IncludeWeeklys().IronCondor(0, 20, 40))
        self._symbol2 = self.option2.Symbol

        # Risk and trade management parameters
        self.max_portfolio_risk = 0.20
        self.profit_target = 1.5
        self.stop_loss = 0.75
        self.trade_open = False
        self.initial_credit = 0
        self.max_potential_loss = 0
        self.target_delta = 0.2

        self.kurtosis_threshold = 3
        self.current_date = None
        self.kurtosis_condition_met = False
        self.computed_kurtosis_today = False

    def OnData(self, slice):
        # Check if a new day has started
        if self.current_date != self.Time.date():
            self.current_date = self.Time.date()
            self.trade_open = False
            self.kurtosis_condition_met = False
            self.computed_kurtosis_today = False

        # Compute kurtosis from option1 (broad filter)
        if not self.computed_kurtosis_today:
            chain1 = slice.OptionChains.get(self._symbol1)
            if chain1:
                iv_values = [x.ImpliedVolatility for x in chain1 if x.ImpliedVolatility and 0 < x.ImpliedVolatility < 5]
                if len(iv_values) > 3:
                    daily_kurtosis = kurtosis(iv_values)
                    if daily_kurtosis > self.kurtosis_threshold:
                        self.Debug(f"{self.Time} Kurtosis={daily_kurtosis:.2f} condition met.")
                        self.kurtosis_condition_met = True
                    else:
                        self.Debug(f"{self.Time} Kurtosis={daily_kurtosis:.2f} condition not met.")
                    self.computed_kurtosis_today = True

        # Manage open position
        if self.trade_open:
            self.CheckPositionManagement()

        # If kurtosis is met and not invested, try to open Iron Condor from option2
        if not self.Portfolio.Invested and self.kurtosis_condition_met:
            self.OpenIronCondor(slice)

    def OpenIronCondor(self, slice):
        chain2 = slice.OptionChains.get(self._symbol2)
        if not chain2:
            return

        # Select the nearest expiry from the option2 chain
        expiries = [x.Expiry for x in chain2]
        if len(expiries) == 0:
            return
        expiry = max(expiries)

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

        put_contracts = [x for x in chain2 if x.Right == OptionRight.PUT and x.Greeks and abs(x.Greeks.Delta) <= self.target_delta]
        call_contracts = [x for x in chain2 if x.Right == OptionRight.CALL and x.Greeks 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_candidates = [x for x in call_contracts if x.Strike > near_call.Strike]
        if not far_call_candidates:
            return
        far_call = min(far_call_candidates, 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_candidates = [x for x in put_contracts if x.Strike < near_put.Strike]
        if not far_put_candidates:
            return
        far_put = min(far_put_candidates, 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._symbol2,
            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