import gradio as gr
import matplotlib.pyplot as plt
import os


def update_num_variable(dataset_choice):
    reset_value = 1
    if not dataset_choice :
        return gr.update(minimum=1, maximum=10, step=1, value=reset_value)
    if dataset_choice in ['Electricity', 'Traffic']:
        num_variable_range = (1,10,1)
    elif dataset_choice in ['ETTh1', 'ETTh2', 'ETTm1', 'ETTm2', 'National_illness']:
        num_variable_range = (1,7,1)
    elif dataset_choice == "Weather":
        num_variable_range = (1,21,1)
    elif dataset_choice == "Exchange":
        num_variable_range = (1,8,1)
    return gr.update(minimum=num_variable_range[0], maximum=num_variable_range[1], step=num_variable_range[2],
                     value=reset_value)
def update_time_horizon(dataset_choice):
    print(f"[DEBUG] update_time_horizon called with dataset_choice: {dataset_choice}")

    if not dataset_choice:
        return gr.update(choices=[], value=None)

    if dataset_choice == "National_illness":
        choices = [24, 36, 48, 60]
    else:
        choices = [96, 192, 336, 720]

    default_value = max(choices)
    print(f"[DEBUG] Returning time horizon choices: {choices}, default: {default_value}")
    return gr.update(choices=choices, value=default_value)

def update_starting_point(model_choice, dataset_choice, time_horizon):
    reset_value = 0
    starting_point_range = (0, 300, 20)  # Standardwert
    if not model_choice or not dataset_choice or not time_horizon:
        return gr.update(minimum=0, maximum=300, step=20, value = reset_value)

    if model_choice == "NLinear":
        if dataset_choice == "Electricity":
            if time_horizon == 96:
                starting_point_range = (0, 160, 20)
            elif time_horizon in [192, 336, 720]:
                starting_point_range = (0, 140, 20)
        elif dataset_choice == "ETTh1":
            if time_horizon in [96, 192]:
                starting_point_range = (0, 80, 20)
            else:
                starting_point_range = (0, 60, 20)
        elif dataset_choice == "ETTh2":
            if time_horizon in [96, 192]:
                starting_point_range = (0, 80, 20)
            else:
                starting_point_range = (0, 60, 20)
        elif dataset_choice == "Traffic":
            if time_horizon in [96, 192]:
                starting_point_range = (0, 100, 20)
            else:
                starting_point_range = (0, 80, 20)
        elif dataset_choice == "Weather":
            if time_horizon in [96, 192]:
                starting_point_range = (0, 320, 20)
            else:
                starting_point_range = (0, 300, 20)
        elif dataset_choice == "Exchange":
            if time_horizon in [96, 192]:
                starting_point_range = (0,40,20)
            else:
                starting_point_range = (0, 20, 20)
        elif dataset_choice == "National_illness":
            starting_point_range = (0, 0, 0)
        elif dataset_choice == "ETTm2":
            if time_horizon == 720:
                starting_point_range = (0, 320, 20)
            else:
                starting_point_range = (0,340,20)
        elif dataset_choice == "ETTm1":
            if time_horizon == 720:
                starting_point_range = (0, 320, 20)
            else:
                starting_point_range = (0,340,20)
    else:
        if dataset_choice == "National_illness":
            if time_horizon == 24:
                starting_point_range = (0, 160, 20)
            elif time_horizon in [36, 48]:
                starting_point_range = (0, 140, 20)
            elif time_horizon == 60:
                starting_point_range = (0, 120, 20)
        elif dataset_choice == "Electricity":
            starting_point_range = (0,4320, 720)
        elif dataset_choice in ["ETTh1", "ETTh2"]:
            if time_horizon == 96:
                starting_point_range = (0, 2700, 100)
            elif time_horizon == 192:
                starting_point_range = (0, 2600, 100)
            elif time_horizon == 336:
                starting_point_range = (0, 2300, 100)
            elif time_horizon == 720:
                starting_point_range = (0, 2100, 100)
        elif dataset_choice in ["ETTm1", "ETTm2"]:
            starting_point_range = (0,10800, 720)
        elif dataset_choice == "Exchange":
            if time_horizon == 96:
                starting_point_range = (0, 1300, 100)
            elif time_horizon == 192:
                starting_point_range = (0, 1100, 100)
            elif time_horizon == 336:
                starting_point_range = (0, 1000, 100)
            elif time_horizon == 720:
                starting_point_range = (0, 700, 100)
        elif dataset_choice == "Traffic":
            if time_horizon == 96:
                starting_point_range = (0, 3300, 100)
            elif time_horizon == 192:
                starting_point_range = (0, 3200, 100)
            elif time_horizon == 336:
                starting_point_range = (0, 3000, 100)
            elif time_horizon == 720:
                starting_point_range = (0, 2100, 100)
        elif dataset_choice == "Weather":
            starting_point_range = (0, 9360, 720)
        else:
            starting_point_range = (0, 300, 20)
    return gr.update(minimum = starting_point_range[0], maximum = starting_point_range[1], step = starting_point_range[2], value = reset_value)



