from AlgorithmImports import *
from QuantConnect.DataSource import *
from statistics import stdev
import time
# endregion

class MeasuredTanTermite(QCAlgorithm):

    def Initialize(self):
        self.SetStartDate(2021, 6, 10)
        self.SetCash(100000)
        self.rebalanceTime = datetime.min
        self.AddUniverse(self.CoarseSelectionFilter, self.FineSelectionFilter)
        self.UniverseSettings.Resolution = Resolution.Daily #resolution of the securities in the universe, default is minute
        self.Data = {}
        self.ValueCheck= {}
        self.banned_tickers = {}

        self.spy = self.AddEquity("SPY", Resolution.Daily).Symbol

        #initializes bond yeild data from FRED
        self.fred_peak_to_trough = self.AddData(Fred, Fred.ICEBofAML.USHighYieldMasterIIOptionAdjustedSpread, Resolution.Daily).Symbol

        self.nasdaqCode = "FRED/UNRATE"
        self.unrate = self.AddData(NasdaqDataLink, self.nasdaqCode, Resolution.Minute).Symbol

    

    def OnData(self, data):
        if data.ContainsKey(self.nasdaqCode):
            self.Plot("Unemployment Rate", "Value", self.Securities[self.nasdaqCode].Price)

        for symbol in self.Data.keys():
            symbolData = self.Data[symbol]

            if not symbolData.IsReady: 
               continue

            
            

            #self.Log("Symbol " + str(symbol))

            #checks for availability of FRED data by using SPY trading days, plots spread
            if data.ContainsKey(self.fred_peak_to_trough) and data.ContainsKey(self.spy):
                peak_to_trough = data.Get(Fred, self.fred_peak_to_trough).Value
                self.Plot('FRED HY Spread', 'Bond Spread', peak_to_trough)

 # The CoarseFundamental class provides access to the following information that you can base your selection upon: self.Market, self.DollarVolume (for the day), self.Volume (for the day), self.HasFundamentalData, self.PriceFactor, self.SplitFactor, self.PriceScaleFactor, self.AdjustedPrice, self.EndTime.
 
    def CoarseSelectionFilter(self, coarse): #updates at midnight during backtesting, and around 6-7am in live trading, based on previous day data
        
        if self.Time <= self.rebalanceTime: #checks if it has been 1 day since rebalancing
            return self.Universe.Unchanged #if not do not change the universe
        self.rebalanceTime = self.Time + timedelta(35) #set rebalance time in days

        symbols_by_price = [c for c in coarse if c.HasFundamentalData and len(c.Symbol.Value) <= 4]
        self.filteredByPrice = symbols_by_price[:10]

        #for c in symbols_by_price[:30]: 
	    #    self.Log(",Symbol," + str(c.Symbol.Value) + ",DollarVolume," + str(c.DollarVolume) +  ",Price," + str(c.Price) + ",Volume," + str(c.Volume))
        
        return [c.Symbol for c in self.filteredByPrice]

    def FineSelectionFilter(self, fine):
        sortedByMarketCap = sorted(fine, key=lambda c: c.MarketCap)
        symbols_by_marketcap = [c for c in sortedByMarketCap if c.MarketCap > 0 and c.CompanyReference.IsREIT is False]
        self.filteredBymarketcap = symbols_by_marketcap

        #return self.filteredBymarketcap
        return [c.Symbol for c in self.filteredBymarketcap]

    def OnSecuritiesChanged(self, changes):
        for security in changes.AddedSecurities:
            symbol = security.Symbol
            if symbol not in self.Data:
                self.Data[symbol] = SymbolData(self, symbol)
        
        for security in changes.RemovedSecurities:
            symbol = security.Symbol
            if symbol in self.Data:
                symbolData = self.Data.pop(symbol, None)
                self.SubscriptionManager.RemoveConsolidator(symbol, symbolData.consolidator)

class SymbolData:
    def __init__(self, algorithm, symbol):
        self.algorithm = algorithm
        self.symbol = symbol
        self.Bars = RollingWindow[TradeBar](7)
        self.consolidator = TradeBarConsolidator(timedelta(days=1))
        self.consolidator.DataConsolidated += self.OnDataConsolidated
        algorithm.SubscriptionManager.AddConsolidator(symbol, self.consolidator)
    
    def OnDataConsolidated(self, sender, bar):
        self.Bars.Add(bar)
    
    @property
    def IsReady(self):
        return self.Bars.IsReady

class NasdaqCustomColumns(NasdaqDataLink):
    '''Custom nasdaq data type for setting customized value column name. Value column is used for the primary trading calculations and charting.'''
    def __init__(self):
        self.ValueColumnName = "Unrate"