import os
import numpy as np
import tensorflow as tf
from flask import Flask, request, render_template, jsonify
from tensorflow.keras.utils import load_img, img_to_array
from werkzeug.utils import secure_filename
from datetime import datetime

app = Flask(__name__)

# Load the trained model
MODEL_PATH = r"model.keras"  # Update to correct path
model = tf.keras.models.load_model(MODEL_PATH)

# Configurations
UPLOAD_FOLDER = os.path.join('static', 'uploads')
ALLOWED_EXTENSIONS = {'jpg', 'jpeg', 'png'}
app.config['UPLOAD_FOLDER'] = UPLOAD_FOLDER
os.makedirs(UPLOAD_FOLDER, exist_ok=True)

def allowed_file(filename):
    return '.' in filename and filename.rsplit('.', 1)[1].lower() in ALLOWED_EXTENSIONS

def preprocess_image(image_path):
    img = load_img(image_path, target_size=(224, 224))
    img_array = img_to_array(img) / 255.0
    return np.expand_dims(img_array, axis=0)

@app.route('/')
def index():
    return render_template('home.html')

@app.route('/tool')
def tool():
    return render_template('tool.html')

@app.route('/about')
def about():
    return render_template('about.html')

@app.route('/contact') 
def contact():
    return render_template('contact.html')

@app.route('/predict', methods=['POST'])
def predict():
    if 'file' not in request.files:
        return jsonify({'error': 'No files uploaded'}), 400

    files = request.files.getlist('file')
    if not files or all(f.filename == '' for f in files):
        return jsonify({'error': 'No files selected'}), 400

    results = []
    for file in files:
        if file and allowed_file(file.filename):
            filename = secure_filename(file.filename)
            timestamp = datetime.now().strftime("%Y%m%d%H%M%S%f")
            unique_filename = f"{timestamp}_{filename}"
            file_path = os.path.join(app.config['UPLOAD_FOLDER'], unique_filename)
            file.save(file_path)

            try:
                img_array = preprocess_image(file_path)
                prediction = model.predict(img_array)[0][0]
                label = "Dirty" if prediction > 0.5 else "Clean"
                confidence = prediction if label == "Dirty" else 1 - prediction

                results.append({
                    'label': label,
                    'confidence': f"{confidence:.2%}",
                    'image_url': f"/static/uploads/{unique_filename}"
                })
            except Exception as e:
                results.append({
                    'label': 'Error',
                    'confidence': 'N/A',
                    'image_url': None,
                    'error': str(e)
                })
        else:
            results.append({
                'label': 'Error',
                'confidence': 'N/A',
                'image_url': None,
                'error': f"Invalid file type: {file.filename}"
            })

    # Render a results page and pass results into it
    return render_template('results.html', results=results)

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