def find_csv_file(dataset_choice, time_horizon, look_back_window):
    dir = "../../../Desktop/Project/Models/LTBoost"
    for root, dirs, files in os.walk(dir):
        for dir_name in dirs:
            if dataset_choice.lower() in dir_name.lower():
                full_dir_path = os.path.join(root, dir_name)
                for file_name in os.listdir(full_dir_path):
                    if f"pl_{time_horizon}" in file_name:
                        csv_file_path = os.path.join(full_dir_path, file_name)
                        if os.path.exists(csv_file_path):
                            return csv_file_path
    return None

def plot_forecast_vs_ground_truth(df_forecast, df_ground_truth, dataset_choice, model_choice, lookback_window,time_horizon, starting_point, num_variable):
    plt.figure(figsize=(12, 6))

    plt.plot(df_ground_truth, label='Forecast')
    plt.plot(df_forecast, label='Ground Truth')
    # Marking the transition between look-back window and forecast horizon
    transition_point = lookback_window
    plt.axvline(x=transition_point, color='black', linestyle='-')
    plt.title(f'Forecast vs Ground Truth for {dataset_choice} using {model_choice} Model')
    #plt.xlabel('Time')
    #plt.ylabel('Value')
    plt.legend()
    plt.grid(True)
    # Create directory path based on dataset name, lookback window, and time horizon
    save_dir = f"results_ltboost/{dataset_choice}_sl{lookback_window}_pl{time_horizon}/{num_variable}"
    os.makedirs(save_dir, exist_ok=True)  # Ensure directory exists, create if not

    # Save the plot as PNG file
    path = os.path.join(save_dir, f"img_{starting_point}.png")
    plt.savefig(path)

    plt.close()
    return path

def find_image_path(dataset_choice, model_choice, time_horizon, look_back_window, starting_point, num_variable):
    directory = "results_nlinear"
    for root, dirs, files in os.walk(directory):
        for dir_name in dirs:
            if dataset_choice.lower() in dir_name.lower() and model_choice in dir_name and f"pl{time_horizon}" in dir_name and f"sl{look_back_window}" in dir_name:
                variable_dir = os.path.join(root, dir_name, str(num_variable))
                if os.path.exists(variable_dir):
                    image_path = os.path.join(variable_dir, f"{starting_point}.png")
                    if os.path.exists(image_path):
                        return image_path
    return None

def find_image_path_LTBoost(dataset_choice, time_horizon, look_back_window, starting_point, num_variable):
    directory = "results_ltboost"
    for root, dirs, files in os.walk(directory):
        for dir_name in dirs:
            # Check if the folder name contains the desired dataset, time horizon, and look-back window
            if (dataset_choice.lower() in dir_name.lower() and
                    f"pl{time_horizon}" in dir_name and
                    f"sl{look_back_window}" in dir_name):

                variable_dir = os.path.join(root, dir_name, str(num_variable))

                if os.path.exists(variable_dir):
                    image_path = os.path.join(variable_dir, f"img_{starting_point}.png")

                    if os.path.exists(image_path):
                        return image_path

    return None


def perform_forecasting(dataset_choice, model_choice, time_horizon, starting_point, num_variable):
    if model_choice == "NLinear":
        look_back_window = 104 if dataset_choice == "National_illness" else 336
        img = find_image_path(dataset_choice, model_choice, time_horizon, look_back_window, starting_point, num_variable)
        return img if img else None  
    
    elif model_choice == "LTBoost":
        look_back_window = 104 if dataset_choice == "National_illness" else (336 if dataset_choice in ["Exchange", "ETTh1"] else 720)
        img = find_image_path_LTBoost(dataset_choice, time_horizon, look_back_window, starting_point, num_variable)
        return img if img else None



