import traceback
import numpy as np
import pandas as pd

from QuantConnect.Data.Custom import Quandl
from QuantConnect.Python import PythonQuandl
from QuantConnect.Data.Custom.USTreasury import *

# ref 
# https://www.quantconnect.com/forum/discussion/2445/using-quandl-data-w-python/p1
# https://github.com/QuantConnect/Lean/blob/master/Algorithm.Python/AltData/USTreasuryYieldCurveRateAlgorithm.py
# https://www.quandl.com/data/USTREASURY/YIELD-Treasury-Yield-Curve-Rates
# https://www.quantconnect.com/docs/algorithm-reference/importing-custom-data
# https://fred.stlouisfed.org/graph/fredgraph.csv?id=DGS10

from QuantConnect.Python import PythonData
from QuantConnect.Data import SubscriptionDataSource
from datetime import datetime, timedelta
import decimal


class MyYield(PythonData):
    def GetSource(self, config, date, isLiveMode):
        url = "na"
        return SubscriptionDataSource(url,SubscriptionTransportMedium.RemoteFile)
    def Reader(self, config, line, date, isLiveMode):
        if not (line.strip() and line[0].isdigit()): return None
        inst = MyYield()
        inst.Symbol = config.Symbol
        try:
            # Example File Format:
            #DATE,DGS10,DGS2
            #1954-07-01,0.80,0.80
            data = line.split(',')
            # # Make sure we only get this data AFTER trading day - don't want forward bias.
            inst.Time = datetime.strptime(data[0], '%Y-%m-%d')+timedelta(hours=20)
            inst.Value = decimal.Decimal(data[1])
        except:
            return None
        return inst
        
class MyYield2yr(MyYield):
    def GetSource(self, config, date, isLiveMode):
        url = "https://fred.stlouisfed.org/graph/fredgraph.csv?id=DGS2"
        return SubscriptionDataSource(url,SubscriptionTransportMedium.RemoteFile)
class MyYield10yr(MyYield):
    def GetSource(self, config, date, isLiveMode):
        url = "https://fred.stlouisfed.org/graph/fredgraph.csv?id=DGS10"
        return SubscriptionDataSource(url,SubscriptionTransportMedium.RemoteFile)

class QuandlYield2yr(PythonQuandl):
    def __init__(self):
        self.ValueColumnName = "2 yr"
class QuandlYield10yr(PythonQuandl):
    def __init__(self):
        self.ValueColumnName = "10 yr"
        
class QuandlAlgo(QCAlgorithm):

    def Initialize(self):
        
        self.SetStartDate(datetime.now().date() - timedelta(100))
        self.SetEndDate(datetime.now().date() - timedelta(1))

        self.SetCash(10000)
        self.SetBrokerageModel(AlphaStreamsBrokerageModel())
        
        self.spy = self.AddEquity('SPY', Resolution.Daily).Symbol
        self.vix = self.AddData(Quandl,"CHRIS/CBOE_VX1", Resolution.Daily).Symbol
        
        # method one
        self.yieldCurve = self.AddData(USTreasuryYieldCurveRate, "USTYCR", Resolution.Daily).Symbol
        self.History(USTreasuryYieldCurveRate, self.yieldCurve, 1, Resolution.Daily)
        
        # method two        
        self.yieldCurveTwo = self.AddData(QuandlYield2yr,"USTREASURY/YIELD", Resolution.Daily).Symbol
        self.yieldCurveTen = self.AddData(QuandlYield10yr,"USTREASURY/YIELD", Resolution.Daily).Symbol
        self.History(self.yieldCurveTwo, 1, Resolution.Daily)
        self.History(self.yieldCurveTen, 1, Resolution.Daily)
        
        # method three
        self.myyieldTwo = self.AddData(MyYield2yr, "MYYIELD").Symbol
        self.myyieldTen = self.AddData(MyYield10yr, "MYYIELD").Symbol
        self.History(self.myyieldTwo, 1, Resolution.Daily)
        self.History(self.myyieldTen, 1, Resolution.Daily)
        
        self.twoyear0 = -1
        self.tenyear0 = -1
        
        self.twoyear1 = -1
        self.tenyear1 = -1
        
        self.twoyear2 = -1
        self.tenyear2 = -1
        
        self.Schedule.On(self.DateRules.EveryDay(self.spy), self.TimeRules.AfterMarketOpen(self.spy, 1), self.MyBalance)
        self.Schedule.On(self.DateRules.EveryDay(self.spy), self.TimeRules.BeforeMarketClose(self.spy, 1), self.MyPlot)
        
        # Create custom charts
        myplot = Chart('vix')
        myplot.AddSeries(Series('vix', SeriesType.Line, 0))
        myplot.AddSeries(Series('ma', SeriesType.Line, 0))
        
        myplot = Chart('yield')
        myplot.AddSeries(Series('2yr0', SeriesType.Line, 0))
        myplot.AddSeries(Series('10yr0', SeriesType.Line, 0))
        myplot.AddSeries(Series('2yr1', SeriesType.Line, 0))
        myplot.AddSeries(Series('10yr1', SeriesType.Line, 0))
        myplot.AddSeries(Series('2yr2', SeriesType.Line, 0))
        myplot.AddSeries(Series('10yr2', SeriesType.Line, 0))
        
    def MyBalance(self):
        self.SetHoldings(self.spy,1.0)
        
    def OnData(self, data):

        if data.ContainsKey(self.myyieldTwo):
            self.twoyear2 = data[self.myyieldTwo].Value
            
        if data.ContainsKey(self.myyieldTen):
            self.tenyear2 = data[self.myyieldTen].Value

        if data.ContainsKey(self.yieldCurveTwo):
            self.twoyear1 = data[self.yieldCurveTwo].Value
            
        if data.ContainsKey(self.yieldCurveTen):
            self.tenyear1 = data[self.yieldCurveTen].Value
            
        if data.ContainsKey(self.yieldCurve):
            rates = data[self.yieldCurve]
            
            # Check for None before using the values
            if rates.TenYear is None or rates.TwoYear is None:
                pass
            else:
                self.twoyear0 = rates.TwoYear
                self.tenyear0 = rates.TenYear
                
    
    def MyPlot(self):
        
        self.Plot('yield', '2yr0', self.twoyear0)
        self.Plot('yield', '10yr0', self.tenyear0)
        
        self.Plot('yield', '2yr1', self.twoyear1)
        self.Plot('yield', '10yr1', self.tenyear1)
        
        self.Plot('yield', '2yr2', self.twoyear2)
        self.Plot('yield', '10yr2', self.tenyear2)
        
        data = self.History(self.vix,timedelta(days = 20),Resolution.Daily)
        if len(data) > 10:
            values = data['close'].values
            self.Plot('vix', 'vix', values[-1])
            self.Plot('vix', 'ma', np.nanmean(values))