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(2024, 12, 1)
        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.05
        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.25

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

    def OnData(self, slice):

        # Manage open position
        if self.trade_open:
            self.CheckPositionManagement()
            
        # 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), WITH ROLLING FILTER
        if not self.computed_kurtosis_today and self.time.hour == 9 and self.time.minute >= 31 and self.time.minute <= 36:
            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) > 10: #CHANGE HERE, 3 to 10
                    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

        # If kurtosis is met and not invested, try to open Iron Condor from option2
        if not self.Portfolio.Invested and self.kurtosis_condition_met and self.time.hour == 15 and self.time.minute >= 0 and self.time.minute <= 55:
            self.OpenIronCondor(slice)


#this is the same as the old code:

    def OpenIronCondor(self, slice):

        chain2 = slice.OptionChains.get(self._symbol2)

        if not chain2:
            return

        expiry = max([x.Expiry for x in chain2])
        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 abs(x.Greeks.Delta) <= self.target_delta]
        call_contracts = [x for x in chain2 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._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.Log(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.Log(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.Log(f"Closed position at stop loss on {self.Time}")
            self.trade_open = False