# %%
# imports
import gradio as gr
import matplotlib.pyplot as plt
import numpy as np
import pandas as pd
import math

# %%
# import data
data = pd.read_excel("graph_data.xlsx",
                     sheet_name = "Sheet1",
                     header = 0
)

# clean data
# rename year column
data = data.rename(columns = {"Unnamed: 0": "Year"})
# drop rows without data
# data = data.drop([0,1,2,3,4,5,6,7,8,9,10,11])
# keep columns of interest
data = data[[
    "Year",
    "Total Revenues",
    "Debt Balance",
    "Revenues ex SS OASDI",
    "Average Rate on Federal Debt", 
    "GDP", 
    "Net Interest"
    ]]

data.set_index("Year", inplace=True)

print(data)


# %%
baseline_interest_rate = 3.67
#baseline_revenues = 15928.73
baseline_cagr_revenues = 4.02
total_revenues_2024 = data.loc[2024, "Total Revenues"]

def plot_interest_coverage(interest_rate, cagr_revenues):

    # calculate the yearly increase in the interest rate based on the projected interest rate in 2054
    interest_rate_yearly_increase = (interest_rate - baseline_interest_rate) / (2054 - 2025) / 100
    # calculate the yearly increase in revenues based on the projected interest rate in 2054
    # revenues_yearly_increase = (revenues - baseline_revenues) / (2054 - 2025)

    # add a baseline net interest / revenues column
    data["Net Interest / Revenues (Baseline)"] = data["Net Interest"] / data["Total Revenues"]

    # add a baseline net interest / revenues ex SS OASDI column
    data["Net Interest / Revenues ex SS OASDI (Baseline)"] = data["Net Interest"] / data["Revenues ex SS OASDI"]

    # add a baseline SS OASDI revenues column
    data["SS OASDI Revenues"] = data["Total Revenues"] - data["Revenues ex SS OASDI"]

    # add a projected average rate on federal debt column
    data["Average Rate on Federal Debt (Projected)"] = np.where(
        data.index.astype(int) < 2025, 
        data["Average Rate on Federal Debt"], 
        data["Average Rate on Federal Debt"] + (interest_rate_yearly_increase * (data.index.astype(int) - 2024)))
    
    # add a projected revenues column
    data["Total Revenues (Projected)"] = np.where(
        data.index.astype(int) < 2025, 
        data["Total Revenues"], 
        total_revenues_2024 * ((1 + (cagr_revenues / 100)) ** (data.index.astype(int) - 2024)))
    # data["Total Revenues"] + (revenues_yearly_increase * (data.index.astype(int) - 2025)))
    
    print(data)
    
    # add a projected revenues ex SS OASDI column
    data["Revenues ex SS OASDI (Projected)"] = np.where(
        data.index.astype(int) < 2025, 
        data["Revenues ex SS OASDI"], 
        data["Total Revenues (Projected)"] - data["SS OASDI Revenues"])
    
    # add a projected interest / revenues column
    data["Net Interest / Revenues (Projected)"] = np.where(
        data.index.astype(int) < 2025, 
        data["Net Interest / Revenues (Baseline)"], 
        data["Average Rate on Federal Debt (Projected)"] * data["Debt Balance"] / data["Total Revenues (Projected)"])
    
    # add a projected interest / revenues ex SS OASDI column
    data["Net Interest / Revenues ex SS OASDI (Projected)"] = np.where(
        data.index.astype(int) < 2025, 
        data["Net Interest / Revenues ex SS OASDI (Baseline)"], 
        data["Average Rate on Federal Debt (Projected)"] * data["Debt Balance"] / data["Revenues ex SS OASDI (Projected)"])

    # Create the plot
    plt.figure(figsize = (10, 4.8))
    # plot average rate on federal debt
    plt.plot(
        data.index, 
        data["Average Rate on Federal Debt"], 
        color = "Green", 
        label = "Average Rate on Federal Debt"
    )
    # plot average rate on federal debt projected
    if cagr_revenues != baseline_cagr_revenues or interest_rate != baseline_interest_rate:
        plt.plot(
            data.index, 
            data["Average Rate on Federal Debt (Projected)"], 
            color = "Green", 
            label = "Average Rate on Federal Debt (Projected)",
            linestyle = "--"
        )
    else:
        plt.plot(
            data.index, 
            data["Average Rate on Federal Debt"], 
            color = "Green", 
            label = "Average Rate on Federal Debt (Projected)",
            linestyle = "--"
        )

    # plot interest / revenues (baseline)
    plt.plot(
        data.index, 
        data["Net Interest / Revenues (Baseline)"], 
        color = "Blue", 
        label = "Net Interest / Revenues (Baseline)"
    )
    # plot interest / revenues (projected)
    if cagr_revenues != baseline_cagr_revenues or interest_rate != baseline_interest_rate:
        plt.plot(
            data.index, 
            data["Net Interest / Revenues (Projected)"], 
            color = "Blue", 
            label = "Net Interest / Revenues (Projected)",
            linestyle = "--"
        )
    else:
        plt.plot(
            data.index, 
            data["Net Interest / Revenues (Baseline)"], 
            color = "Blue", 
            label = "Net Interest / Revenues (Projected)",
            linestyle = "--"
        )

    
    # plot interest / revenues ex ss oasdi (baseline)
    plt.plot(
        data.index, 
        data["Net Interest / Revenues ex SS OASDI (Baseline)"], 
        color = "Orange", 
        label = "Net Interest / Revenues ex SS OASDI (Baseline)"
    )
    # plot interest / revenues ex ss oasdi (projected)
    if cagr_revenues != baseline_cagr_revenues or interest_rate != baseline_interest_rate:
        plt.plot(
            data.index, 
            data["Net Interest / Revenues ex SS OASDI (Projected)"], 
            color = "Orange", 
            label = "Net Interest / Revenues ex SS OASDI (Projected)",
            linestyle = "--"
        )
    else:
        plt.plot(
            data.index, 
            data["Net Interest / Revenues ex SS OASDI (Baseline)"], 
            color = "Orange", 
            label = "Net Interest / Revenues ex SS OASDI (Projected)",
            linestyle = "--"
        )
    
    plt.title("Interest as Share of Revenues Through 2054")
    plt.legend(loc = "upper left")
    plt.ylim(0, 1.05)
    plt.yticks(np.arange(0, 1.1, 0.1))
    plt.xticks(range(1940,2055,4), rotation = 45)
    plt.axvline(x = 2024.5, color = "black", linestyle = '--')
    plt.grid(True, axis = 'y', linestyle = '--', linewidth = 0.7)
    
    # Save the plot to a file
    plt.savefig("interest_coverage.png")
    plt.close()
    
    # Return the path to the image file
    return "interest_coverage.png"

