from AlgorithmImports import *
import json
from collections import defaultdict

class GammaExposureAlgorithm(QCAlgorithm):
    
    def Initialize(self):
        self.SetStartDate(2025, 8, 5)
        self.SetEndDate(2025, 8, 5)
        self.SetCash(100000)
        
        self.spx = self.AddIndex("SPX", Resolution.Minute).Symbol
        option = self.AddIndexOption(self.spx, "SPXW", Resolution.Minute)
        option.SetFilter(lambda x: x.IncludeWeeklys().Strikes(-100, 100).Expiration(0, 7))
        
        self.Schedule.On(self.DateRules.EveryDay(), self.TimeRules.At(10, 30), self.CalculateGEX)
        self.object_store_key = f"gex_test_{self.Time.strftime('%Y%m%d')}"
        
    def OnData(self, data):
        pass
    
    def CalculateGEX(self):
        spot = self.Securities[self.spx].Price
        if spot == 0:
            self.Log("ERROR: No SPX price")
            return
            
        chains = self.CurrentSlice.OptionChains
        if not chains:
            self.Log("ERROR: No option chains")
            return
        
        gex_data = defaultdict(lambda: {'call_gex': 0, 'put_gex': 0, 'total_gex': 0})
        total_gex = contracts = 0
        today = self.Time.date()
        
        for chain in chains.values():
            for c in chain:
                if c.Expiry.date() != today or not c.Greeks or c.Greeks.Gamma == 0:
                    continue
                
                strike = float(c.Strike)
                gamma = c.Greeks.Gamma
                oi = getattr(c, 'OpenInterest', None) or c.Volume or 1000
                
                contract_gex = gamma * oi * 100 * (spot ** 2)
                adjusted_gex = -contract_gex if c.Right == OptionRight.Call else contract_gex
                
                if c.Right == OptionRight.Call:
                    gex_data[strike]['call_gex'] += adjusted_gex
                else:
                    gex_data[strike]['put_gex'] += adjusted_gex
                    
                gex_data[strike]['total_gex'] += adjusted_gex
                total_gex += adjusted_gex
                contracts += 1
        
        self.Log(f"SPX: {spot}, Contracts: {contracts}, Total GEX: ${total_gex:,.0f}")
        self.ExportData(gex_data, spot, total_gex)
        
    def ExportData(self, gex_data, spot, total_gex):
        strikes = sorted(gex_data.keys())
        
        export_data = {
            'metadata': {
                'timestamp': self.Time.isoformat(),
                'spot_price': float(spot),
                'total_gamma_exposure': float(total_gex),
                'num_strikes': len(strikes)
            },
            'strike_data': [{
                'strike': float(s),
                'total_gex': float(gex_data[s]['total_gex']),
                'call_gex': float(gex_data[s]['call_gex']),
                'put_gex': float(gex_data[s]['put_gex']),
                'distance_from_spot': float(s - spot)
            } for s in strikes]
        }
        
        try:
            success = self.ObjectStore.Save(self.object_store_key, json.dumps(export_data))
            self.Log(f"{'✓' if success else '✗'} Export: {len(strikes)} strikes")
        except Exception as e:
            self.Log(f"Export error: {e}")
            
        # Log top 3 GEX strikes
        top_strikes = sorted(gex_data.items(), key=lambda x: abs(x[1]['total_gex']), reverse=True)[:3]
        for i, (strike, data) in enumerate(top_strikes):
            self.Log(f"{i+1}. {strike}: ${data['total_gex']:,.0f}")
        
    def OnEndOfAlgorithm(self):
        self.Log(f"Complete. ObjectStore key: {self.object_store_key}")