from AlgorithmImports import *
from datetime import timedelta, datetime, time
import pandas as pd


class DataTracker(object):
    def __init__(self, symbol: Symbol, prior_close):
        self.Symbol = symbol
        self.prior_close = prior_close
        self.pre_market_vol = 0.0
        self.market_open_price = 0.0
        self.asset = None

    def add_pre_market_vol(self, volume):
        self.pre_market_vol += volume

    def __str__(self):
        return f"{self.Symbol.Value} - prior close: {self.prior_close} - pre vol: {self.pre_market_vol} - market open: {self.market_open_price}"


class MicroCapTrades(QCAlgorithm):

    def Initialize(self):
        self.SetStartDate(2021, 12, 31)  # Set Start Date
        self.SetEndDate(2021, 12, 31)  # Set Start Date
        self.SetCash(60000)  # Set Strategy Cash
        self.SetTimeZone("America/New_York")
        self.UniverseSettings.Resolution = Resolution.Minute
        self.UniverseSettings.DataNormalizationMode = DataNormalizationMode.Raw
        self.UniverseSettings.ExtendedMarketHours = True
        self.AddUniverse(self.Coarse, self.Fine)
        self.SetBrokerageModel(BrokerageName.InteractiveBrokersBrokerage)
        self.lot_size_dollar = 5000.0
        self._data_tracker = {}
        self.spy = self.AddEquity("SPY", Resolution.Minute, Market.USA, True, 1, extendedMarketHours=True)
        self.spy.SetDataNormalizationMode(DataNormalizationMode.Raw)
        self.Schedule.On(self.DateRules.EveryDay("SPY"), self.TimeRules.BeforeMarketClose("SPY", 1),  self.CloseAllPosition)

    def Coarse(self, coarse):
        self.Debug(f"Coarse is run with {len([x.Symbol for x in coarse])} tickers entered")
        filtered = [
            x.Symbol for x in coarse if x.HasFundamentalData and x.DollarVolume > 0 and x.Value >= 1 and x.Value <= 20]
        self._data_tracker = {
            x.Symbol: DataTracker(x.Symbol, x.Value) for x in coarse if x.Symbol in filtered}
        return filtered

    def Fine(self, fine):
        self.Debug(f"Fine is run with {len([x.Symbol for x in fine])} tickers entered")
        small_cap = [x.Symbol for x in fine if x.MarketCap <= 300 * 1e6]
        for s in list(self._data_tracker.keys()):
            if s not in small_cap:
                self._data_tracker.pop(s)
        for symbol in small_cap:
            self._data_tracker[symbol].asset = self.AddEquity(symbol, Resolution.Minute, Market.USA, True, 1, True)
            self._data_tracker[symbol].asset.SetDataNormalizationMode(DataNormalizationMode.Raw)
        return []

    def OnData(self, data: Slice):
        tradeBars = data.Bars
        
        ticker_fed = [s for s in tradeBars.Keys]
        
        if self.Time.minute == 0:
            self.Debug(f"OnData with {len(ticker_fed)} tickers in data feed")

        if self.Time.hour < 9 or (self.Time.hour == 9 and self.Time.minute <= 30):
            for s in self._data_tracker:
                if s in tradeBars:
                    self._data_tracker[s].add_pre_market_vol(
                        tradeBars[s].Volume)
            return

        if self.Time.hour == 9 and self.Time.minute == 31:
            has_no_data = [
                self._remove_tickers(s) for s in list(self._data_tracker.keys()) if s not in tradeBars]
            has_event = [
                self._remove_tickers(s) for s in list(self._data_tracker.keys()) if data.Dividends.ContainsKey(s) or data.Splits.ContainsKey(s) or data.Delistings.ContainsKey(s) or data.SymbolChangedEvents.ContainsKey(s)]
            no_gap_up = [
                self._remove_tickers(s) for s in list(self._data_tracker.keys()) if tradeBars[s].Open < self._data_tracker[s].prior_close * 1.19]
            no_vol_1mm = [
                self._remove_tickers(s) for s in list(self._data_tracker.keys()) if self._data_tracker[s].pre_market_vol < 1e6]

            if self._data_tracker:
                self.Debug(f"{len(self._data_tracker)} tickers are selected to trade")

    def OnSecuritiesChanged(self, changes: SecurityChanges):
        active_symbol = [s.Symbol for s in self.ActiveSecurities.Values]
        self.Debug(f"OnSecuritiesChanged called with {len(active_symbol)} tickers in ActiveSecurities")

    def CloseAllPosition(self):
        if self.Portfolio.Invested:
            self.Liquidate()

    def _remove_tickers(self, s):
        self.RemoveSecurity(s)
        return self._data_tracker.pop(s)