import tensorflow as tf
from tensorflow import keras
import numpy as np
import cv2
from PIL import Image
import json

# Constants
IMG_SIZE = 224
MODEL_PATH = "model/plant_disease_efficientnetb0.weights.h5"
CLASS_NAMES_PATH = "model/class_names.json"

# Load CLASS_NAMES
with open(CLASS_NAMES_PATH, "r") as f:
    CLASS_NAMES = json.load(f)

def build_model(num_classes, img_size=IMG_SIZE):
    inputs = keras.Input(shape=(img_size, img_size, 3))
    base_model = keras.applications.EfficientNetB0(
        include_top=False,
        weights=None,
        input_shape=(img_size, img_size, 3)
    )
    base_model.trainable = False
    x = keras.applications.efficientnet.preprocess_input(inputs)
    x = base_model(x, training=False)
    x = keras.layers.GlobalAveragePooling2D()(x)
    x = keras.layers.BatchNormalization()(x)
    x = keras.layers.Dropout(0.3)(x)
    outputs = keras.layers.Dense(num_classes, activation="softmax")(x)
    return keras.Model(inputs, outputs)

NUM_CLASSES = len(CLASS_NAMES)
model = build_model(NUM_CLASSES)
model.load_weights(MODEL_PATH)



# Preprocess image
def preprocess_image(image_path):
    img = tf.keras.preprocessing.image.load_img(
        image_path, target_size=(IMG_SIZE, IMG_SIZE)
    )
    img_array = tf.keras.preprocessing.image.img_to_array(img)
    # Preprocessing will be done by the model's input layer
    return np.expand_dims(img_array, axis=0)


# Inference
def predict_plant_disease(image_path):
    img_array = preprocess_image(image_path)
    preds = model.predict(img_array)[0]

    class_index = int(np.argmax(preds))
    confidence = float(preds[class_index])
    label = CLASS_NAMES[class_index]

    return {label: confidence}


if __name__ == "__main__":
    print("Model loaded. Enter image paths to classify (Ctrl+C to exit):\n")
    try:
        while True:
            image_path = input("Enter image path: ").strip()
            if not image_path:
                print("Please enter a valid path.\n")
                continue
            try:
                result = predict_plant_disease(image_path)
                for label, confidence in result.items():
                    print(f"Label: {label}, Confidence: {confidence:.4f}\n")
            except Exception as e:
                print(f"Error processing image: {e}\n")
    except KeyboardInterrupt:
        print("\n\nExiting...")