import os
import shutil
import cv2
import numpy as np
from PIL import Image
from retinaface import RetinaFace
from transformers import pipeline

# Initialize the pipeline for image classification
pipe = pipeline("image-classification", model="rizvandwiki/gender-classification-2")

def is_face_blurry(face_image, threshold_factor=0.7, edge_threshold=100):
    if not ENABLE_BLURRY_DETECTION:
        return 0, False, 0, 0

    if face_image is None or face_image.size == 0:
        raise ValueError("Invalid or empty face image provided.")

    # Convert to grayscale
    gray = cv2.cvtColor(face_image, cv2.COLOR_BGR2GRAY)

    # Compute the Laplacian of the image and then the variance
    laplacian = cv2.Laplacian(gray, cv2.CV_64F)
    variance = laplacian.var()

    # Calculate the mean pixel value and set a dynamic threshold based on it
    mean_val = np.mean(gray)
    dynamic_threshold = mean_val * threshold_factor

    # Edge intensity check
    edge_intensity = np.mean(np.abs(laplacian))
    is_edge_strong = edge_intensity > edge_threshold

    # Check if the image is blurry or not
    is_blurry = variance < dynamic_threshold

    return variance, is_blurry and not is_edge_strong, dynamic_threshold, edge_intensity

def is_face_blurry_old(face_image, threshold_factor=1.0):
    if face_image is None or face_image.size == 0:
        raise ValueError("Invalid or empty face image provided.")

    # Convert to grayscale
    gray = cv2.cvtColor(face_image, cv2.COLOR_BGR2GRAY)
    #gray = cv2.GaussianBlur(gray, (3, 3), 0)

    # Compute the Laplacian of the image and then the variance
    variance = cv2.Laplacian(gray, cv2.CV_64F).var()

    mean_val = np.mean(gray)
    dynamic_threshold = mean_val * threshold_factor

    return variance, variance < dynamic_threshold, dynamic_threshold

def extract_face(image_path, facial_area):
    img = cv2.imread(image_path)
    x1, y1, x2, y2 = facial_area
    return img[y1:y2, x1:x2]

def classify_gender(face_image):
    predictions = pipe(face_image)
    return predictions[0]['label'], predictions[0]['score']

def save_debug_image(original_image_path, facial_area, blur_score, dynamic_threshold, edge_intensity):
    image = cv2.imread(original_image_path)
    x1, y1, x2, y2 = facial_area
    cv2.rectangle(image, (x1, y1), (x2, y2), (0, 255, 0), 2)
    cv2.putText(image, f"Blur Score: {blur_score:.2f}", (x1, y1 - 50), cv2.FONT_HERSHEY_SIMPLEX, 0.9, (0, 255, 0), 2)
    cv2.putText(image, f"Dynamic Threshold: {dynamic_threshold:.2f}", (x1, y1 - 30), cv2.FONT_HERSHEY_SIMPLEX, 0.9, (0, 255, 0), 2)
    cv2.putText(image, f"Edge Intensity: {edge_intensity:.2f}", (x1, y1 - 10), cv2.FONT_HERSHEY_SIMPLEX, 0.9, (0, 255, 0), 2)
    debug_image_path = os.path.join(DEBUG_FOLDER, os.path.basename(original_image_path))
    cv2.imwrite(debug_image_path, image)

def is_image_black_and_white(image_path, saturation_threshold=10):
    image = cv2.imread(image_path)
    image_hsv = cv2.cvtColor(image, cv2.COLOR_BGR2HSV)
    _, saturation, _ = cv2.split(image_hsv)
    std_dev_saturation = np.std(saturation)
    return std_dev_saturation < saturation_threshold

