# Grab each next quarter
for quarter in range(1,5):
    print(quarter % 4 + 1)

#Define a function that processes the data given the year/quarter we want to use as of

def process_data_yield_curve(yield_curve: DataFrame, year: int, quarter: int) -> DataFrame:
    """
    The function for processing of the yield curve data

    Parameters
    ----------
    yield_curve : DataFrame
        The raw yield curve data
    year : int
        The year for the report
    quarter : int
        The quarter for the report

    Returns
    -------
    DataFrame
        Yield table

    """
    # Copy the yield curve to avoid overwriting
    yield_curve = yield_curve.copy()
    
    #Find the next quarter
    end_quarter = quarter % 4 + 1
    
    #Convert to the month
    end_month = (end_quarter-1) * 3 + 1
    
    # Find the end year which is the same unless the quarter is 4
    end_year = year
    if quarter == 4:
        end_year += 1
    # Finally create the datetime and subtract one day from it for the last day
    end_date = datetime(end_year, end_month, 1) - pd.Timedelta("1D")
    
    # Index the yield 
    yield_curve = yield_curve.loc[:end_date]
    
    # Create the yield table
    yield_table = create_yield_table(yield_curve, end_date)
    
    return yield_table
yield_table = process_data_yield_curve(yield_curve, 2019, 1)
print(yield_table)

#Define a similar function for the real estate data

def process_data_real_estate(data: DataFrame, year: int, quarter: int) -> Tuple[DataFrame, DataFrame, DataFrame]:
    """
    Function which process the real estate data

    Parameters
    ----------
    data : DataFrame
        Real estate data
    year : int
        The year for the last quarter to use
    quarter : int
        The last quarter to use

    Returns
    -------
    Tuple[DataFrame, DataFrame, DataFrame]
        1. Dataframe of returns
        2. Quarterly returns
        3. Quarterly returns adjusted for inflation

    """
    # Divide all data by the starting values at the end of 2000
    returns = data / data.loc[(1999,4)]
    
    # Only use 2000 and on
    returns = returns.loc[(2000,1):]
    
    # Add adjustments for inflation
    returns["Real Estate Inflation Adjusted"] = returns["Real Estate"] / returns["CPI"]
    
    # Cut the data to be only up to the last quarter/year requested
    returns = returns.loc[:(year, quarter)]
    
    # Get the percentage changes for real estate
    quarterly_returns = data["Real Estate"].pct_change()
    quarterly_returns = quarterly_returns.loc[(2000,1):(year, quarter)]
    quarterly_returns = quarterly_returns.unstack()
    
    # Get the adjusted quarterly returns
    quarterly_returns_adj = returns["Real Estate Inflation Adjusted"].pct_change()
    
    #Because there is no baseline to compare to, the percent change will be negative in the first year
    #We can set it to be equal to the index - 1
    quarterly_returns_adj.iloc[0] = returns["Real Estate Inflation Adjusted"].iloc[0] - 1
    quarterly_returns_adj = quarterly_returns_adj.unstack()
    
    return returns, quarterly_returns, quarterly_returns_adj

# Run our function and check the outputs
returns_RE, quarterly_returns_RE, quarterly_returns_adj_RE = process_data_real_estate(real_estate_data, 2019, 1)
print(returns_RE.head(5))
print()
print(quarterly_returns_RE.head(5))
print()
print(quarterly_returns_adj_RE.head(5))