from AlgorithmImports import *
import datetime
import pandas as pd
from io import StringIO 


class AlertBlueLemur(QCAlgorithm):

    def Initialize(self):
        # set up
        self.SetStartDate(2021, 12, 14)
        self.SetEndDate(2022, 3, 5)
        self.SetCash(1000000)

        # selection will run on the begining of the month
        self.spy = self.AddEquity("SPY", Resolution.Daily).Symbol
        self.SetBenchmark(self.spy)
        self.UniverseSettings.Resolution = Resolution.Daily
        self.AddUniverse("my-blob-universe", self.SelectSymbols)

        # init
        self.backtestSymbolsPerDay = {}
        self.backtestSymbolsPerDayLong = {}
        self.current_universe = []
        self.current_universe_long = []

    def SelectSymbols(self, date):
        
        # handle live mode file format
        if self.LiveMode:
            # fetch the file from dropbox
            str = self.Download("https://www.dropbox.com/s/2l73mu97gcehmh7/daily-stock-picker-live.csv?dl=1")
            # if we have a file for today, return symbols, else leave universe unchanged
            self.current_universe = str.split(',') if len(str) > 0 else self.current_universe
            return self.current_universe

        # backtest - first cache the entire file
        if len(self.backtestSymbolsPerDay) == 0:

            # str = self.Download("https://www.dropbox.com/s/ae1couew5ir3z9y/daily-stock-picker-backtest.csv?dl=1", headers)
            str_ = self.Download("https://snpmarketdata.blob.core.windows.net/qc-backtest/fundamental_stat_arb.csv")
            for line in str_.splitlines():
                data = line.split(',')
                pairs_short = data[1].split("|")
                pairs_long = data[2].split("|")
                self.backtestSymbolsPerDay[data[0]] = pairs_short + pairs_long
                self.backtestSymbolsPerDayLong[data[0]] = pairs_long

        index = date.strftime("%Y-%m-%d")
        if index not in self.backtestSymbolsPerDay:
            return Universe.Unchanged
        # self.current_universe = [Symbol.Create(x, SecurityType.Equity, Market.USA) for x in tickers if x not in ["|", " "]]
        self.current_universe = self.backtestSymbolsPerDay.get(index, self.current_universe)
        self.current_universe_long = self.backtestSymbolsPerDayLong.get(index, self.current_universe_long)

        self.Debug(f" Universe {self.current_universe}")
        self.Debug(f" Universe long {self.current_universe_long}")

        return self.current_universe

    def OnSecuritiesChanged(self, changes):
        self.changes = changes

    def OnData(self, slice):

        # if slice.Bars.Count == 0: return
        # if self.changes is None: return

        # start fresh
        # self.Liquidate()

        # trade
        if len(self.current_universe_long) == 0: return
        weight = 1 / len(self.current_universe_long)
        self.Debug(f"weight {weight}")

        self.Debug(f"Universe long {self.current_universe_long}")
        for symbol in self.current_universe_long:
            self.Debug(f"Symbol 1 {symbol}")
            if not slice.ContainsKey(symbol): continue
            self.Debug(f"Symbol 2 {symbol}")
            if slice[symbol] is None: continue
            self.Debug(f"Symbol 3 {symbol}")
            self.SetHoldings(symbol, weight)

        self.changes = None