from AlgorithmImports import *
from datetime import datetime, timedelta


def get_third_friday(year, month):
    """
    Get the date of the third Friday of the given month and year.
    """
    # Start with the first day of the month
    first_day = datetime(year, month, 1)
    first_friday = first_day + timedelta(days=(4 - first_day.weekday() + 7) % 7)
    
    # Add two weeks to get to the third Friday
    third_friday = first_friday + timedelta(weeks=2)
    
    return third_friday.date()

def is_monthly_expiration_week(current_date):
    """
    Check if current date is in the monthly options expiration week.
    Returns:
        tuple: (bool, int, str) - (is_expiration_week, days_until_expiration, day_description)
              days_until_expiration will be:
              2 for Wednesday before expiration
              1 for Thursday before expiration
              0 for expiration Friday
              -1 if not in expiration week
              day_description will describe the current day (e.g., 'Wednesday before expiration').
    """

    third_friday = get_third_friday(current_date.year, current_date.month)
    
    # Find the Wednesday of expiration week
    expiration_wednesday = third_friday - timedelta(days=2)  # Wednesday
    expiration_thursday = third_friday - timedelta(days=1)    # Thursday
    
    # Calculate days until expiration
    days_until = (third_friday - current_date).days
    
    # Check if current date is within Wednesday to Friday of expiration week
    if expiration_wednesday <= current_date <= third_friday:
        if days_until == 2:
            return 2
        elif days_until == 1:
            return 1
    
    return False

"""
Options Flow Analysis Algorithm
Aligns with Monthly IV Algorithm structure for consistency and reduced errors
"""
from AlgorithmImports import *
import numpy as np
import pandas as pd
from collections import defaultdict
from date_utils import is_monthly_expiration_week

