from AlgorithmImports import *

class ObjsectStoreUniverseSelectionAlgorithm(QCAlgorithm):

    def Initialize(self):
        self.SetStartDate(2023, 12, 5)
        self.SetEndDate(2023, 12, 7)
        # Upload files to ObjectStore programmatically
        if not self.ObjectStore.ContainsKey("trades_2023_12_06.csv"):
            content = self.Download("https://www.dropbox.com/scl/fi/v7hbemm1haa9xuz9o9bto/trades_2023_12_06.csv?rlkey=ia2sbot56q9962s5oxwot7xi5&dl=1")
            self.ObjectStore.Save("trades_2023_12_06.csv", content)
        self.symbols = []
        self.universe = self.AddUniverse(StockDataSource, "my-stock-data-source", Resolution.Daily, self.Filter)
        benchmark = self.AddEquity("SPY", Resolution.Minute).Symbol
        self.Schedule.On(self.DateRules.EveryDay(benchmark), self.TimeRules.BeforeMarketClose(benchmark, 30), self.Trade)

        self.Settings.MinimumOrderMarginPortfolioPercentage = 0.01

    def Filter(self, data):
        '''The filter returns all symbols'''
        self.symbols = [x.Symbol for x in data if x.Value > 0]
        return self.symbols

    def Trade(self):
        if not self.symbols:
            return

        weight = 1/len(self.symbols)
        for symbol in self.symbols:
            quantity = self.CalculateOrderQuantity(symbol, weight)
            if quantity != 0:
                self.MarketOnCloseOrder(symbol, quantity)

class StockDataSource(PythonData):
    '''
    This data source looks for data in the ObjectStore
    https://www.quantconnect.com/docs/v2/cloud-platform/organizations/object-store
    trades_2023_12_06.csv contains data in the format:
    ,symbol,date,fiscal_date_ending,predictions,spx_return_coef
    0,PNC US,2023-12-06,2023-12-31,1,1.1914152093473287
    1,USB US,2023-12-06,2023-12-31,1,1.5266932633773835
    '''
    market_hours_data_base = MarketHoursDatabase.FromDataFolder()
    hours = market_hours_data_base.GetExchangeHours(Market.USA, "[*]", SecurityType.Equity)    
    symbols = []
    def GetSource(self, config, date, isLiveMode):
        key = 'trades_2023_12_06.csv'
        return SubscriptionDataSource(key, SubscriptionTransportMedium.ObjectStore)

    def Reader(self, config, line, date, isLiveMode):
        csv = line.split(',')
        if not csv[0]:
            return None
        localDateTime = datetime.strptime(csv[2], "%Y-%m-%d") - timedelta(0)
        while not self.hours.IsDateOpen(localDateTime):
            localDateTime = localDateTime - timedelta(1)
        ticker = csv[1].split(' ')[0]
        sid = SecurityIdentifier.GenerateEquity(ticker, Market.USA, mappingResolveDate=date)
        data = StockDataSource()
        data.Symbol = Symbol(sid, ticker)
        data.EndTime = localDateTime
        data.Value = float(csv[4])
        return data