def process_images(source_folder, no_face_folder, multi_face_folder, blurry_folder, low_color_folder, male_folder, female_folder, unknown_folder, log_file_path):
    os.makedirs(no_face_folder, exist_ok=True)
    os.makedirs(multi_face_folder, exist_ok=True)
    os.makedirs(blurry_folder, exist_ok=True)
    os.makedirs(low_color_folder, exist_ok=True)
    os.makedirs(male_folder, exist_ok=True)
    os.makedirs(female_folder, exist_ok=True)
    os.makedirs(unknown_folder, exist_ok=True)
    if DEBUG_IMAGES:
        os.makedirs(DEBUG_FOLDER, exist_ok=True)

    total_files = sum([len(files) for r, d, files in os.walk(source_folder)])
    processed_files = 0

    with open(log_file_path, 'a') as log_file:
        for root, dirs, files in os.walk(source_folder):
            for file in files:
                if file.lower().endswith(('.png', '.jpg', '.jpeg', '.gif', '.bmp')):
                    file_path = os.path.join(root, file)
                    processed_files += 1
                    try:
                        if is_image_black_and_white(file_path):
                            shutil.move(file_path, os.path.join(low_color_folder, file))
                            log_file.write(f"Moved {file} to {low_color_folder}: Low colored image detected\n")
                            continue

                        faces = RetinaFace.detect_faces(file_path)
                        face_count = 0 if not faces else len(faces)

                        if face_count == 0:
                            shutil.move(file_path, os.path.join(no_face_folder, file))
                            log_file.write(f"Moved {file} to {no_face_folder}: No face detected\n")
                        elif face_count > 1:
                            shutil.move(file_path, os.path.join(multi_face_folder, file))
                            log_file.write(f"Moved {file} to {multi_face_folder}: Multiple faces detected\n")
                        else:
                            main_face = list(faces.values())[0]
                            face_img = extract_face(file_path, main_face['facial_area'])
                            blur_score, main_face_blurry, dynamic_threshold, edge_intensity = is_face_blurry(face_img)

                            if DEBUG_IMAGES:
                                save_debug_image(file_path, main_face['facial_area'], blur_score, dynamic_threshold, edge_intensity)

                            log_file.write(f"Edge Intensity for {file}: {edge_intensity:.2f}\n")  # Log edge intensity

                            if main_face_blurry:
                                shutil.move(file_path, os.path.join(blurry_folder, file))
                                log_file.write(f"Moved {file} to {blurry_folder}: Blurry face detected\n")
                            else:
                                face_pil = Image.fromarray(face_img)
                                predicted_gender, confidence = classify_gender(face_pil)

                                if predicted_gender == 'male':
                                    shutil.move(file_path, os.path.join(male_folder, file))
                                    log_file.write(f"Moved {file} to {male_folder}: Male with confidence {confidence:.2f}\n")
                                elif predicted_gender == 'female':
                                    shutil.move(file_path, os.path.join(female_folder, file))
                                    log_file.write(f"Moved {file} to {female_folder}: Female with confidence {confidence:.2f}\n")
                                else:
                                    shutil.move(file_path, os.path.join(unknown_folder, file))
                                    log_file.write(f"Moved {file} to {unknown_folder}: Gender not recognized\n")
                    except Exception as e:
                        shutil.move(file_path, os.path.join(unknown_folder, file))
                        log_file.write(f"Error processing {file}, moved to {unknown_folder}: {e}\n")

                    print(f"Processed {processed_files}/{total_files} files.")
    os.makedirs(no_face_folder, exist_ok=True)
    os.makedirs(multi_face_folder, exist_ok=True)
    os.makedirs(blurry_folder, exist_ok=True)
    os.makedirs(low_color_folder, exist_ok=True)
    os.makedirs(male_folder, exist_ok=True)
    os.makedirs(female_folder, exist_ok=True)
    os.makedirs(unknown_folder, exist_ok=True)
    if DEBUG_IMAGES:
        os.makedirs(DEBUG_FOLDER, exist_ok=True)

    total_files = sum([len(files) for r, d, files in os.walk(source_folder)])
    processed_files = 0

    with open(log_file_path, 'a') as log_file:
        for root, dirs, files in os.walk(source_folder):
            for file in files:
                if file.lower().endswith(('.png', '.jpg', '.jpeg', '.gif', '.bmp')):
                    file_path = os.path.join(root, file)
                    processed_files += 1
                    try:
                        if is_image_black_and_white(file_path):
                            shutil.move(file_path, os.path.join(low_color_folder, file))
                            log_file.write(f"Moved {file} to {low_color_folder}: Low colored image detected\n")
                            continue

                        faces = RetinaFace.detect_faces(file_path)
                        face_count = 0 if not faces else len(faces)

                        if face_count == 0:
                            shutil.move(file_path, os.path.join(no_face_folder, file))
                            log_file.write(f"Moved {file} to {no_face_folder}: No face detected\n")
                        elif face_count > 1:
                            shutil.move(file_path, os.path.join(multi_face_folder, file))
                            log_file.write(f"Moved {file} to {multi_face_folder}: Multiple faces detected\n")
                        else:
                            main_face = list(faces.values())[0]
                            face_img = extract_face(file_path, main_face['facial_area'])
                            blur_score, main_face_blurry, dynamic_threshold = is_face_blurry(face_img)

                            if DEBUG_IMAGES:
                                save_debug_image(file_path, main_face['facial_area'], blur_score, dynamic_threshold)

                            if main_face_blurry:
                                shutil.move(file_path, os.path.join(blurry_folder, file))
                                log_file.write(f"Moved {file} to {blurry_folder}: Blurry face detected\n")
                            else:
                                face_pil = Image.fromarray(face_img)
                                predicted_gender, confidence = classify_gender(face_pil)

                                if predicted_gender == 'male':
                                    shutil.move(file_path, os.path.join(male_folder, file))
                                    log_file.write(f"Moved {file} to {male_folder}: Male with confidence {confidence:.2f}\n")
                                elif predicted_gender == 'female':
                                    shutil.move(file_path, os.path.join(female_folder, file))
                                    log_file.write(f"Moved {file} to {female_folder}: Female with confidence {confidence:.2f}\n")
                                else:
                                    shutil.move(file_path, os.path.join(unknown_folder, file))
                                    log_file.write(f"Moved {file} to {unknown_folder}: Gender not recognized\n")
                    except Exception as e:
                        shutil.move(file_path, os.path.join(unknown_folder, file))
                        log_file.write(f"Error processing {file}, moved to {unknown_folder}: {e}\n")

                    print(f"Processed {processed_files}/{total_files} files.")

# Set your folder paths here
source_folder = r'G:\gender_classifier\blur_test'
no_face_folder = r'G:\man dataset\no_face_detected'
multi_face_folder = r'G:\man dataset\multi_faces_detected'
blurry_folder = r'G:\man dataset\blurry_faces_detected'
low_color_folder = r'G:\man dataset\low_colored_images'
male_folder = r'G:\man dataset\male_new_images'
female_folder = r'G:\man dataset\female_new_images'
unknown_folder = r'G:\man dataset\unknown_images'
log_file_path = 'logs.txt'
DEBUG_IMAGES = True
ENABLE_BLURRY_DETECTION = False
DEBUG_FOLDER = 'debug_folder_path'  # Set the path for the debug folder

# Run the function
process_images(source_folder, no_face_folder, multi_face_folder, blurry_folder, low_color_folder, male_folder, female_folder, unknown_folder, log_file_path)