#region imports
from AlgorithmImports import *
#endregion


class CPIData(PythonData):
    # 12-month unadjusted CPI data
    # Source: https://www.bls.gov/charts/consumer-price-index/consumer-price-index-by-category-line-chart.htm
    # Release dates source: https://www.bls.gov/bls/news-release/cpi.htm

    def GetSource(self,
         config: SubscriptionDataConfig,
         date: datetime,
         isLive: bool) -> SubscriptionDataSource:
        return SubscriptionDataSource("https://www.dropbox.com/s/f02a9htg6pyhf9p/CPI%20data%201.csv?dl=1", SubscriptionTransportMedium.RemoteFile)

    def Reader(self,
         config: SubscriptionDataConfig,
         line: str,
         date: datetime,
         isLive: bool) -> BaseData:

        if not (line.strip()):
            return None

        cpi = CPIData()
        cpi.Symbol = config.Symbol

        try:
            def parse(pct):
                return float(pct[:-1]) / 100

            data = line.split(',')
            cpi.EndTime = datetime.strptime(data[0], "%m%d%Y %H:%M %p")
            cpi["month"] = data[1]
            cpi['all-items'] = parse(data[2])
            cpi['food'] = parse(data[3])
            cpi['food-at-home'] = parse(data[4])
            cpi['food-away-from-home'] = parse(data[5])
            cpi['energy'] = parse(data[6])
            cpi['gasoline'] = parse(data[7])
            cpi['electricity'] = parse(data[8])
            cpi['natural-gas'] = parse(data[9])
            cpi['all-items-less-food-and-energy'] = parse(data[10])
            cpi['commodities-less-food-and-energy-commodities'] = parse(data[11])
            cpi['apparel'] = parse(data[12])
            cpi['new-vehicles'] = parse(data[13])
            cpi['medical-car-commodities'] = parse(data[14])
            cpi['services-less-energy-services'] = parse(data[15])
            cpi['shelter'] = parse(data[16])
            cpi['medical-care-services'] = parse(data[17])
            cpi['education-and-communication'] = parse(data[18])
            
            cpi.Value = cpi['all-items']


        except ValueError:
            # Do nothing
            return None

        return cpi

# region imports
from AlgorithmImports import *
from data import CPIData
from symbol import SymbolData
# endregion

class Modified_60_40(QCAlgorithm):

    def Initialize(self):
        self.SetStartDate(2017, 1, 1)  # Set Start Date
        self.SetEndDate(2019, 1, 1) # Set End Date
        self.SetCash(100000)  # Set Strategy Cash
        self.AddEquity("SPY", Resolution.Minute)
        self.AddEquity("BND", Resolution.Minute)
        self.AddEquity("GLD", Resolution.Minute)
        self.AddEquity("STIP", Resolution.Minute)

        # requesting data for yield curve
        self.yield_curve_symbol = self.AddData(USTreasuryYieldCurveRate, "USTYCR").Symbol

        # historical data for yield curve
        history = self.History(USTreasuryYieldCurveRate, self.yield_curve_symbol, 60, Resolution.Daily)

        # requesting data for CPI
        self.cpi_symbol = self.AddData(CPIData, "CPIData").Symbol

        # historical data for CPI
        #history = self.History(CPIData, self.cpi_symbol, 60, Resolution.Daily)

        # Warm up CPI data history
        self.cpi_lookback = timedelta(days=395)
        self.cpi_history = self.History(self.cpi_symbol, self.StartDate - self.cpi_lookback, self.StartDate)
        if not self.cpi_history.empty:
            self.cpi_history = self.cpi_history.loc[self.cpi_symbol]['value']   

        self.RotationInterval = timedelta(days=30)

        self.first = True

    def OnData(self, data: Slice):
        
        #look for first
        if self.first:
            self.first = False
            self.LastRotationTime = self.Time
            return

        #invest in short term treasuries if not invested
        if not self.Portfolio.Invested:
            self.SetHoldings("STIP", 1)

        #check for null for yield curve
        if not data.ContainsKey(self.yield_curve_symbol):
            return
        
        #define rates
        rates = data[self.yield_curve_symbol]
        
        # Check for null before using the values
        if not (rates.TenYear is not None and rates.ThreeMonth is not None):
            return

        #Update CPI historical data
        if data.ContainsKey(self.cpi_symbol):
            self.rebalance = True
            self.Debug("test")
            self.cpi_history.loc[self.Time] = data[self.cpi_symbol].Value
            self.cpi_history = self.cpi_history.loc[self.cpi_history.index >= self.Time]
        
        #look for rebalance window
        delta = self.Time - self.LastRotationTime
        if delta > self.RotationInterval:

            #reset rebalance window
            self.LastRotationTime = self.Time
            
            # determine if risk on (10Y-3M is >1) or off (10Y-3M is <1)
            if (rates.TenYear - rates.ThreeMonth) > 1:
                risk = 1
            else:
                risk = 0

            # if inflation is > X% then inflation = 1 otherwise inflation = 0 
            # determine if inflationary (expected CPI > 2.5%) or deflationary (expected CPI <2.5%)            
            # Get the current environment and environment history
            
            self.Debug(f"{self.Time}")
            self.Debug(self.cpi_history)
            self.Debug(self.cpi_history[-13])
            if self.cpi_history[-1] > 3:
                inflation = 1  
            else:
                inflation = 0

            if risk == 1 and inflation == 1:
                #risk on, inflationary
                self.Debug("risk on inflationary")
                self.Liquidate()
                self.SetHoldings("SPY", 0.6)
                self.SetHoldings("GLD", 0.4)
            elif risk == 1 and inflation == 0:
                #risk on, deflationary
                self.Debug("risk on deflationary")
                self.Liquidate()
                self.SetHoldings("SPY", 0.6)
                self.SetHoldings("BND", 0.4)
            elif risk == 0 and inflation == 1:
                #risk off, inflationary
                self.Debug("risk off inflationary")
                self.Liquidate()
                self.SetHoldings("STIP", 0.6)
                self.SetHoldings("GLD", 0.4)
            elif risk == 0 and inflation == 0:
                #risk off, deflationary
                self.Debug("risk off deflationary")
                self.Liquidate()
                self.SetHoldings("BND", 0.6)
                self.SetHoldings("STIP", 0.4)    

#region imports
from AlgorithmImports import *
#endregion

  
class SymbolData:
    def __init__(self, algorithm, security, lookback):
        self.algorithm = algorithm
        self.symbol = security.Symbol
        self.security = security

        # Set up consolidators to collect pricing data
        self.consolidator = TradeBarConsolidator(1)
        self.consolidator.DataConsolidated += self.consolidation_handler
        algorithm.SubscriptionManager.AddConsolidator(self.symbol, self.consolidator)

        self.history = pd.Series()
        self.lookback = lookback

        # Get historical data
        history = algorithm.History(self.symbol, self.lookback, Resolution.Daily)
        if not history.empty:
            self.history = history.loc[self.symbol].open

    def consolidation_handler(self, sender: object, consolidated_bar: TradeBar) -> None:
        self.history.loc[consolidated_bar.EndTime] = consolidated_bar.Open
        self.history = self.history.loc[self.history.index > consolidated_bar.EndTime - self.lookback]

    def dispose(self):
        self.algorithm.SubscriptionManager.RemoveConsolidator(self.symbol, self.consolidator)