import os
import numpy as np
import cv2
from sklearn.model_selection import train_test_split
from tensorflow.keras.models import Sequential
from tensorflow.keras.layers import Conv2D, MaxPooling2D, Flatten, Dense
from tensorflow.keras.callbacks import ModelCheckpoint

# Function to load images from specified folders
def load_images_from_folders(folders):
    images = []
    for folder in folders:
        for filename in os.listdir(folder):
            if filename.lower().endswith(('.png', '.jpg', '.jpeg')):  # Check for valid image extensions
                img = cv2.imread(os.path.join(folder, filename), cv2.IMREAD_GRAYSCALE)  # Read image as grayscale
                if img is not None:
                    img = cv2.resize(img, (224, 224))  # Resize to 224x224 pixels
                    img = img.astype(np.float32)  # Ensure the image is in float32 format
                    img /= 255.0  # Normalize to [0, 1]
                    images.append(img)
                else:
                    print(f"Failed to load image: {filename}")
    return np.array(images)

# Load normal and pneumonia images
normal_folders = [
    os.path.join('chest-xray-pneumonia', 'chest_xray', 'test', 'NORMAL'),
    os.path.join('chest-xray-pneumonia', 'chest_xray', 'train', 'NORMAL'),
    os.path.join('chest-xray-pneumonia', 'chest_xray', 'val', 'NORMAL'),
]
pneumonia_folders = [
    os.path.join('chest-xray-pneumonia', 'chest_xray', 'test', 'PNEUMONIA'),
    os.path.join('chest-xray-pneumonia', 'chest_xray', 'train', 'PNEUMONIA'),
    os.path.join('chest-xray-pneumonia', 'chest_xray', 'val', 'PNEUMONIA'),
]

normal_images = load_images_from_folders(normal_folders)
pneumonia_images = load_images_from_folders(pneumonia_folders)

# Print the number of images loaded for each category
print(f"Loaded {len(normal_images)} normal images.")
print(f"Loaded {len(pneumonia_images)} pneumonia images.")

# Prepare dataset for machine learning
images = np.concatenate([normal_images, pneumonia_images])  # Combine normal and pneumonia images
labels = np.array([0] * len(normal_images) + [1] * len(pneumonia_images))  # Create labels

# Split the dataset into training and testing sets
X_train, X_test, y_train, y_test = train_test_split(images, labels, test_size=0.2, random_state=42)

# Convert to RGB format
X_train_rgb = np.stack([cv2.cvtColor((img * 255).astype(np.uint8), cv2.COLOR_GRAY2RGB) for img in X_train])
X_test_rgb = np.stack([cv2.cvtColor((img * 255).astype(np.uint8), cv2.COLOR_GRAY2RGB) for img in X_test])

print(f'Training data shape: {X_train_rgb.shape}, Training labels shape: {y_train.shape}')
print(f'Testing data shape: {X_test_rgb.shape}, Testing labels shape: {y_test.shape}')

# Build a simple CNN model
model = Sequential([
    Conv2D(32, (3, 3), activation='relu', input_shape=(224, 224, 3)),
    MaxPooling2D(pool_size=(2, 2)),
    Conv2D(64, (3, 3), activation='relu'),
    MaxPooling2D(pool_size=(2, 2)),
    Flatten(),
    Dense(128, activation='relu'),
    Dense(1, activation='sigmoid')  # Binary classification
])

model.compile(optimizer='adam', loss='binary_crossentropy', metrics=['accuracy'])

# Set up model checkpointing with the new file extension
checkpoint = ModelCheckpoint('pneumonia_model.keras', save_best_only=True, monitor='val_loss', mode='min')

# Train the model
model.fit(X_train_rgb, y_train, validation_data=(X_test_rgb, y_test), epochs=10, callbacks=[checkpoint])

# Save the model in the new format
model.save('pneumonia_model_final.keras')
print("Model saved as 'pneumonia_model_final.keras'")