import numpy as np
import joblib
import gradio as gr
import matplotlib.pyplot as plt

# Load the pre-trained model
model = joblib.load("XGBoost_best_of_10.joblib")  

def predict_hc(mt_content, liquid_limit, plastic_limit, specific_gravity, 
              initial_wc, temperature, dry_density_values):
    warnings = []
    predictions = []
    single_pred = None
    
    # Process dry density values
    try:
        density_values = [float(x.strip()) for x in dry_density_values.split(',')]
        if len(density_values) < 4 or len(density_values) > 8:
            raise ValueError("Please enter 4 to 8 values separated by commas")
    except ValueError as e:
        return f"Error in dry density input: {str(e)}", None, None, None
    
    # Fixed features (same for all predictions)
    fixed_features = np.array([
        mt_content,
        liquid_limit,
        plastic_limit,
        specific_gravity,
        initial_wc,
        temperature
    ])
    
    # Define trained ranges for soil properties
    trained_ranges = {
        "Montmorillonite Content": (33.60, 92.00),
        "Liquid Limit": (50.00, 500.00),
        "Plastic Limit": (20.00, 200.00),
        "Specific Gravity": (2.65, 2.82),
        "Initial Water Content": (5.40, 21.10),
        "Temperature": (20.00, 40.00),
        "Dry Density": (1.30, 2.05),
    }
    
    # Check input values against trained ranges
    input_values = {
        "Montmorillonite Content": mt_content,
        "Liquid Limit": liquid_limit,
        "Plastic Limit": plastic_limit,
        "Specific Gravity": specific_gravity,
        "Initial Water Content": initial_wc,
        "Temperature": temperature,
    }
    
    for feature, value in input_values.items():
        min_val, max_val = trained_ranges[feature]
        if value < min_val:
            warnings.append(f"⚠️ {feature} is below the trained range ({min_val}-{max_val}).")
        elif value > max_val:
            warnings.append(f"⚠️ {feature} is above the trained range ({min_val}-{max_val}).")
    
    # Make predictions for all density values
    for density in density_values:
        # Check dry density range
        min_density, max_density = trained_ranges["Dry Density"]
        if density < min_density:
            warnings.append(f"⚠️ Dry Density {density} is below the trained range ({min_density}-{max_density}).")
        elif density > max_density:
            warnings.append(f"⚠️ Dry Density {density} is above the trained range ({min_density}-{max_density}).")
        
        # Create input array
        input_array = np.append(fixed_features, density).reshape(1, -1)
        pred = model.predict(input_array)[0]
        predictions.append(pred)
    
    # If only one density value provided, show single prediction
    if len(density_values) == 1:
        single_pred = predictions[0]
    
    # Create plot if multiple density values
    plot_path_png = None
    if len(density_values) > 1:
        fig, ax = plt.subplots(figsize=(5, 4))
        ax.plot(density_values, predictions, marker="o", linestyle="-", color="blue", markersize=8)
        ax.set_xlabel("Dry Density [g/cm³]")
        ax.set_ylabel("Saturated Hydraulic Conductivity (HC) [-]")
        ax.grid(True, linestyle="--", alpha=0.5)
        
        plot_path_png = "HC_plot.png"
        fig.savefig(plot_path_png, dpi=600, bbox_inches="tight")
        plt.close(fig)
    
    # Format output
    pred_text = ""
    if single_pred is not None:
        pred_text += f"Predicted Saturated HC: {single_pred:.4f} [-]\n\n"
    else:
        pred_text += "Predicted Saturated HC values:\n"
        for d, p in zip(density_values, predictions):
            pred_text += f"• At ρ = {d:.2f} g/cm³: HC = {p:.4f} [-]\n"
        pred_text += "\n"
    
    if warnings:
        pred_text += "\n".join(warnings)
    
    # Team text
    marketing_text = (
        "**Developed by:**  \nMuntasir Shehab*, Reza Taherdangkoo and Christoph Butscher  \n\n"
        "**Institution:**  \nTU Bergakademie Freiberg, Institute of Geotechnics  \n"
        "Gustav-Zeuner-Str. 1, Freiberg, 09599, Germany"
    )
    
    return pred_text, plot_path_png, marketing_text

# Input components
inputs = [
    gr.Number(label="Montmorillonite Content [%]", value=60, minimum=5.60, maximum=96.00, step=0.1),
    gr.Number(label="Liquid Limit [%]", value=150, minimum=50.00, maximum=500.00, step=1),
    gr.Number(label="Plastic Limit [%]", value=50, minimum=20.00, maximum=200.00, step=1),
    gr.Number(label="Specific Gravity", value=2.7, minimum=2.15, maximum=2.92, step=0.01),
    gr.Number(label="Initial Water Content [%]", value=10, minimum=5.40, maximum=21.10, step=0.1),
    gr.Number(label="Temperature [°C]", value=25, minimum=20.00, maximum=40.00, step=0.1),
    gr.Textbox(label="Dry Density values [g/cm³] (4-8 values, comma separated)", 
              value="1.3, 1.5, 1.7, 1.9", 
              placeholder="Enter 4-8 values separated by commas (e.g., 1.3, 1.5, 1.7, 1.9)"),
]

# Output components
outputs = [
    gr.Textbox(label="Prediction Results & Warnings"),
    gr.Image(label="HC vs Dry Density Plot", type="filepath"),
    gr.Markdown(label="About the Team")
]

# Launch interface
gr.Interface(
    fn=predict_hc,
    inputs=inputs,
    outputs=outputs,
    title="Saturated Hydraulic Conductivity Prediction Model",
    description="Predict saturated hydraulic conductivity (HC) based on soil properties. Enter 4-8 dry density values to see the trend."
).launch(share=True)