import gradio as gr
from scipy.optimize import differential_evolution
import pandas as pd
import joblib
import matplotlib.pyplot as plt

# Load trained model
best_model = joblib.load("xgb_model.pkl")

# Predictor names and bounds
predictors = [
    'loaded_drv_time_percycle',
    'empty_drv_time_percycle',
    'Eng_Speed_Ave',
    'empty_stop_time_percycle',
    'loadingstoptime_percycle',
    'loaded_stop_time_percycle'
]

bounds_dict = {
    'loaded_drv_time_percycle': (3, 60),
    'empty_drv_time_percycle': (2, 51),
    'Eng_Speed_Ave': (1051, 1596),
    'empty_stop_time_percycle': (0.2, 24.6),
    'loadingstoptime_percycle': (2, 18),
    'loaded_stop_time_percycle': (0.4, 9)
}

# Optimization function
def optimize_dynamic(
    loaded_drv, empty_drv, eng_speed, empty_stop, loading_stop, loaded_stop,
    fix_loaded_drv, fix_empty_drv, fix_eng_speed, fix_empty_stop, fix_loading_stop, fix_loaded_stop
):
    input_values = {
        'loaded_drv_time_percycle': loaded_drv,
        'empty_drv_time_percycle': empty_drv,
        'Eng_Speed_Ave': eng_speed,
        'empty_stop_time_percycle': empty_stop,
        'loadingstoptime_percycle': loading_stop,
        'loaded_stop_time_percycle': loaded_stop
    }

    fixed_flags = {
        'loaded_drv_time_percycle': fix_loaded_drv,
        'empty_drv_time_percycle': fix_empty_drv,
        'Eng_Speed_Ave': fix_eng_speed,
        'empty_stop_time_percycle': fix_empty_stop,
        'loadingstoptime_percycle': fix_loading_stop,
        'loaded_stop_time_percycle': fix_loaded_stop
    }

    bounds = []
    variable_names = []
    for name in predictors:
        if not fixed_flags[name]:
            bounds.append(bounds_dict[name])
            variable_names.append(name)

    if len(variable_names) == 0:
        pred = best_model.predict(pd.DataFrame([input_values]))[0]
        return f"✅ All inputs fixed.\nPredicted Fuel Rate: {pred:.2f} L/cycle"

    def objective(x):
        current_input = input_values.copy()
        for i, name in enumerate(variable_names):
            current_input[name] = x[i]
        return best_model.predict(pd.DataFrame([current_input]))[0]

    result = differential_evolution(objective, bounds=bounds, seed=42, maxiter=100)

    final_input = input_values.copy()
    for i, name in enumerate(variable_names):
        final_input[name] = result.x[i]

    changes = []
    for name in variable_names:
        original = input_values[name]
        optimized = final_input[name]
        if abs(original - optimized) > 0.01:
            changes.append(f"{name}: {original:.2f} → {optimized:.2f}")
        else:
            changes.append(f"{name}: unchanged ({original:.2f})")

    result_text = "\n".join([f"{k}: {v:.2f}" for k, v in final_input.items()])
    result_text += "\n\n🔁 Optimized Inputs:\n" + "\n".join(changes)
    result_text += f"\n\n⚡️ Predicted Fuel Rate: {result.fun:.2f} L/cycle"

    # Sensitivity analysis for Engine Speed
    # sensitivity_lines = ["\n📊 Suggested RPM vs. Fuel Rate:"]
    start_rpm = int(round(final_input['Eng_Speed_Ave'] / 50.0) * 50)
    end_rpm = min(start_rpm + 300, bounds_dict['Eng_Speed_Ave'][1])
    start_rpm = max(start_rpm, bounds_dict['Eng_Speed_Ave'][0])

    rpm_values = []
    fuel_rates = []

    for rpm in range(start_rpm, end_rpm + 1, 50):
        temp_input = final_input.copy()
        temp_input['Eng_Speed_Ave'] = rpm

        temp_bounds = []
        temp_names = []

        for name in predictors:
            if name != 'Eng_Speed_Ave' and not fixed_flags.get(name, False):
                temp_bounds.append(bounds_dict[name])
                temp_names.append(name)

        def temp_objective(x):
            t_input = temp_input.copy()
            for i, name in enumerate(temp_names):
                t_input[name] = x[i]
            return best_model.predict(pd.DataFrame([t_input]))[0]

        if temp_bounds:
            temp_result = differential_evolution(temp_objective, bounds=temp_bounds, seed=42, maxiter=50)
            fuel_rate = temp_result.fun
        else:
            fuel_rate = best_model.predict(pd.DataFrame([temp_input]))[0]

        rpm_values.append(rpm)
        fuel_rates.append(fuel_rate)

    # Create plot
    plt.figure(figsize=(6, 4))
    plt.plot(rpm_values, fuel_rates, marker='o')
    plt.title("Sensitivity of Fuel Rate to Engine Speed")
    plt.xlabel("Engine Speed (RPM)")
    plt.ylabel("Predicted Fuel Rate (L/cycle)")
    plt.grid(True)

    # Save and return plot image
    plot_path = "sensitivity_plot.png"
    plt.savefig(plot_path)
    plt.close()

    return result_text, plot_path

# Gradio Interface (Interface-style)
interface = gr.Interface(
    fn=optimize_dynamic,
    inputs=[
        gr.Slider(3, 60, value=30, label="Loaded Drive Time"),
        gr.Slider(2, 51, value=23.8, label="Empty Drive Time"),
        gr.Slider(1051, 1596, value=1300, label="Engine Speed"),
        gr.Slider(0.2, 24.6, value=10.0, label="Empty Stop Time"),
        gr.Slider(2, 18, value=2.0, label="Loading Stop Time"),
        gr.Slider(0.4, 9, value=0.4, label="Loaded Stop Time"),
        gr.Checkbox(label="Fix Loaded Drive"),
        gr.Checkbox(label="Fix Empty Drive"),
        gr.Checkbox(label="Fix Engine Speed"),
        gr.Checkbox(label="Fix Empty Stop"),
        gr.Checkbox(label="Fix Loading Stop"),
        gr.Checkbox(label="Fix Loaded Stop"),
    ],
        outputs=[
        gr.Textbox(label="Optimization Result"),
        gr.Image(type="filepath", label="Sensitivity Plot")
    ],
    title="⚙️ Fuel Rate What-If Optimizer",
    description="Perform global optimization of fuel rate with optional fixed inputs and sensitivity on engine RPM."
)

interface.launch()