class OptionsFlowAlgorithm(QCAlgorithm):
    def Initialize(self):
        # Set start and end dates
        self.SetStartDate(2022, 12, 1)
        self.SetEndDate(2025, 3, 1)
        
        # Initial cash
        self.InitCash = 100000
        self.SetCash(self.InitCash)
        
        # Track first run
        self.first_run = True

        # Strategy Parameters
        self.params = {
            # Portfolio Construction
            'low_flow_allocation': 0.0,    # Portfolio allocation to low options flow assets
            'mid_flow_allocation': 0.0,    # Portfolio allocation to mid options flow assets
            'high_flow_allocation': 1.0,   # Portfolio allocation to high options flow assets
            
            # Flow Categorization
            'low_flow_percentile': 25,     # Threshold for low flow category (0-25th percentile)
            'high_flow_percentile': 75,    # Threshold for high flow category (75-100th percentile)
            
            # Risk Management
            'min_sharpe_ratio': 1.0,       # Minimum Sharpe ratio for inclusion
            'max_position_size': 0.80,     # Maximum allocation % to any single position
            'min_position_size': 0.01,     # Minimum allocation % to any single position
            
            # Volatility Parameters
            'lookback_period': 180,        # Days for calculating Sharpe ratio
            'risk_free_rate': 0.04,        # Annual risk-free rate for Sharpe calculation
            
            # Rebalancing Parameters
            'days_before_expiry': 1,       # Days before expiry to rebalance
            
            # Technical Indicators
            'use_trend_filter': False,     # Use trend filter for entries
            'trend_lookback': 20,          # Days for trend calculation
            'momentum_lookback': 20,       # Days for momentum calculation
            'rsi_oversold': 30,            # RSI oversold threshold
            'rsi_overbought': 70,          # RSI overbought threshold
        }

        # Market benchmark
        self.MKT = self.AddEquity("SPY", Resolution.DAILY).Symbol
        self.spy = []

        # Tickers to monitor

        self.tickers = [
            "TSLA", "NVDA", "AMD", "PLTR", "MSTR", "AAPL", "AMZN", "MO", 
            "MARA", "SMCI", "META", "HOOD", "INTC", "IBIT", "SOFI", "MSFT", 
            "KSS", "MU", "GOOGL", "COIN", "DIS", "ZM", "SNAP", "XLF", "XOM",
            "PYPL", "WMT", "NFLX", "XLY", "QQQ", "GS", "BABA", "WMT", "BA",
            "XOM", "SPY", "SOXS", "GLD"
        ]
        # Add symbols
        self.symbols = [self.AddEquity(ticker, Resolution.DAILY).Symbol for ticker in self.tickers]

        # Set universe settings
        self.UniverseSettings.Asynchronous = True

        # Add options for each symbol
        for symbol in self.symbols:
            option = self.AddOption(symbol, Resolution.DAILY)
            option.SetFilter(self._contract_selector)

        # Schedule daily recording
        self.Schedule.On(
            self.DateRules.EveryDay(), 
            self.TimeRules.BeforeMarketClose('SPY', 0), 
            self.record_vars
        )

    def _contract_selector(self, option_filter_universe: OptionFilterUniverse) -> OptionFilterUniverse:
        """Option contract selection criteria"""
        return (
            option_filter_universe
            .IncludeWeeklys()
            .Expiration(10, 30)
            .IV(0, 100)
        )

    def get_options_flow(self, underlying_symbol):
        """
        Calculate options flow metrics for a given symbol
        
        Returns:
        - Put/call volume ratio
        - Large trade count
        - Unusual activity indicator
        """
        try:
            # Fetch the option chain
            option_chain = self.OptionChain(underlying_symbol, flatten=True).DataFrame

            if option_chain is None or option_chain.empty:
                self.Debug(f"Option chain empty for {underlying_symbol}")
                return None

            # Calculate flow metrics
            call_volume = option_chain[option_chain['right'] == OptionRight.Call]['volume'].sum()
            put_volume = option_chain[option_chain['right'] == OptionRight.Put]['volume'].sum()
            
            # Calculate put/call ratio
            total_volume = call_volume + put_volume
            put_call_ratio = put_volume / total_volume if total_volume > 0 else 1.0

            # Large trade detection (trades above 100 contracts)
            large_call_trades = option_chain[
                (option_chain['right'] == OptionRight.Call) & 
                (option_chain['volume'] > 100)
            ]
            large_put_trades = option_chain[
                (option_chain['right'] == OptionRight.Put) & 
                (option_chain['volume'] > 100)
            ]
            
            # Unusual activity (volume above 2x average)
            historical_volumes = list(self.History(underlying_symbol, 5, Resolution.Daily)['volume'])
            avg_volume = np.mean(historical_volumes) if historical_volumes else 0
            unusual_activity = total_volume > 2 * avg_volume

            return {
                'put_call_ratio': put_call_ratio,
                'large_call_trades': len(large_call_trades),
                'large_put_trades': len(large_put_trades),
                'unusual_activity': unusual_activity
            }
        except Exception as e:
            self.Debug(f"Error calculating options flow for {underlying_symbol}: {e}")
            return None

    def get_sharpe_ratio(self, symbol, lookback_period=None):
        """
        Calculate the Sharpe Ratio for a given symbol
        
        Args:
            symbol: The stock symbol to calculate Sharpe Ratio for
            lookback_period: Number of trading days to look back
        
        Returns:
            float: The Sharpe Ratio value
        """
        if lookback_period is None:
            lookback_period = self.params.get('lookback_period', 252)

        # Get historical daily price data
        history = self.History(symbol, lookback_period, Resolution.Daily)
        if history.empty:
            self.Debug(f"No historical data found for {symbol}")
            return None
        
        # Calculate daily returns
        returns = history['close'].pct_change().dropna()
        
        # Calculate excess returns over risk-free rate
        risk_free_rate = self.params['risk_free_rate']
        excess_returns = returns - (risk_free_rate / self.params['lookback_period'])  # Daily risk-free rate
        
        # Calculate Sharpe Ratio
        sharpe_ratio = np.sqrt(self.params['lookback_period']) * (excess_returns.mean() / excess_returns.std())
        
        return sharpe_ratio

    def categorize_by_flow(self, flow_values):
        """
        Categorize tickers based on options flow percentiles
        """
        # Extract numeric flow values (put/call ratio)
        put_call_ratios = np.array([flow['put_call_ratio'] for flow in flow_values.values()])

        # Calculate percentiles
        p_low = np.percentile(put_call_ratios, self.params['low_flow_percentile'])
        p_high = np.percentile(put_call_ratios, self.params['high_flow_percentile'])

        # Initialize dictionaries for categorization
        low_flow, mid_flow, high_flow = {}, {}, {}

        # Categorize tickers based on put/call ratio
        for ticker, flow in flow_values.items():
            if flow['put_call_ratio'] <= p_low:
                low_flow[ticker] = flow
            elif flow['put_call_ratio'] <= p_high:
                mid_flow[ticker] = flow
            else:
                high_flow[ticker] = flow

        self.Debug(f"Low Flow (0-25th): {len(low_flow)} tickers")
        self.Debug(f"Mid Flow (26-75th): {len(mid_flow)} tickers")
        self.Debug(f"High Flow (76-100th): {len(high_flow)} tickers")

        return low_flow, mid_flow, high_flow


    def get_technical_signals(self, symbol):
        """
        Technical signal calculation similar to the original algorithm
        Reuses the same implementation from the Monthly IV Algorithm
        """
        if not self.params['use_trend_filter']:
            return True

        # Implementation of technical trend indicator here
        # Placeholder return - would need full implementation
        return True

    def OnData(self, data):
        """
        Main trading logic for options flow strategy
        """
        # Similar to Monthly IV Algorithm, check if it's day before expiration
        
        #days_until_exp = is_monthly_expiration_week(self.Time.date())

        # allocate on first day of deployment
        if self.first_run:
            self.first_run = False
         
        # For monthly exp
        # elif days_until_exp != 1:
        #     return

        # Check if today is Monday (0 = Monday, 6 = Sunday)
        elif self.Time.weekday() != 0:
            return

        # Calculate options flow and Sharpe ratios
        options_flow = {}
        sharpe_ratios = {}
        valid_tickers = []

        for symbol in self.symbols:
            # Calculate options flow
            flow = self.get_options_flow(symbol)
            
            # Calculate Sharpe ratio
            sharpe = self.get_sharpe_ratio(symbol)
            
            # Apply filtering criteria
            if (sharpe is not None and 
                sharpe > self.params['min_sharpe_ratio'] and 
                flow is not None):
                
                # Check technical signals if option is set to True
                if self.get_technical_signals(symbol.Value):
                    sharpe_ratios[symbol.Value] = sharpe
                    options_flow[symbol.Value] = flow
                    valid_tickers.append(symbol.Value)

        # Portfolio allocation if we have flow data
        if options_flow:
            # Categorize tickers by flow
            low_flow, mid_flow, high_flow = self.categorize_by_flow(options_flow)
            
            # Portfolio allocations for each category
            category_allocations = {
                'low_flow': self.params['low_flow_allocation'],
                'mid_flow': self.params['mid_flow_allocation'],
                'high_flow': self.params['high_flow_allocation']
            }

            # Iterate through categories and allocate
            for category, tickers in [('low_flow', low_flow), ('mid_flow', mid_flow), ('high_flow', high_flow)]:
                category_weight = category_allocations[category]
                
                # Allocation logic similar to Monthly IV Algorithm
                for ticker, flow_data in tickers.items():
                    symbol = self.Securities[ticker].Symbol
                    allocation = category_weight * (sharpe_ratios[ticker] / sum(sharpe_ratios.values()))
                    
                    # Apply position size limits
                    allocation = min(allocation, self.params['max_position_size'])
                    if allocation >= self.params['min_position_size']:
                        self.SetHoldings(symbol, allocation)
                        
                        # Logging
                        self.Log(f"{ticker}: Category={category}, " +
                                 f"Put/Call Ratio={flow_data['put_call_ratio']:.4f}, " +
                                 f"Sharpe={sharpe_ratios.get(ticker, 'N/A')}, " +
                                 f"Weight={allocation:.4f}")

            # Liquidate positions no longer meeting criteria
            for symbol in self.symbols:
                if (symbol.Value not in valid_tickers and 
                    self.Portfolio[symbol].Invested):
                    self.Liquidate(symbol)
                    self.Log(f"Liquidating {symbol.Value} due to criteria")

    def record_vars(self):             
        """
        Record performance variables, similar to original algorithm
        """
        hist = self.History(self.MKT, 2, Resolution.Daily)['close'].unstack(level=0).dropna() 
        self.spy.append(hist[self.MKT].iloc[-1])
        spy_perf = self.spy[-1] / self.spy[0] * self.InitCash
        self.Plot('Strategy Equity', 'SPY', spy_perf)