from AlgorithmImports import QCAlgorithm, Resolution
import pandas as pd
from datetime import timedelta
from io import StringIO
from pandas.tseries.offsets import BDay
import math

class RuleBasedEquityLongStrategy(QCAlgorithm):
    def Initialize(self):
        # Set the backtest timeframe and starting cash
        self.SetStartDate(2023, 1, 1)
        self.SetEndDate(2023, 12, 31)
        self.SetCash(100000)
        self.settings.daily_precise_end_time = True  # Ensures accurate daily timing

        # Add benchmark
        self.AddEquity("SPY", Resolution.Daily)
        self.SetBenchmark("SPY")

        # Positions tracking
        self.positions = {}
        self.processed_tickers = {}  # Tracks processed tickers per date

        # Load signal data from Object Store
        if self.ObjectStore.ContainsKey("StrategyA-2012-2023.csv"):
            try:
                csv_file = self.ObjectStore.Read("StrategyA-2012-2023.csv")
                self.signal_data = self.ParseCsv(csv_file)
            except Exception as e:
                self.Debug(f"Error reading CSV: {e}")
                self.signal_data = pd.DataFrame()
        else:
            self.Debug("2012-2023.csv not found in the Object Store.")
            self.signal_data = pd.DataFrame()

        # Validate signal data
        if self.signal_data.empty:
            self.Debug("No valid signal data. Algorithm will run without processing signals.")
        else:
            self.Debug(f"Loaded {len(self.signal_data)} signals.")

        self.capital_per_position = 10000

        # Add all tickers in advance
        for _, signal in self.signal_data.iterrows():
            ticker = signal['ticker']
            self.AddEquity(ticker, Resolution.Daily)

    def OnData(self, data):
        today = self.Time.date()

        # Initialize processed tickers for today if not already done
        if today not in self.processed_tickers:
            self.processed_tickers[today] = set()

        positions_to_open = 0
        # Process signals
        # Compute the number of signals for the current date
        daily_signals = self.signal_data[self.signal_data['date'] == pd.Timestamp(today)]
        positions_to_open = len(daily_signals)
        self.Debug(f"Number of signals for {today}: {positions_to_open}")

        if positions_to_open > 0:
            available_cash = self.Portfolio.Cash
            capital_per_position_today = available_cash / positions_to_open
            position_size = min(self.capital_per_position, capital_per_position_today)
            
            for _, signal in self.signal_data.iterrows():
                ticker = signal['ticker']
                signal_date = signal['date']

                if today == signal_date.date() and ticker not in self.processed_tickers[today]:
                    ticker_open = self.Securities[ticker].Open

                    if ticker_open == 0 or ticker_open is None:
                        self.Debug(f"Error on {ticker} Open.")

                    if position_size > 0 and ticker_open > 0:
                        qty = math.floor(position_size/ticker_open)
                        self.MarketOnOpenOrder(ticker, qty)
                        self.positions[ticker] = self.Time
                        self.processed_tickers[today].add(ticker)  # Mark ticker as processed

            # Manage exits
            to_remove = []
            for ticker, entry_date in self.positions.items():
                # Exit if holding period exceeds 5 business days
                if self.IsHoldingPeriodExceeded(entry_date, 5):
                    self.Liquidate(ticker)
                    to_remove.append(ticker)

                # Stop-loss logic
                current_price = self.Securities[ticker].Price
                entry_price = self.Portfolio[ticker].AveragePrice
                if current_price <= entry_price * 0.85:
                    self.Liquidate(ticker)
                    to_remove.append(ticker)

            # Remove exited positions
            for ticker in to_remove:
                try:
                    del self.positions[ticker]
                except KeyError:
                    continue

    def IsHoldingPeriodExceeded(self, entry_date, num_business_days):
        # Calculate the expected exit date
        exit_date = pd.date_range(entry_date, periods=num_business_days + 1, freq=BDay())[num_business_days]
        return self.Time.date() >= exit_date.date()

    def ParseCsv(self, csv_data):
        try:
            # Parse CSV
            df = pd.read_csv(StringIO(csv_data), usecols=['index', 'issuerTicker'], sep=",")
            df.rename(columns={'index': 'date', 'issuerTicker': 'ticker'}, inplace=True)
            df['date'] = pd.to_datetime(df['date'], format='%Y-%m-%d')
            df.sort_values(by='date', inplace=True)
            return df
        except Exception as e:
            self.Debug(f"CSV Parsing Error: {e}")
            return pd.DataFrame()