import gradio as gr
import numpy as np
from skimage.color import rgb2gray
from skimage.transform import resize
import matplotlib.pyplot as plt
import matplotlib.patches as patches
from joblib import load
import tempfile
import pandas as pd
from sklearn.preprocessing import StandardScaler
import os

# Load models
mlp_model = load("mlp_bbox_model.pkl")
model_x = load("linear_model_x.pkl")
model_y = load("linear_model_y.pkl")
model_w = load("linear_model_w.pkl")
model_h = load("linear_model_h.pkl")

# Load annotations and fit scaler (for MLP model)
# IMPORTANT: Load the scaler saved during training so inverse_transform uses
# the exact same mean/std. If target_scaler.pkl is missing, fall back to
# fitting on the full CSV (less accurate — save the scaler from training!).
annotation_data = pd.read_csv("image_annotation.csv")
if os.path.exists("target_scaler.pkl"):
    target_scaler = load("target_scaler.pkl")
else:
    target_scaler = StandardScaler()
    target_scaler.fit(annotation_data[['x', 'y', 'width', 'height']])

# Preprocess function
def preprocess_image(image, image_size=(64, 64)):
    if image.ndim == 3 and image.shape[2] == 4:  # RGBA -> RGB
        image = image[:, :, :3]
    if image.ndim == 3 and image.shape[2] == 3:
        image = rgb2gray(image)
    image_resized = resize(image, image_size, anti_aliasing=True)
    return image_resized.flatten(), image

# Predict and draw function
def predict(image, model_type):
    x_input, original_image = preprocess_image(image)
    x_input = x_input.reshape(1, -1)

    if model_type == "MLP":
        y_scaled_pred = mlp_model.predict(x_input)
        y_pred = target_scaler.inverse_transform(y_scaled_pred)[0]
    else:  # Linear Regression
        x_pred = model_x.predict(x_input)[0]
        y_pred_ = model_y.predict(x_input)[0]
        w_pred = model_w.predict(x_input)[0]
        h_pred = model_h.predict(x_input)[0]
        y_pred = [x_pred, y_pred_, w_pred, h_pred]

    # Clip width and height
    y_pred[2] = np.clip(y_pred[2], 1, original_image.shape[1])
    y_pred[3] = np.clip(y_pred[3], 1, original_image.shape[0])

    # Attempt to get filename (for ground truth matching)
    # NOTE: Gradio passes numpy arrays; filename matching requires the original
    # filename to be known externally. Ground truth display is unavailable here.
    gt_row = pd.DataFrame()

    # Draw image and boxes
    fig, ax = plt.subplots()
    ax.imshow(original_image, cmap='gray')

    # Predicted box in red
    rect_pred = patches.Rectangle((y_pred[0], y_pred[1]), y_pred[2], y_pred[3],
                                  linewidth=2, edgecolor='red', facecolor='none', label="Prediction")
    ax.add_patch(rect_pred)

    # Ground truth box in green
    if not gt_row.empty:
        x_gt = gt_row.iloc[0]['x']
        y_gt = gt_row.iloc[0]['y']
        w_gt = gt_row.iloc[0]['width']
        h_gt = gt_row.iloc[0]['height']
        rect_gt = patches.Rectangle((x_gt, y_gt), w_gt, h_gt,
                                    linewidth=2, edgecolor='green', facecolor='none', label="Ground Truth")
        ax.add_patch(rect_gt)

    ax.legend()
    plt.axis('off')

    # Save to temporary file
    tmpfile = tempfile.NamedTemporaryFile(suffix=".png", delete=False)
    tmp_path = tmpfile.name
    tmpfile.close()
    plt.savefig(tmp_path, bbox_inches='tight', pad_inches=0)
    plt.close(fig)

    return tmp_path

# Gradio interface
interface = gr.Interface(
    fn=predict,
    inputs=[
        gr.Image(label="Upload an image", interactive=True, type="numpy"),
        gr.Radio(choices=["MLP", "Linear Regression"], label="Select Model")
    ],
    outputs=gr.Image(type="filepath", label="Predicted vs Ground Truth"),
    title="Object Localization: Predicted vs Ground Truth",
    description="Upload an image and select a model to predict the bounding box. Ground truth is shown in green if available."
)

interface.launch()