# %%
interest_rate_lowerbound = 0
interest_rate_upperbound = 10

# revenues_lower_bound = 10000
# revenues_upper_bound = 20000

cagr_revenues_lower_bound = 0
cagr_revenues_upper_bound = 8

with gr.Blocks() as interface:
    # Create the image output
    graph = gr.Image(type="filepath", label = "Graph", value = plot_interest_coverage(baseline_interest_rate, baseline_cagr_revenues))
    
    # Create the slider input below the image for projected interest rate
    interest_rate_slider = gr.Slider(
        interest_rate_lowerbound, 
        interest_rate_upperbound, 
        step = .1,
        value = baseline_interest_rate, 
        label = "2054 Projected Average Interest Rate on Federal Debt"
    )

    # Create the slider input below the image for projected revenues
    cagr_revenues_slider = gr.Slider(
        cagr_revenues_lower_bound, 
        cagr_revenues_upper_bound, 
        step = 0.01,
        value = baseline_cagr_revenues, 
        label = "Compound Annual Growth Rate of Revenues through 2054"
    )

    gr.Markdown('<p style="font-size:11px;">Source: CBO June 2024  An Update to the Budget and Economic Outlook: 2024 to 2034, '
    'March 2024 report The Long-Term Budget Outlook: 2024 to 2054, and author\'s calculations. Scenario of higher interest rate and revenues '
    'is author\'s calculations. Data points for 2035-2054 calculated to build on CBO\'s June 2024 10-year update and be consistent '
    'with the March 2024 long-term report. Historical Social Security OASDI payroll tax revenue from Table 4-3 of the SSA\'s '
    'Trust Fund Data, and projections from the CBO\'s August 2024 Long Term Projections for Social Security.</p>')

    # Set up an action that updates the graph when the interest rate slider value changes
    interest_rate_slider.change(
        plot_interest_coverage, 
        inputs = [interest_rate_slider, cagr_revenues_slider], 
        outputs = graph
    )

    # Set up an action that updates the graph when the revenues slider value changes
    cagr_revenues_slider.change(
        plot_interest_coverage, 
        inputs = [interest_rate_slider, cagr_revenues_slider], 
        outputs = graph
    )

# Launch the interface
interface.launch(share = True)