app.py

import gradio as gr
import pandas as pd
import requests
import json
import time
import base64
from io import StringIO
import getpass
import os

pd.set_option('display.max_columns', None)
pd.set_option('display.max_rows', None)

DATABRICKS_INSTANCE = 'https://136286767906206.6.gcp.databricks.com'
ACCESS_TOKEN = os.getenv("ACCESS_TOKEN_")
JOB_ID = os.getenv("JOB_ID_")

# Function that controls the visibility of the panels
def check_words(user_text, input_text):
    # If the entered text is "user" and "casa", show Panel 2 and hide Panel 1
    if user_text == os.getenv("USER_1_") and input_text == os.getenv("PASS_1_"):
        return gr.update(visible=True), gr.update(visible=False)#, user_text#, gr.update(value=""), gr.update(value="")
    elif user_text == os.getenv("USER_2_") and input_text == os.getenv("PASS_2_"):
        return gr.update(visible=True), gr.update(visible=False)#, user_text#, gr.update(value=""), gr.update(value="")
    elif user_text == os.getenv("USER_3_") and input_text == os.getenv("PASS_3_"):
        return gr.update(visible=True), gr.update(visible=False)#, user_text#, gr.update(value=""), gr.update(value="")
    elif user_text == os.getenv("USER_4_") and input_text == os.getenv("PASS_4_"):
        return gr.update(visible=True), gr.update(visible=False)#, user_text#, gr.update(value=""), gr.update(value="")
    else:
        return gr.update(visible=False), gr.update(visible=True)#, user_text#, user_text, input_text  # Keep Panel 2 hidden, show Panel 1

# Function to hide Panel 2
def hide_panel():
    return gr.update(visible=False)  # Hide Panel 2

# Function to show Panel 1
def show_panel_1():
    return gr.update(visible=True)  # Show Panel 1

# Wrapper function to hide Panel 2 and show Panel 1 when btn_2 is clicked
def hide_and_show_panel():
    return hide_panel(), show_panel_1(), gr.update(value=""), gr.update(value="")

#def update_message(request: gr.Request):
#    return f"Welcome, {request.username}"

def get_run_status(run_id):
    url = f"{DATABRICKS_INSTANCE}/api/2.1/jobs/runs/get?run_id={run_id}"
    headers = {
        "Authorization": f"Bearer {ACCESS_TOKEN}"
    }
    response = requests.get(url, headers=headers)
    return response.json()

def get_databricks_file(file_path):
    # Construct the DBFS URL for the file
    dbfs_url = f"{DATABRICKS_INSTANCE}/api/2.0/dbfs/read"
    headers = {
        "Authorization": f"Bearer {ACCESS_TOKEN}",
        "Content-Type": "application/json"
    }
    data = {
        "path": file_path
    }
    response = requests.get(dbfs_url, headers=headers, json=data)
    #response.raise_for_status()  # Raise an exception for HTTP errors
    content = response.json()
    file_data = base64.b64decode(content['data'])
    return file_data

def get_month_label(i):
    return f'M{i+1}'

#data_predictions_ = get_databricks_file("dbfs:/dbfs/FileStore/forecast_alleviatetax_predictions.csv")
#df = pd.read_csv(StringIO(data_predictions_.decode('utf-8')))
#df['revenue'] = (df['vintage_unique_cases'] / df['predicted_monthly_payment_rate']).round(2)
#df['Month'] = [get_month_label(i) for i in range(len(df))]
df = pd.DataFrame()
#df_2 = pd.DataFrame()
    
