from AlgorithmImports import *

import random
from datetime import timedelta, datetime


class HookStrategyAlgorithm(QCAlgorithm):
    
    def Initialize(self):
        self.SetStartDate(2020, 2, 21)  # 4 years ago from current date
        self.SetEndDate(2025, 2, 21)    # Current date 
        self.SetCash(100000)            # Starting capital

        
        # Add your desired asset (example with SPY)
        self.symbol = self.AddEquity("SPY", Resolution.Minute).Symbol
        
        # Set trading parameters
        self.lookback_period = 4 * 252 * 24  # 4 years of 15-minute bars (approximate)
        self.risk_percentage = 0.02          # 2% risk per trade
        self.reward_risk_ratio = 2           # 1:2 risk-reward ratio
        
        # Initialize variables for position tracking
        self.is_in_position = False
        self.position_type = None
        self.entry_price = 0
        self.stop_loss = 0
        self.target_price = 0
        
        # Function parameters
        self.lower_band = 0.85
        self.upper_band = 0.98
        self.pulse_length = 100
        self.pulse_a = 0.0001
        self.pulse_b = 0.01
        self.pulse_c = 0.5
        self.pulse_breakout_threshold = 2
        
        # Consolidator for 15-minute bars
        self.consolidator = self.consolidate(self.symbol, timedelta(minutes=15), self.OnDataConsolidated)

        
        # Initialize dataframe for historical data
        self.history_df = None
        self.scheduled_history_pull = True
        
        # Schedule the initial history pull
        self.schedule.on(self.date_rules.every_day(self.symbol), 
                         self.time_rules.after_market_open(self.symbol, 5), 
                         self.PullHistoricalData)
    
    def PullHistoricalData(self):
        if self.scheduled_history_pull:
            # Pull 4 years of 15-minute bar history
            history = self.history(self.symbol, self.lookback_period, Resolution.Minute)
            self.Debug(f"Fetched Historical data. Got {len(history)} bars.")
            history = history.reset_index(level='symbol', drop=True)
            history.index = pd.to_datetime(history.index)
            if history.empty:
                self.Debug("No historical data available yet. Will try again.")
                return
            
            # Convert to 15-minute bars
            history_bars = history.resample('15T').agg({
                'open': 'first',
                'high': 'max',
                'low': 'min',
                'close': 'last',
                'volume': 'sum'
            }).dropna()
            
            # Create DataFrame with required format
            df = history_bars[['open', 'high', 'low', 'close', 'volume']].copy()
            df['time'] = history_bars.index
            df['index'] = np.arange(len(df))
            
            # Rename columns as specified
            df.rename(columns={
                'time': 'Time', 
                'open': 'Open', 
                'high': 'High', 
                'low': 'Low', 
                'close': 'Close', 
                'volume': 'Volume',
                'index': 'index'
            }, inplace=True)
            
            self.history_df = df

            self.scheduled_history_pull = False
            self.Debug(f"Historical data processed. Got {len(self.history_df)} bars.")
    
    def OnDataConsolidated(self, consolidated_bar: TradeBar) -> None:
        """Handler for consolidated 15-minute bars"""
        if self.history_df is None:
            return
        
        # Append the new bar to our historical dataframe
        new_bar = pd.DataFrame([{
            'Time': consolidated_bar.Time,
            'Open': consolidated_bar.Open,
            'High': consolidated_bar.High,
            'Low': consolidated_bar.Low,
            'Close': consolidated_bar.Close,
            'Volume': consolidated_bar.Volume
        }])
        
        self.history_df = pd.concat([self.history_df, new_bar], ignore_index=True)
        #self.Debug(f"New bar added: {new_bar.iloc[0].to_dict()}")

        # Process the strategy with the new bar
        self.ProcessStrategy(consolidated_bar.Close, len(self.history_df) - 1)
        
        # Check if we need to exit existing positions
        self.CheckPositionExit(consolidated_bar.Close)
    
    def ProcessStrategy(self, current_price, current_index):
        """Process strategy logic for each new 15-minute bar"""
        if current_index < 5:  # Need some bars to calculate indicators
            return
            
        # Check if we're already in a position
        if self.is_in_position:
            return
            
        # Get pulse guardian signal
        pulse_result = self.Pulse_Guardian_with_Quadratic_Volatility(
            self.history_df, 
            length=self.pulse_length,
            a=self.pulse_a,
            b=self.pulse_b,
            c=self.pulse_c,
            breakout_threshold=self.pulse_breakout_threshold
        )

        #self.Debug(f"current breakout= {pulse_result}")

        if not pulse_result:
            return  # No breakout signal
            
        # Check for long signal
        long_signal, long_stop_loss, *long_rest = self.Is_Pos_Hook2(
            self.history_df,
            current_price,
            current_index,
            self.lower_band,
            self.upper_band
        )

        # Check for short signal
        short_signal, short_stop_loss, *short_rest = self.Is_Neg_Hook2(
            self.history_df,
            current_price,
            current_index,
            self.lower_band,
            self.upper_band
        )

        # Debug the results, printing long_rest and short_rest separately
        self.Debug(f"Long signal result: {long_signal}, Stop Loss: {long_stop_loss}, Rest: {long_rest}")
        self.Debug(f"Short signal result: {short_signal}, Stop Loss: {short_stop_loss}, Rest: {short_rest}")


        # Process long signals
        if long_signal in ['Type A hook. Buy!', 'Type B hook. Buy!']:
            # Calculate position size based on risk
            stop_distance = abs(current_price - long_stop_loss)
            if stop_distance == 0:
                self.Debug("Warning: Stop distance is zero, skipping trade")
                return
                
            risk_amount = self.portfolio.Cash * self.risk_percentage
            shares_to_buy = int(risk_amount / stop_distance)
            
            if shares_to_buy == 0:
                self.Debug("Warning: Not enough capital to take position with current risk parameters")
                return
                
            # Calculate target based on risk:reward ratio
            target_price = current_price + (stop_distance * self.reward_risk_ratio)
            
            # Enter long position
            self.market_order(self.symbol, shares_to_buy)
            self.is_in_position = True
            self.position_type = "LONG"
            self.entry_price = current_price
            self.stop_loss = long_stop_loss
            self.target_price = target_price
            
            # Log trade information
            self.Debug(f"LONG ENTRY: Price=${current_price:.2f}, Stop=${long_stop_loss:.2f}, Target=${target_price:.2f}, Shares={shares_to_buy}")
            
        # Process short signals
        elif short_signal in ['Type A hook. Sell!', 'Type B hook. Sell!']:
            # Calculate position size based on risk
            stop_distance = abs(current_price - short_stop_loss)
            if stop_distance == 0:
                self.Debug("Warning: Stop distance is zero, skipping trade")
                return
                
            risk_amount = self.portfolio.Cash * self.risk_percentage
            shares_to_short = int(risk_amount / stop_distance)
            
            if shares_to_short == 0:
                self.Debug("Warning: Not enough capital to take position with current risk parameters")
                return
                
            # Calculate target based on risk:reward ratio
            target_price = current_price - (stop_distance * self.reward_risk_ratio)
            
            # Enter short position
            self.market_order(self.symbol, -shares_to_short)
            self.is_in_position = True
            self.position_type = "SHORT"
            self.entry_price = current_price
            self.stop_loss = short_stop_loss
            self.target_price = target_price
            
            # Log trade information
            self.Debug(f"SHORT ENTRY: Price=${current_price:.2f}, Stop=${short_stop_loss:.2f}, Target=${target_price:.2f}, Shares={shares_to_short}")
    
    def CheckPositionExit(self, current_price):
        """Check if we need to exit current position based on stop loss or target price"""
        if not self.is_in_position:
            return
            
        if self.position_type == "LONG":
            # Check if stop loss or target hit
            if current_price <= self.stop_loss:
                self.market_order(self.symbol, -self.portfolio[self.symbol].Quantity)
                self.Debug(f"LONG EXIT - STOP LOSS: Entry=${self.entry_price:.2f}, Exit=${current_price:.2f}")
                self.ResetPosition()
            elif current_price >= self.target_price:
                self.market_order(self.symbol, -self.portfolio[self.symbol].Quantity)
                self.Debug(f"LONG EXIT - TARGET: Entry=${self.entry_price:.2f}, Exit=${current_price:.2f}")
                self.ResetPosition()
                
        elif self.position_type == "SHORT":
            # Check if stop loss or target hit
            if current_price >= self.stop_loss:
                self.market_order(self.symbol, -self.portfolio[self.symbol].Quantity)
                self.Debug(f"SHORT EXIT - STOP LOSS: Entry=${self.entry_price:.2f}, Exit=${current_price:.2f}")
                self.ResetPosition()
            elif current_price <= self.target_price:
                self.market_order(self.symbol, -self.portfolio[self.symbol].Quantity)
                self.Debug(f"SHORT EXIT - TARGET: Entry=${self.entry_price:.2f}, Exit=${current_price:.2f}")
                self.ResetPosition()
    
    def ResetPosition(self):
        """Reset position tracking variables"""
        self.is_in_position = False
        self.position_type = None
        self.entry_price = 0
        self.stop_loss = 0
        self.target_price = 0
    

    def Is_Pos_Hook2(self, df, Current_Closed_Price, Index_Current, Lower_band, Upper_band):
        """
        Implementation of the Is_Pos_Hook2 function for long signals
        Note: Placeholder implementation as the actual logic wasn't shared
        """
        # Since the actual logic wasn't shared, we'll return placeholder values
        # In your actual implementation, replace this with your proprietary logic
        
        # Check the last few candles for a pattern (placeholder logic)
        if Index_Current < 5:
            return "No signal", Current_Closed_Price, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, Index_Current, Current_Closed_Price, Current_Closed_Price
        
        # Placeholder for your proprietary logic
        # Here we just use a simple example - replace with your actual logic
        last_candles = df.iloc[Index_Current-5:Index_Current+1]
        
        # Example placeholder logic (replace with your actual proprietary logic)
        if (last_candles['Close'].iloc[-1] > last_candles['Close'].iloc[-2] and
            last_candles['Close'].iloc[-2] > last_candles['Close'].iloc[-3] and
            last_candles['Low'].iloc[-1] > last_candles['Low'].iloc[-3]):
            signal = "Type A hook. Buy!"
            stop_loss = last_candles['Low'].min() * 0.99  # Placeholder stop loss
        elif (last_candles['Close'].iloc[-1] > last_candles['Open'].iloc[-1] and
              last_candles['Close'].iloc[-2] < last_candles['Open'].iloc[-2]):
            signal = "Type B hook. Buy!"
            stop_loss = last_candles['Low'].min() * 0.99  # Placeholder stop loss
        else:
            signal = "No signal"
            stop_loss = Current_Closed_Price * 0.95  # Placeholder
        
        # Return placeholder values for all required outputs
        return (signal, stop_loss, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 
                Index_Current, Current_Closed_Price, Current_Closed_Price)
    
    def Is_Neg_Hook2(self, df, Current_Closed_Price, Index_Current, Lower_band, Upper_band):
        """
        Implementation of the Is_Neg_Hook2 function for short signals
        Note: Placeholder implementation as the actual logic wasn't shared
        """
        # Since the actual logic wasn't shared, we'll return placeholder values
        # In your actual implementation, replace this with your proprietary logic
        
        # Check the last few candles for a pattern (placeholder logic)
        if Index_Current < 5:
            return "No signal", Current_Closed_Price, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, Index_Current, Current_Closed_Price, Current_Closed_Price
        
        # Placeholder for your proprietary logic
        # Here we just use a simple example - replace with your actual logic
        last_candles = df.iloc[Index_Current-5:Index_Current+1]
        
        # Example placeholder logic (replace with your actual proprietary logic)
        if (last_candles['Close'].iloc[-1] < last_candles['Close'].iloc[-2] and
            last_candles['Close'].iloc[-2] < last_candles['Close'].iloc[-3] and
            last_candles['High'].iloc[-1] < last_candles['High'].iloc[-3]):
            signal = "Type A hook. Sell!"
            stop_loss = last_candles['High'].max() * 1.01  # Placeholder stop loss
        elif (last_candles['Close'].iloc[-1] < last_candles['Open'].iloc[-1] and
              last_candles['Close'].iloc[-2] > last_candles['Open'].iloc[-2]):
            signal = "Type B hook. Sell!"
            stop_loss = last_candles['High'].max() * 1.01  # Placeholder stop loss
        else:
            signal = "No signal"
            stop_loss = Current_Closed_Price * 1.05  # Placeholder
        
        # Return placeholder values for all required outputs
        return (signal, stop_loss, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,
                Index_Current, Current_Closed_Price, Current_Closed_Price)
    
    def Pulse_Guardian_with_Quadratic_Volatility(self, df, length=100, a=0.0001, b=0.01, c=0.5, breakout_threshold=2):
        """
        Implementation of the Pulse_Guardian_with_Quadratic_Volatility function
        Note: Placeholder implementation as the actual logic wasn't shared
        """
        # Since the actual logic wasn't shared, we'll create a placeholder implementation
        # In your actual implementation, replace this with your proprietary logic
        
        return random.choice([True, False])