import os
import base64
from flask import Flask, render_template, request, jsonify
from flask_cors import CORS
import tensorflow as tf
from tensorflow.keras.preprocessing.image import load_img, img_to_array
import numpy as np
from skimage.feature import graycomatrix, graycoprops
from skimage.color import rgb2gray
from sklearn.preprocessing import StandardScaler
from scipy.stats import entropy

# Disable TensorFlow OneDNN optimization
os.environ['TF_ENABLE_ONEDNN_OPTS'] = '0'

# Set thresholds
CONFIDENCE_THRESHOLD = 0.6
num_classes = 6  
ENTROPY_THRESHOLD = np.log2(num_classes) * 0.2

# Model settings
MODEL_PATH = "models/glcm_classweight_InceptionV3_model.h5"
IMAGE_SIZE = (227, 227)

def extract_glcm_features(image):
    """Extracts GLCM features from an image."""
    img_gray = rgb2gray(image)
    img_gray = (img_gray * 255).astype(np.uint8)
    
    # Compute GLCM
    glcm = graycomatrix(img_gray, distances=[50], angles=[np.pi/2], levels=256)
    
    # Extract statistical properties
    features = np.array([
        graycoprops(glcm, 'contrast')[0, 0],
        graycoprops(glcm, 'dissimilarity')[0, 0],
        graycoprops(glcm, 'homogeneity')[0, 0],
        graycoprops(glcm, 'energy')[0, 0],
        graycoprops(glcm, 'correlation')[0, 0]
    ]).reshape(1, -1)
    
    # Normalize GLCM features
    scaler = StandardScaler()
    features = scaler.fit_transform(features)
    return features

def create_app():
    """Creates Flask app instance."""
    app = Flask(__name__)
    CORS(app)
    return app

app = create_app()

@app.route('/', methods=['GET', 'POST'])
def home():
    """Renders the home page."""
    return render_template('home.html')

@app.route('/classify', methods=['POST'])
def classify_images():
    """Handles multiple image classification requests."""
    images = request.files.getlist('images')

    results = []

    try:
        # Load the optimized model
        model = tf.keras.models.load_model(MODEL_PATH)

        for image in images:
            image_path = os.path.join("uploads", image.filename)
            os.makedirs("uploads", exist_ok=True)
            image.save(image_path)

            # Read image and convert to base64 for display
            with open(image_path, "rb") as img_file:
                image_base64 = base64.b64encode(img_file.read()).decode('utf-8')

            # Preprocess image
            img = load_img(image_path, target_size=IMAGE_SIZE)
            img_array = img_to_array(img) / 255.0
            img_array = np.expand_dims(img_array, axis=0)

            # Extract GLCM features
            glcm_features = extract_glcm_features(img_array[0])  
            predictions = model.predict([img_array, glcm_features])

            probs = predictions[0]
            max_prob = float(np.max(probs))
            predicted_class = int(np.argmax(probs))

            # Compute entropy of the probability distribution
            ent = float(entropy(probs))

            # Classify as "Unknown" if entropy is high or max_prob is too low
            if ent > ENTROPY_THRESHOLD or max_prob < CONFIDENCE_THRESHOLD:
                predicted_class = "Unknown 🫥"
                confidence = 0.0
            else:
                confidence = max_prob * 100  # Convert to percentage

            results.append({
                "image_name": image.filename,
                "predicted_label": predicted_class,
                "confidence": round(confidence, 2),
                "image_data": image_base64
            })
            
            print(f"ENTROPY_THRESHOLD: {ENTROPY_THRESHOLD}")
            print(f"Entropy: {ent}")
            print(f"ent > ENTROPY_THRESHOLD: {ent > ENTROPY_THRESHOLD}")
            print(f"Low confidence: {max_prob is not None and max_prob < CONFIDENCE_THRESHOLD}")

        return jsonify(results)

    except Exception as e:
        print(f"Error during classification: {str(e)}")
        return jsonify({'error': str(e)}), 500


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