def update_table(start_date, end_date, window, user_text):
    #############################################################################################
    try:
        DATABRICKS_INSTANCE = 'https://136286767906206.6.gcp.databricks.com'
        ACCESS_TOKEN = os.getenv("ACCESS_TOKEN_")
        JOB_ID = os.getenv("JOB_ID_")
        
        # Parameters to pass to the job
        params = {
            "events_api_base_url": "https://beta.datascore.ai",
            "month_train_starting": str(start_date),
            "month_for_prediction": str(end_date),
            "moving_average_period": str(window),
            "username": str(user_text) 
        }
        
        url = f"{DATABRICKS_INSTANCE}/api/2.1/jobs/run-now"
        headers = {
            "Authorization": f"Bearer {ACCESS_TOKEN}",
            "Content-Type": "application/json"
        }
        data = {
            "job_id": JOB_ID,
            "notebook_params": params  # Pass parameters here
        }
        response = requests.post(url, headers=headers, data=json.dumps(data))
        #return response.json()
        # Run the Job
        #run_info = post_run_job(JOB_ID, params)
        run_info = response.json()
        run_id = run_info['run_id']
        #print(f"Run ID: {run_id}")
        
        while True:
            status = get_run_status(run_id)
            state = status['state']['life_cycle_state']
            if state in ['TERMINATED', 'SKIPPED', 'FAILED', 'ERROR']:
                print(f"Job finished with state: {state}")
                break
            elif state in ['PROGRES']:
                print(f"Job progress: {state}")
            time.sleep(30)
        
        print(f"Run ID: {run_id}")
        #data_payments = get_databricks_file("dbfs:/dbfs/FileStore/forecast_alleviatetax_payments_"+str(user_text).split('@')[0]+".csv")
        #data_metrics = get_databricks_file("dbfs:/dbfs/FileStore/forecast_alleviatetax_metrics.csv")
        data_predictions = get_databricks_file("dbfs:/dbfs/FileStore/forecast_alleviatetax_predictions_"+str(user_text).split('@')[0]+".csv")
        #df_payments = pd.read_csv(StringIO(data_payments.decode('utf-8')))
        #new_columns = ['vintage', 'vintage_unique_cases'] + [f'M{col}' for col in df_payments.columns[2:]] 
        #df_metrics = pd.read_csv(StringIO(data_metrics.decode('utf-8')))
        df_predictions = pd.read_csv(StringIO(data_predictions.decode('utf-8')))
        #df_predictions['revenue'] = (df_predictions['vintage_unique_cases'] / df_predictions['predicted_monthly_payment_rate']).round(2)
        #df_predictions['Month'] = [get_month_label(i) for i in range(len(df_predictions))]
    except Exception as e:
        return {error: str(e)}
    #############################################################################################
    global df  # Use global variable
    #global df_2
    df = df_predictions.copy()
    #df_2 = df_payments.copy() 
    return df#, df_2

def save_csv(file_name):
    global df  # Use global variable
    if file_name:
        df.to_csv(file_name, index=False)
        return f"File '{file_name}' saved!"
    else:
        return "Please set file name!"

with gr.Blocks(fill_height=True) as demo:
    # Create the first panel (visible from the start)
    with gr.Column(visible=True) as panel_1:
        gr.HTML(f"<h1>Forecast UI</h1>")
        user_text = gr.Textbox(label="Username:", placeholder="Type here...", scale=0)
        input_text = gr.Textbox(label="Password:", placeholder="Type here...", scale=0, type="password")
    
        # First button, visible from the start
        btn_1 = gr.Button("Login", scale=0)

    
    with gr.Column(visible=False) as panel_2:
        #m = gr.Markdown()
        gr.HTML(f"<h1>Forecast UI</h1>")
        with gr.Row():
        #    demo.load(update_message, None, gr.Markdown())
        #    logout_button = gr.Button("Logout", link="/logout", scale=0)
            # Add a second button inside Panel 2 to hide the panel and show Panel 1 again
            #tu = gr.Textbox(interactive=False, scale=0) 
            btn_2 = gr.Button("Logout", visible=True, scale=0)
            btn_2.click(hide_and_show_panel, outputs=[panel_2, panel_1, user_text, input_text])  # Hide Panel 2 and show Panel 1
        with gr.Row():
            with gr.Column():               
                start_input = gr.Textbox(label="Start Date:", placeholder="Enter date (e.g., 2023-09-01)", max_lines=1)
                gr.Markdown("Start of the data range of months that will be used as input by the model for revenue prediction.<br><br>The end month is automatically set to one month before the prediction month.")
            with gr.Column():  
                prediction_input = gr.Text(label="Month Prediction:", placeholder="Enter date (e.g., 2024-10-31)", max_lines=1)
                gr.Markdown("Case creation month to predict its revenue.")
            with gr.Column():  
                range_input = gr.Slider(3, 12, 6, label="Window / Moving Average Period")
                gr.Markdown("Window = 3-period/months, the predictive model reacts more quickly to recent monthly payment fluctuations, but it may also include more noise.<br><br>Window = 12-period/months, the Forecast adjusts more slowly and is less sensitive to small fluctuations, making it more reliable, but also slower to react to sharp changes.")
        
        table_1 = gr.DataFrame(value=df, label="Predictions (consult 'revenue' column):")
        btn_update = gr.Button("Run Forecast")
       
        #table_2 = gr.DataFrame(value=df_2, label="Forecast Inputs:")
        btn_update.click(fn=update_table, inputs=[start_input, prediction_input, range_input, user_text], outputs=[table_1#,table_2
                                                                                                                  ]) 

    
    # Configure the buttons and the panel visibility
    btn_1.click(check_words, inputs=[user_text, input_text], outputs=[panel_2, panel_1])

demo.launch()