# region imports
from AlgorithmImports import *
import numpy as np
from datetime import timedelta
# endregion

class AssenagonDispersionAlgorithm(QCAlgorithm):

    def Initialize(self):
        # 1. Strategy Timelines
        self.SetStartDate(2021, 1, 1) 
        self.SetEndDate(2026, 1, 1)
        self.SetCash(1000000)

        # 2. WarmUp and Universe Settings
        self.SetWarmUp(30, Resolution.Daily)
        self.UniverseSettings.Resolution = Resolution.Minute
        self.AddUniverse(self.FundamentalSelectionFunction)
        
        # 3. Index Setup + Assign Price Model for Greeks
        self.index_symbol = self.AddEquity("SPY", Resolution.Minute).Symbol
        self.spy_option = self.AddOption("SPY", Resolution.Minute)
        self.spy_option.PriceModel = OptionPriceModels.BlackScholes()
        self.spy_option.SetFilter(-2, 2, 25, 45) 
        
        # 4. Strategy State
        self.basket_symbols = []
        self.target_vega = 1000  
        self.next_rebalance = self.Time
        self.entry_prices = {}   

        # 5. Monitoring Charts
        # These will appear in the "Custom Charts" tab of your backtest
        plot = Chart("Portfolio Risk")
        plot.AddSeries(Series("Net Delta", SeriesType.Line, 0))
        plot.AddSeries(Series("Net Vega", SeriesType.Line, 1))
        self.AddChart(plot)

        # 6. Daily Delta Hedge (At the Close)
        self.Schedule.On(self.DateRules.EveryDay(self.index_symbol), 
                         self.TimeRules.BeforeMarketClose(self.index_symbol, 10), 
                         self.DeltaHedge)

        self.set_brokerage_model(BrokerageName.ALPACA)

    def FundamentalSelectionFunction(self, fundamental):
        if self.Time < self.next_rebalance and len(self.basket_symbols) > 0:
            return self.basket_symbols
        # Financials (103), Energy (309), Not Tech (311)
        filtered = [x for x in fundamental if x.HasFundamentalData and x.Price > 15 
                    and x.AssetClassification.MorningstarSectorCode in [103, 309]
                    and x.AssetClassification.MorningstarSectorCode != 311
                    and x.ValuationRatios.PERatio > 0]
        sorted_by_value = sorted(filtered, key=lambda x: 1 / x.ValuationRatios.PERatio, reverse=True)
        self.basket_symbols = [x.Symbol for x in sorted_by_value[:10]]
        return self.basket_symbols

    def OnData(self, slice):
        if self.IsWarmingUp: return
        self.ApplyCappedRiskManagement()
        
        # Update Risk Logs every hour (optional, but good for monitoring)
        if self.Time.minute == 0:
            self.LogPortfolioRisk()

        if self.Time < self.next_rebalance: return
        
        index_chain = slice.OptionChains.get(self.spy_option.Symbol)
        if not index_chain or not self.basket_symbols: return

        if self.Portfolio.Invested:
            self.Liquidate()
            self.entry_prices.clear()

        # --- Entry ---
        expiry = self.GetTargetExpiry(index_chain)
        if not expiry: return
        
        # Short Index Leg
        contracts = [o for o in index_chain if o.Expiry == expiry]
        atm_index = sorted(contracts, key=lambda x: abs(x.Strike - index_chain.Underlying.Price))[0]
        
        if atm_index.Greeks.Vega > 0:
            index_qty = self.target_vega / (atm_index.Greeks.Vega * 100)
            self.Sell(atm_index.Symbol, index_qty)

        # Buy Basket Legs
        vega_per_stock = self.target_vega / len(self.basket_symbols)
        for symbol in self.basket_symbols:
            if not self.Securities.ContainsKey(symbol):
                opt = self.AddOption(symbol, Resolution.Minute)
                opt.PriceModel = OptionPriceModels.BlackScholes()
                opt.SetFilter(-2, 2, 25, 45)
            
            chain = slice.OptionChains.get(symbol)
            if chain:
                stock_contracts = [o for o in chain if o.Expiry == expiry]
                if not stock_contracts: continue
                atm_stock = sorted(stock_contracts, key=lambda x: abs(x.Strike - chain.Underlying.Price))[0]
                if atm_stock.Greeks.Vega > 0:
                    stock_qty = vega_per_stock / (atm_stock.Greeks.Vega * 100)
                    self.Buy(atm_stock.Symbol, stock_qty)
                    self.entry_prices[atm_stock.Symbol] = atm_stock.LastPrice

        self.next_rebalance = self.Time + timedelta(days=21)

    def LogPortfolioRisk(self):
        """Calculates Net Delta and Net Vega to confirm neutrality."""
        net_delta = 0
        net_vega = 0
        
        for symbol, holdings in self.Portfolio.items():
            if not holdings.Invested: continue
            
            if symbol.SecurityType == SecurityType.Option:
                security = self.Securities[symbol]
                if hasattr(security, "Greeks"):
                    # Option Delta (x100 multiplier)
                    if security.Greeks.Delta:
                        net_delta += security.Greeks.Delta * holdings.Quantity * 100
                    # Option Vega (x100 multiplier)
                    if security.Greeks.Vega:
                        net_vega += security.Greeks.Vega * holdings.Quantity * 100
            
            elif symbol.SecurityType == SecurityType.Equity:
                # Stock Delta (1 share = 1 delta)
                net_delta += holdings.Quantity

        self.Plot("Portfolio Risk", "Net Delta", net_delta)
        self.Plot("Portfolio Risk", "Net Vega", net_vega)

    def ApplyCappedRiskManagement(self):
        for symbol, entry_p in list(self.entry_prices.items()):
            if self.Portfolio[symbol].Invested:
                if self.Securities[symbol].Price >= (entry_p * 2.5):
                    self.Liquidate(symbol)
                    self.Liquidate(symbol.Underlying)
                    del self.entry_prices[symbol]

    def GetTargetExpiry(self, chain):
        expiries = sorted(list(set([x.Expiry for x in chain])))
        for e in expiries:
            if 25 <= (e - self.Time).days <= 45: return e
        return None

    def DeltaHedge(self):
        """Zeroes out Delta daily and logs the risk result."""
        if self.IsWarmingUp: return
        for opt_symbol in list(self.Portfolio.Keys):
            if opt_symbol.SecurityType == SecurityType.Option:
                holdings = self.Portfolio[opt_symbol]
                if holdings.Quantity == 0: continue
                security = self.Securities[opt_symbol]
                if hasattr(security, "Greeks") and security.Greeks.Delta is not None:
                    target_shares = round(-(security.Greeks.Delta * holdings.Quantity * 100))
                    current_shares = self.Portfolio[opt_symbol.Underlying].Quantity
                    diff = target_shares - current_shares
                    if diff != 0:
                        self.MarketOrder(opt_symbol.Underlying, diff)
        
        self.LogPortfolioRisk()