from QuantConnect.Data.Custom.TradingEconomics import *

class CustomDataIndexAlgorithm(QCAlgorithm):

    def Initialize(self):
        self.SetStartDate(2013, 8, 1)  # Set Start Date
        self.SetEndDate(2020, 1, 1)  # Set End Date
        self.SetCash(100000)  # Set Strategy Cash
        
        # Add data sources
        self.AddData(IRX, "IRX")
        self.AddData(TNX, "TNX")
        self.AddData(TradingEconomicsCalendar, TradingEconomics.Calendar.UnitedStates.ConsumerConfidence)
        self.consumerConfidence = None
        
        # Setup plots
        confidencePlot = Chart("Confidence Plot")
        confidencePlot.AddSeries(Series('Confidence', SeriesType.Line, 0))
        self.AddChart(confidencePlot)
        differencePlot = Chart("Difference Plot")
        differencePlot.AddSeries(Series('TNX-IRX', SeriesType.Line, 0))
        self.AddChart(differencePlot)


    def OnData(self, data):
        # Get latest consumer confidence
        confidence = data.Get(TradingEconomicsCalendar)
        for c in confidence.Values:
            self.consumerConfidence = c.Actual
        
        # Ensure we have data from all three sources
        if (self.consumerConfidence is None or \
            not data.ContainsKey("IRX") or \
            not data.ContainsKey("TNX")): 
            return

        # Get values for IRX and TNX
        irx = data["IRX"].Close
        tnx = data["TNX"].Close

        # Add to plots
        self.Plot("Confidence Plot", "Confidence", self.consumerConfidence)
        self.Plot("Difference Plot", "TNX-IRX", tnx-irx)

        # Check condition
        if tnx - irx < 0 and self.consumerConfidence > 100:
            self.Log(f"Do something")

def parse(line, index, config):
    """Helper function to parse custom data sources"""
    
    # If first character is not digit, pass
    if not (line.strip() and line[0].isdigit()): return None

    try:
        data = line.split(',')
        
        if data[1] == 'null': return None
        
        tnx = index()
        tnx.Symbol = config.Symbol
        
        tnx.Time = datetime.strptime(data[0], '%Y-%m-%d') + timedelta(hours=20) # Make sure we only get this data AFTER trading day - don't want forward bias.
        tnx.Value = float(data[5])
        
        tnx['Open'] = float(data[1])
        tnx['High'] = float(data[2])
        tnx['Low'] = float(data[3])
        tnx['Close'] = float(data[4])
        tnx['Adj Close'] = float(data[5])
        
        return tnx
    except ValueError:
        # Do nothing, possible error in json decoding
        return None

    
class IRX(PythonData):
    '''Custom Data Type: IRX data from Yahoo Finance - https://www.dropbox.com/s/qa0rlszexthwa8g/IRX.csv?dl=0'''
    
    def GetSource(self, config, date, isLiveMode):
        source = "https://www.dropbox.com/s/qa0rlszexthwa8g/IRX.csv?dl=1"
        return SubscriptionDataSource(source, SubscriptionTransportMedium.RemoteFile)
    
    def Reader(self, config, line, date, isLive):
        return parse(line, IRX, config)
    
class TNX(PythonData):
    '''Custom Data Type: TNX data from Yahoo Finance - https://www.dropbox.com/s/rgmqyuu400cgsad/TNX.csv?dl=0'''
    
    def GetSource(self, config, date, isLiveMode):
        source = "https://www.dropbox.com/s/rgmqyuu400cgsad/TNX.csv?dl=1"
        return SubscriptionDataSource(source, SubscriptionTransportMedium.RemoteFile)
    
    def Reader(self, config, line, date, isLive):
        return parse(line, TNX, config)