import os
import torch
import torchvision.transforms as transforms
import torchvision.models.detection as detection
from flask import Flask, request, render_template, send_from_directory, url_for, Response
from PIL import Image
import cv2
import numpy as np

# Initialize Flask app
app = Flask(__name__)
UPLOAD_FOLDER = os.path.abspath(os.path.join("deployment", "static", "uploads"))
OUTPUT_FOLDER = os.path.abspath(os.path.join("deployment", "outputs"))

# Ensure directories exist
os.makedirs(UPLOAD_FOLDER, exist_ok=True)
os.makedirs(OUTPUT_FOLDER, exist_ok=True)

# Define class names
CLASS_NAMES = {
    0: "vehicle", 1: "bicycle", 2: "bus", 3: "car", 4: "lorry"
}
NUM_CLASSES = len(CLASS_NAMES) + 1  # Include background class

# Load Faster R-CNN model
def load_model(model_path, num_classes):
    model = detection.fasterrcnn_resnet50_fpn(weights=None, num_classes=num_classes)
    model.load_state_dict(torch.load(model_path, map_location='cpu'))
    model.eval()
    return model

model_path = os.path.join("models", "fasterrcnn_model.pth")
model = load_model(model_path, NUM_CLASSES)

# Image preprocessing
def preprocess_image(image):
    transform = transforms.Compose([transforms.ToTensor()])
    return transform(image).unsqueeze(0)

# Draw predictions
def draw_predictions(image, prediction):
    boxes, scores, labels = prediction['boxes'], prediction['scores'], prediction['labels']
    for i in range(len(boxes)):
        if scores[i] > 0.5:  # Confidence threshold
            x1, y1, x2, y2 = map(int, boxes[i].tolist())
            class_name = CLASS_NAMES.get(labels[i].item(), f"Class {labels[i].item()}")
            cv2.rectangle(image, (x1, y1), (x2, y2), (0, 255, 0), 2)
            cv2.putText(image, f"{class_name}: {scores[i]:.2f}", (x1, y1 - 10),
                        cv2.FONT_HERSHEY_SIMPLEX, 0.5, (0, 255, 0), 2)
    return image

# Home route (upload and detect)
@app.route('/', methods=['GET', 'POST'])
def upload_and_predict():
    if request.method == 'POST':
        file = request.files['file']
        if file:
            filename = file.filename
            upload_path = os.path.join(UPLOAD_FOLDER, filename)
            result_path = os.path.join(OUTPUT_FOLDER, filename)
            file.save(upload_path)

            print(f"✅ Uploaded file saved at: {upload_path}")

            # Run inference
            img = Image.open(upload_path).convert("RGB")
            img_tensor = preprocess_image(img)

            with torch.no_grad():
                prediction = model(img_tensor)[0]

            # Draw results
            image_cv = cv2.imread(upload_path)
            if image_cv is None:
                print(f"❌ ERROR: OpenCV could not read {upload_path}")
            else:
                result_image = draw_predictions(image_cv, prediction)
                cv2.imwrite(result_path, result_image)
                print(f"✅ Result image saved at: {result_path}")

            return render_template('index.html', 
                                   uploaded_image=url_for('uploaded_file', filename=filename), 
                                   result_image=url_for('output_file', filename=filename))

    return render_template('index.html')

# Serve uploaded images correctly
@app.route('/uploads/<filename>')
def uploaded_file(filename):
    return send_from_directory(os.path.abspath(UPLOAD_FOLDER), filename)

# Serve output images correctly
@app.route('/outputs/<filename>')
def output_file(filename):
    return send_from_directory(os.path.abspath(OUTPUT_FOLDER), filename)


# Real-time webcam detection
def generate_frames():
    cap = cv2.VideoCapture(0)
    while cap.isOpened():
        success, frame = cap.read()
        if not success:
            break

        # Convert frame and preprocess
        img = Image.fromarray(cv2.cvtColor(frame, cv2.COLOR_BGR2RGB))
        img_tensor = preprocess_image(img)

        # Run inference
        with torch.no_grad():
            prediction = model(img_tensor)[0]

        # Draw results
        frame = draw_predictions(frame, prediction)

        # Convert to JPEG
        _, buffer = cv2.imencode('.jpg', frame)
        frame_bytes = buffer.tobytes()

        yield (b'--frame\r\n'
               b'Content-Type: image/jpeg\r\n\r\n' + frame_bytes + b'\r\n')

    cap.release()

@app.route('/video_feed')
def video_feed():
    return Response(generate_frames(), mimetype='multipart/x-mixed-replace; boundary=frame')

if __name__ == '__main__':
    app.run(debug=True)