from AlgorithmImports import *
import csv
from io import StringIO
from datetime import time, timedelta, date, datetime

class FracSharesAlgorithm(QCAlgorithm):
    def Initialize(self):
        # User-defined variables for time calculations
        self.AdjustedCloseSubtractor = time(0, 15, 0)  # 00:15:00
        self.MarketCloseTime = time(16, 0, 0)  # 16:00:00

        # Calculate the adjusted close time
        close_minutes = self.MarketCloseTime.hour * 60 + self.MarketCloseTime.minute
        subtractor_minutes = self.AdjustedCloseSubtractor.hour * 60 + self.AdjustedCloseSubtractor.minute
        adjusted_minutes = close_minutes - subtractor_minutes
        self.AdjustedCloseTime = time(adjusted_minutes // 60, adjusted_minutes % 60)

        self.Log(f"Market Close Time: {self.MarketCloseTime}")
        self.Log(f"Adjusted Close Subtractor: {self.AdjustedCloseSubtractor}")
        self.Log(f"Adjusted Close Time: {self.AdjustedCloseTime}")

        # Get today's actual date
        today = datetime.now().date()
        self.Log(f"Today's Date: {today}")
        self.Log(f"Today - 1 day: {today - timedelta(days=1)}")
        self.Log(f"Today - 2 days: {today - timedelta(days=2)}")
        self.Log(f"Today - 3 days: {today - timedelta(days=3)}")
        self.Log(f"Today - 4 days: {today - timedelta(days=4)}")
        self.Log(f"Today - 5 days: {today - timedelta(days=5)}")
        self.Log(f"Today - 6 days: {today - timedelta(days=6)}")
        self.Log(f"Today - 7 days: {today - timedelta(days=7)}")
        self.Log(f"Today - 8 days: {today - timedelta(days=8)}")
        self.Log(f"Today - 9 days: {today - timedelta(days=9)}")
        self.Log(f"Today - 10 days: {today - timedelta(days=10)}")
        self.Log(f"Today - 11 days: {today - timedelta(days=11)}")

        # Get FetchEnd and FetchStart dates
        self.FetchEnd = self.GetFetchEnd(today)
        self.FetchStart = self.GetFetchStart(self.FetchEnd)

        # Set the date range
        self.SetStartDate(self.FetchStart)
        self.SetEndDate(today)

        self.Log(f"Start Date: {self.StartDate}")
        self.Log(f"End Date: {self.EndDate}")
        self.Log(f"FetchStart: {self.FetchStart}")
        self.Log(f"FetchEnd: {self.FetchEnd}")

        # Initialize the set to store Fracshares tickers
        self.frac_symbols = set()

        # Schedule the refresh of Fracshares data
        self.Schedule.On(self.DateRules.EveryDay(), self.TimeRules.At(self.AdjustedCloseTime), self.RefreshFracShares)

    def GetFetchEnd(self, start_date):
        if self.IsMarketOpen(start_date):
            return start_date
        elif self.IsMarketOpen(start_date - timedelta(days=1)):
            return start_date - timedelta(days=1)
        elif self.IsMarketOpen(start_date - timedelta(days=2)):
            return start_date - timedelta(days=2)
        elif self.IsMarketOpen(start_date - timedelta(days=3)):
            return start_date - timedelta(days=3)
        elif self.IsMarketOpen(start_date - timedelta(days=4)):
            return start_date - timedelta(days=4)
        elif self.IsMarketOpen(start_date - timedelta(days=5)):
            return start_date - timedelta(days=5)
        elif self.IsMarketOpen(start_date - timedelta(days=6)):
            return start_date - timedelta(days=6)
        elif self.IsMarketOpen(start_date - timedelta(days=7)):
            return start_date - timedelta(days=7)
        elif self.IsMarketOpen(start_date - timedelta(days=8)):
            return start_date - timedelta(days=8)
        elif self.IsMarketOpen(start_date - timedelta(days=9)):
            return start_date - timedelta(days=9)
        else:
            raise ValueError("FetchEnd: No market open days found in the last 10 days")

    def GetFetchStart(self, fetch_end):
        if self.IsMarketOpen(fetch_end - timedelta(days=1)):
            return fetch_end - timedelta(days=1)
        elif self.IsMarketOpen(fetch_end - timedelta(days=2)):
            return fetch_end - timedelta(days=2)
        elif self.IsMarketOpen(fetch_end - timedelta(days=3)):
            return fetch_end - timedelta(days=3)
        elif self.IsMarketOpen(fetch_end - timedelta(days=4)):
            return fetch_end - timedelta(days=4)
        elif self.IsMarketOpen(fetch_end - timedelta(days=5)):
            return fetch_end - timedelta(days=5)
        elif self.IsMarketOpen(fetch_end - timedelta(days=6)):
            return fetch_end - timedelta(days=6)
        elif self.IsMarketOpen(fetch_end - timedelta(days=7)):
            return fetch_end - timedelta(days=7)
        elif self.IsMarketOpen(fetch_end - timedelta(days=8)):
            return fetch_end - timedelta(days=8)
        elif self.IsMarketOpen(fetch_end - timedelta(days=9)):
            return fetch_end - timedelta(days=9)
        elif self.IsMarketOpen(fetch_end - timedelta(days=10)):
            return fetch_end - timedelta(days=10)
        else:
            raise ValueError("FetchStart: No market open days found in the 10 days before FetchEnd")

    def IsMarketOpen(self, date):
        return self.MarketHoursDatabase.GetExchangeHours(Market.USA, None, SecurityType.Equity).IsDateOpen(date)

    def RefreshFracShares(self):
        self.Log(f"Attempting to fetch Fracshares data on {datetime.now().date()} at {datetime.now().time()}")

        url = "http://www.ibkr.com/download/fracshare_stk.csv"
        content = self.Download(url)
        
        if content:
            old_fracshares = self.frac_symbols.copy()
            
            csv_reader = csv.reader(StringIO(content))
            next(csv_reader)  # Skip header
            
            new_tickers = [row[0] for row in csv_reader if row and row[0].isalpha()]
            
            if new_tickers:
                self.frac_symbols.clear()
                self.frac_symbols.update(new_tickers)
                
                self.Log(f"New Fracshares data found. Total number of Fracshares tickers: {len(self.frac_symbols)}")
                self.Log("--- Start of Fracshares Tickers ---")
                for ticker in self.frac_symbols:
                    self.Log(f"{ticker}")
                self.Log("--- End of Fracshares Tickers ---")
            else:
                self.Log("No new Fracshares data found. Using previous data.")
                self.frac_symbols = old_fracshares
        else:
            self.Log("Failed to download Fracshares data. Using previous data.")

        self.Log(f"Current number of Fracshares tickers: {len(self.frac_symbols)}")

    def OnData(self, data):
        # This method is required but can be left empty for this algorithm
        pass