from AlgorithmImports import *
import json
import math
import pandas as pd
from datetime import timedelta
from io import StringIO
import os

class Form4TradingAlgorithm(QCAlgorithm):
    def Initialize(self):
        # self.csv_file_name = "StrategyA-2012-2023.csv"
        # self.csv_file_name = "StrategyA_2024.csv"
        # self.csv_file_name = "StrategyA_2024_2.csv"
        # self.csv_file_name = "StrategyA_2024_V2.csv"
        self.csv_file_name = "StrategyA_2024_V3.csv"
        # self.csv_file_name = "StrategyA_allInsiderTrades-2011-2023.csv"
        start_year = 2024
        end_year = 2024
        starting_capital = 100000
        maxSimultaneousOpenPositions = 10
        leverage = 1

        self.SetStartDate(start_year, 1, 1)  # Start backtesting period
        self.SetEndDate(end_year, 12, 31)  # End backtesting period
        
        self.SetCash(starting_capital)  # Initial capital for each year
        self.maxSimultaneousOpenPositions = maxSimultaneousOpenPositions
        self.maxInitialCapitalPerPosition = starting_capital / self.maxSimultaneousOpenPositions
        self.leverage = leverage # was 3: great results
        
        self.set_brokerage_model(BrokerageName.INTERACTIVE_BROKERS_BROKERAGE, AccountType.Margin)
        self.stopLossThreshold = -0.15  # -15% stop loss

        self.openPositions = []  # Track open positions

        self.allInsiderTrades = self.LoadAllInsiderTrades()

        self.AddEquity("SPY", Resolution.Daily)
        self.SetBenchmark("SPY")

        # Add ticker to universe to receive pricing data
        for _, insiderTrade in self.allInsiderTrades.iterrows():
            ticker = insiderTrade['ticker']
            # QC default leverage: 2
            # this means algo can use twice the value of portfolio (50% margin requirement) when trading equities.
            # self.AddEquity(ticker, Resolution.Daily, leverage=2)
            self.AddEquity(ticker, Resolution.Daily, leverage=self.leverage)


        self.Schedule.On(self.DateRules.EveryDay("SPY"), # Align with the trading calendar of SPY, excluding weekends and holidays
                        # SPY: Symbol of the security whose market open time is being used as a reference
                        # 0: This specifies how many minutes after the market opens the function should be triggered.
                        self.TimeRules.AfterMarketOpen("SPY", 1), 
                        self.OnMarketOpen)

        # realize all unrealized gains/losses at last trading day of the year
        # self.ScheduleLiquidateOnLastTradingDay()

        self.Debug(f"Starting year: {start_year}")


    # def OnData(self, data):
    def OnMarketOpen(self):
        # self.Debug(f"Market open logic executed at: {self.Time}")
        # Check open positions
        self.CheckOpenPositions()

        # Get filings for the previous trading day
        todaysInsiderTrades = self.GetTodaysInsiderTrades()

        if len(todaysInsiderTrades) == 0:
            return

        tickersToSkip = todaysInsiderTrades[todaysInsiderTrades['ticker'].apply(lambda ticker: self.Portfolio[ticker].Invested)]
        if len(tickersToSkip) > 0:
            tickersToSkip = tickersToSkip['ticker'].to_list()
            self.Debug(f"{self.Time} - Skipping tickers: {tickersToSkip}")

        # Filter insider trades to only include tickers that are not already 
        # included as long open position in portfolio
        todaysInsiderTrades = todaysInsiderTrades[~todaysInsiderTrades['ticker'].apply(lambda ticker: self.Portfolio[ticker].Invested)]

        if len(todaysInsiderTrades) == 0:
            return

        # only trade liquid stocks
        # todaysInsiderTrades = todaysInsiderTrades[todaysInsiderTrades['ticker'].apply(lambda ticker: self.IsVolumeAboveThreshold(ticker))]
        # todaysInsiderTrades = todaysInsiderTrades[todaysInsiderTrades['ticker'].apply(lambda ticker: self.IsTradeable(ticker))]
        # todaysInsiderTrades = todaysInsiderTrades[todaysInsiderTrades['ticker'].apply(lambda ticker: self.IsNotTooCheap(ticker))]

        # de-duplicate on ticker, keep first occurence of ticker
        # NOTE: performance drops with de-duplication is applied
        # todaysInsiderTrades = todaysInsiderTrades.drop_duplicates(subset='ticker', keep='first')

        if len(todaysInsiderTrades) == 0:
            return

        # Calculate available slots and capital allocation
        availableSlots = self.maxSimultaneousOpenPositions - len(self.openPositions)

        if availableSlots <= 0:
            return

        availableCash = self.Portfolio.Cash
        availableBuyingPower = availableCash + (self.Portfolio.TotalPortfolioValue - availableCash) * (self.leverage - 1)
        positionSize = availableBuyingPower / availableSlots
        positionSize = min(positionSize, self.maxInitialCapitalPerPosition)

        for _, insiderTrade in todaysInsiderTrades[:availableSlots].iterrows():
            # Fetch current stock price
            ticker = insiderTrade['ticker']
            openPrice = self.Securities[ticker].Open

            if openPrice is None or openPrice <= 0:
                continue

            numShares = math.floor(positionSize / openPrice)
            
            if numShares > 0:
                # self.Debug(f"Buying {ticker} at {self.Time} for ${openPrice} ({numShares} shares)")
                # self.MarketOnOpenOrder(ticker, numShares)
                self.MarketOrder(ticker, numShares)

                self.openPositions.append({
                    "Symbol": ticker,
                    "EntryDate": self.Time,
                    "EntryPrice": openPrice,
                    "NumShares": numShares,
                    "StopLossPrice": openPrice * (1 + self.stopLossThreshold)
                })


    def CheckOpenPositions(self):
        for position in self.openPositions[:]:
            symbol = position["Symbol"]
            stopLossPrice = position["StopLossPrice"]
            entryDate = position["EntryDate"]
            numShares = position["NumShares"]

            currentPrice = self.Securities[symbol].Open

            if currentPrice is None:
                continue

            # calender days    
            # daysHeld = (self.Time - entryDate).days
            # business days
            daysHeld = len(pd.date_range(start=entryDate, end=self.Time, freq='B')) - 1

            if currentPrice <= stopLossPrice or daysHeld >= 5:
                # self.Liquidate(symbol) # Market order at the date/time this function runs
                self.MarketOrder(symbol, -1 * numShares)
                self.openPositions.remove(position)


    def GetTodaysInsiderTrades(self):
        today = self.Time.date()
        return self.allInsiderTrades[self.allInsiderTrades['date'] == pd.Timestamp(today)]


    def LoadAllInsiderTrades(self):
        try:
            # data = self.ObjectStore.ReadBytes(self.objectStoreKey)
            # if data:
                # filings = json.loads(data)
                # return [f for f in filings if f['Date'] == (self.Time - timedelta(days=1)).strftime('%Y-%m-%d')]
            csv_data = self.ObjectStore.Read(self.csv_file_name)
            # df = pd.read_csv(StringIO(csv_data), usecols=['index', 'issuerTicker', 'filedAt'], sep=",")
            df = pd.read_csv(StringIO(csv_data), sep=",")
            df.rename(columns={'index': 'date', 'issuerTicker': 'ticker'}, inplace=True)
            df['date'] = pd.to_datetime(df['date'], format='%Y-%m-%d')
            if "filedAt" in df.columns:
                self.Debug("Sorting insiders by 'filedAt'")
                df.sort_values(by='filedAt', inplace=True)
            else:
                self.Debug("Sorting insiders by 'date'")
                df.sort_values(by='date', inplace=True)

            self.Debug(f"Loaded {len(df):,.0f} insider trades")
            return df
        except Exception as e:
            self.Debug(f"Error reading CSV: {e}")
            return


    def ScheduleLiquidateOnLastTradingDay(self):
        # Get the last calendar day of the year
        lastCalendarDay = datetime(self.Time.year, 12, 31)

        # Find the last trading day
        lastTradingDay = self.Securities["SPY"].Exchange.Hours.GetPreviousTradingDay(lastCalendarDay)

        # Schedule the liquidation at market open on the last trading day
        self.Schedule.On(
            self.DateRules.On(lastTradingDay.year, lastTradingDay.month, lastTradingDay.day),
            # self.TimeRules.AfterMarketOpen("SPY", 2),
            self.TimeRules.BeforeMarketClose("SPY", 10), # 10 min before market close
            self.LiquidateAllPositions
        )

    def LiquidateAllPositions(self):
        self.Liquidate()  # Liquidates all positions
        self.Debug(f"All positions liquidated on {self.Time}.")


    def IsTradeable(self, ticker):
        return self.IsVolumeAboveThreshold(ticker) and self.IsNotTooCheap(ticker)

    
    def IsNotTooCheap(self, ticker, limit_price=5):
        return self.Securities[ticker].Open >= limit_price

    
    def IsVolumeAboveThreshold(self, ticker, volumeThreshold=100000):
        try:
            symbol = Symbol.Create(ticker, SecurityType.Equity, Market.USA)
            history = self.History(symbol, 20, Resolution.Daily)  # Fetch 20 days of historical data
            if history.empty:
                return False
            # Convert history to a DataFrame
            history_df = history.loc[symbol]
            avgVolume = history_df['volume'].mean()
            return avgVolume >= volumeThreshold
        except Exception as e:
            self.Debug(f"Error checking volume for {ticker}: {e}")
            return False