import tensorflow as tf
import joblib
import numpy as np
from PIL import Image, ImageDraw
import gradio as gr

# 1. Load pretrained models
cnn_model = tf.keras.models.load_model('best_cnn_model.keras')   # Load CNN model
ann_model = tf.keras.models.load_model('best_ann_model.keras')   # Load ANN model

# Attempt to load linear regression models (assumed dict of four regressors)
try:
    lr_models = joblib.load('linear_regression_models.pkl')
except Exception as e:
    print(f"[Warning] Failed to load linear regression models: {e}")
    lr_models = None

# 2. Define prediction function with auto‑swap for manual bbox
def predict_bbox(model_choice, input_image, stop_x1, stop_y1, stop_x2, stop_y2):
    errors = []

    # Auto-swap manual bbox so that x0<=x1, y0<=y1
    mx0, mx1 = sorted([stop_x1, stop_x2])
    my0, my1 = sorted([stop_y1, stop_y2])

    # Preprocess image
    gray = input_image.convert('L')
    resized = gray.resize((128, 128))
    norm = np.array(resized) / 255.0

    # Model prediction branch
    if model_choice in ['CNN', 'ANN']:
        arr = norm.reshape(1, 128, 128, 1)
        if model_choice == 'CNN':
            pred = cnn_model.predict(arr)[0]
        else:
            pred = ann_model.predict(arr)[0]
    else:
        if lr_models is None:
            return "Error: Linear Regression models not loaded.", input_image
        flat = norm.flatten().reshape(1, -1)
        x1_p = lr_models['x1'].predict(flat)[0]
        y1_p = lr_models['y1'].predict(flat)[0]
        x2_p = lr_models['x2'].predict(flat)[0]
        y2_p = lr_models['y2'].predict(flat)[0]
        pred = np.array([x1_p, y1_p, x2_p, y2_p])

    # Scale predictions back to original image size
    orig_w, orig_h = input_image.size
    x1_p, y1_p, x2_p, y2_p = pred
    x0_p = int(min(x1_p, x2_p) * orig_w)
    x1_p = int(max(x1_p, x2_p) * orig_w)
    y0_p = int(min(y1_p, y2_p) * orig_h)
    y1_p = int(max(y1_p, y2_p) * orig_h)

    # Validate predicted bbox ordering
    if x1_p < x0_p:
        errors.append("Predicted bbox error: x2 < x1")
    if y1_p < y0_p:
        errors.append("Predicted bbox error: y2 < y1")

    # Create drawable copy
    img_draw = input_image.copy()
    draw = ImageDraw.Draw(img_draw)

    # Draw predicted bounding box in green (only if no critical error)
    if not errors:
        draw.rectangle([x0_p, y0_p, x1_p, y1_p], outline='green', width=2)
    else:
        # If prediction invalid, overlay error text
        err_text = "\n".join(errors)
        draw.text((10, 10), err_text, fill='red')

    # Draw manual Stop Sign box in red
    draw.rectangle([mx0, my0, mx1, my1], outline='red', width=2)

    # Prepare output text
    coords_text = (
        f"Model: {model_choice}\n"
        f"Predicted bbox: ({x0_p}, {y0_p}), ({x1_p}, {y1_p})\n"
        f"Manual Stop Sign bbox: ({mx0}, {my0}), ({mx1}, {my1})"
    )
    if errors:
        coords_text = "Errors in prediction:\n" + "\n".join(errors) + "\n\n" + coords_text

    return coords_text, img_draw

# 3. Build Gradio interface
iface = gr.Interface(
    fn=predict_bbox,
    inputs=[
        gr.Dropdown(choices=['CNN', 'ANN', 'Linear Regression'], label='Select Model'),
        gr.Image(type='pil', label='Upload Test Image'),
        gr.Number(label='Stop Sign x1 or x2'),
        gr.Number(label='Stop Sign y1 or y2'),
        gr.Number(label='Stop Sign x2 or x1'),
        gr.Number(label='Stop Sign y2 or y1'),
    ],
    outputs=[
        gr.Textbox(label='Results / Errors'),
        gr.Image(type='pil', label='Annotated Image'),
    ],
    title='Stop Sign Detection with Auto‑Swap Manual BBox',
    description=(
        'Manual bbox will auto‑swap if coords reversed. '
        'Predicted bbox in green; manual bbox in red. '
        'Prediction errors overlaid as text.'
    )
)

if __name__ == "__main__":
    iface.launch()