import os
import cv2
import xml.etree.ElementTree as ET
import numpy as np
import tensorflow as tf
from tensorflow import keras
import tkinter as tk
from PIL import Image, ImageTk

# Step 1: Parse XML annotations
def parse_xml(xml_file):
    tree = ET.parse(xml_file)
    root = tree.getroot()
    boxes = []
    labels = []
    for obj in root.findall('object'):
        box = obj.find('bndbox')
        xmin = int(box.find('xmin').text)
        ymin = int(box.find('ymin').text)
        xmax = int(box.find('xmax').text)
        ymax = int(box.find('ymax').text)
        label = obj.find('name').text
        boxes.append([xmin, ymin, xmax, ymax])
        labels.append(label)
    return boxes, labels

# Step 2: Preprocess images and annotations
def preprocess_image(image_path):
    image = cv2.imread(image_path)
    image = cv2.resize(image, (224, 224))  # Resize image to desired dimensions
    image = image / 255.0  # Normalize pixel values to [0, 1]
    return image

# Step 3: Build and train model
def build_model():
    model = keras.Sequential([
        keras.applications.MobileNetV2(input_shape=(224, 224, 3), include_top=False),
        keras.layers.GlobalAveragePooling2D(),
        keras.layers.Dense(2, activation='softmax')
    ])
    model.compile(optimizer='adam', loss='sparse_categorical_crossentropy', metrics=['accuracy'])
    return model

def train_model(images_folder, batch_size=32, epochs=10):
    model = build_model()
    image_paths = [os.path.join(images_folder, file) for file in os.listdir(images_folder) if file.endswith('.jpg')]
    for epoch in range(epochs):
        print(f"Epoch {epoch+1}/{epochs}")
        np.random.shuffle(image_paths)
        for i in range(0, len(image_paths), batch_size):
            batch_images = []
            batch_labels = []
            for image_path in image_paths[i:i+batch_size]:
                xml_file = image_path[:-4] + '.xml'
                image = preprocess_image(image_path)
                boxes, labels = parse_xml(xml_file)
                # Convert labels to binary classification: pothole (1) or not pothole (0)
                label = 1 if 'potholes' in labels else 0
                batch_images.append(image)
                batch_labels.append(label)
            batch_images = np.array(batch_images)
            batch_labels = np.array(batch_labels)
            model.train_on_batch(batch_images, batch_labels)
    return model

# Step 4: Create predictor
def predict_image(model, image_path):
    image = preprocess_image(image_path)
    prediction = model.predict(np.expand_dims(image, axis=0))
    return prediction

def show_prediction(image_path, prediction):
    image = Image.open(image_path)
    img_label = f"Prediction: {'Pothole' if prediction[0][1] > 0.6 else 'No Pothole'} (Confidence: {prediction[0][1]*100:.2f}%)"
    root = tk.Tk()
    root.title("Image Prediction")
    img = ImageTk.PhotoImage(image)
    panel = tk.Label(root, image=img)
    panel.image = img
    panel.pack(side="top", padx=10, pady=10)
    label = tk.Label(root, text=img_label)
    label.pack(side="bottom", fill="both", padx=10, pady=10)
    root.mainloop()

def main():
    print("Welcome to Road Pothole Detection System!")
    while True:
        print("\nMenu:")
        print("1. Train Model")
        print("2. Prediction Mode")
        print("3. Exit")
        choice = input("Enter your choice (1/2/3): ")
        
        if choice == '1':
            images_folder = input("Enter the path to the folder containing images and annotations: ")
            trained_model = train_model(images_folder)
            print("Model trained successfully!")
        
        elif choice == '2':
            image_path = input("Enter the path to the image for prediction: ")
            if not os.path.exists(image_path):
                print("Error: Image file not found!")
                continue
            prediction = predict_image(trained_model, image_path)
            show_prediction(image_path, prediction)
        
        elif choice == '3':
            print("Exiting...")
            break
        
        else:
            print("Invalid choice. Please enter a valid option (1/2/3).")

if __name__ == "__main__":
    main()