def main():
    theme = gr.themes.Base(
        primary_hue="blue",
        secondary_hue="blue",
        neutral_hue="gray"
    )
    theme.set(
        body_background_fill="*primary_50",
        body_background_fill_dark="*checkbox_background_color_focus",
        body_text_color_dark="white",
        body_text_color="*neutral_800",
        background_fill_secondary_dark="*checkbox_border_color_hover",
        block_background_fill="*background_fill_primary",
        block_background_fill_dark="*neutral_800",
        block_border_color_dark="*primary_100",
        block_border_width_dark="4px",
        block_border_width="4px",
        block_border_color="*secondary_200"
    )
    with gr.Blocks(theme=theme) as demo:
        gr.Markdown("""
            <h1 style="text-align: center; font-size: 1.5em;">LTBoost</h1>
            <div style="font-size: 16px; letter-spacing: 1.2px; line-height: 1.8; text-align: justify;">
                LTBoost: Boosted Hybrids of Ensemble Linear and Gradient Algorithms for the Long-term Time Series Forecasting<br>
                <strong>LTBoost repository:</strong> <a href="https://github.com/hubtru/LTBoost" target="_blank">LTBoost</a><br>
                You can choose between the dataset, model, time horizon, and starting point. <br>
                <strong>Options:</strong> <br>
                <ul>
                    <li><strong>Datasets:</strong> Electricity, Exchange Rate, Traffic, Weather, ILI, 4 ETTs</li>
                    <li><strong>Models:</strong> LTBoost or NLinear</li>
                    <li><strong>Time horizon:</strong> {96, 192, 336, 720} (different for ILI)</li>
                    <li><strong>Starting point:</strong> {0, ..., end}</li>
                    <li><strong>Variate:</strong> 10 last variates (including the OT which is always the last)</li>
                </ul>
                Number of figures for each of the dataset depends on the dataset_trainset length and the starting_point_resolution (offset).<br>
                Due to the high number of the figures the starting_point_resolution has been limited to {20, 100, 720}.<br>
            </div>
            """)

        with gr.Row():
            with gr.Column():
                dataset_choice = gr.Dropdown(
                    ["Electricity", "Traffic", "Weather", "Exchange", "National_illness", "ETTh1", "ETTh2", "ETTm1", "ETTm2"],
                    label="Dataset Choice"
                )
                model_choice = gr.Dropdown(["NLinear", "LTBoost"], label="Model Choice", value= "LTBoost")
                time_horizon = gr.Dropdown([96, 192, 336, 720], label="Time Horizon", value=720)
                starting_point = gr.Slider(minimum=0, maximum=300, step=20, label="Starting Point", value= 0)
                num_variable = gr.Slider(minimum=1,maximum=10,step=1,label="Variable")

                input_components = [dataset_choice, model_choice, time_horizon, starting_point, num_variable]
                dataset_choice.change(update_time_horizon, inputs=dataset_choice, outputs=time_horizon)

                time_horizon.change(update_starting_point, inputs=[model_choice, dataset_choice, time_horizon], outputs=starting_point)
                model_choice.change(update_starting_point, inputs=[model_choice, dataset_choice, time_horizon], outputs=starting_point)
                dataset_choice.change(update_starting_point, inputs=[model_choice, dataset_choice, time_horizon], outputs=starting_point)

                
                dataset_choice.change(update_time_horizon, inputs=dataset_choice, outputs=time_horizon)
                dataset_choice.change(update_starting_point, inputs=[model_choice, dataset_choice, time_horizon], outputs=starting_point)
                model_choice.change(update_starting_point, inputs=[model_choice, dataset_choice, time_horizon], outputs=starting_point)
                time_horizon.change(update_starting_point, inputs=[model_choice, dataset_choice, time_horizon], outputs=starting_point)

    


                gr.ClearButton(input_components)

            with gr.Column():
                # Automatically update the output image whenever any input component changes
                output_png = gr.Image(label="Forecast Results")
                for component in input_components:
                    component.change(perform_forecasting, inputs=input_components, outputs=output_png)
        with gr.Row():
            with gr.Column():
                gr.Markdown("Examples for the user input:")
                gr.Examples([
                    ['ETTh1', 'NLinear', 720, 0],

                    ['ETTh2', 'NLinear', 720, 0],

                    ['ETTm1', 'NLinear', 720, 0],

                    ['ETTm2', 'NLinear', 720, 0],

                    ['Weather', 'NLinear', 720, 0],

                    ['Traffic', 'NLinear', 720, 0],

                    ['Electricity', 'NLinear', 720, 0],

                    ['ETTh1', 'LTBoost', 720, 0],

                    ['ETTh2', 'LTBoost', 720, 0],

                    ['ETTm1', 'LTBoost', 720, 0],

                    ['ETTm2', 'LTBoost', 720, 0],

                    ['Weather', 'LTBoost', 720, 0],

                    ['Traffic', 'LTBoost', 720, 0],

                    ['Electricity', 'LTBoost', 720, 0],

                            ],

                            inputs=input_components)



    demo.launch()

if __name__ == "__main__":
    main()