import numpy as np
from PIL import Image
from tensorflow.keras.models import Sequential
from tensorflow.keras.layers import Dense, BatchNormalization, Conv2D, MaxPooling2D, Flatten, Activation, Dropout

from sklearn.preprocessing import LabelEncoder
from tensorflow.keras.utils import to_categorical
from typing import List

def getTargets(filepaths: List[str]) -> List[str]:
    """
    Extract class labels from file paths.
    Assumes the file path structure includes the class name as the parent folder.
    Example: /path/to/data/VeryMildDemented/image1.jpg -> VeryMildDemented
    """
    labels = [fp.split('/')[-2] for fp in filepaths]  # Get the parent folder name as the label
    return labels

def encodeLabels(y_train: List, y_test: List):
    """
    Encode class labels into one-hot vectors.
    """
    label_encoder = LabelEncoder()
    y_train_labels = label_encoder.fit_transform(y_train)
    y_test_labels = label_encoder.transform(y_test)

    y_train_1h = to_categorical(y_train_labels)
    y_test_1h = to_categorical(y_test_labels)

    LABELS = label_encoder.classes_
    print(f"Classes: {LABELS} -- Encoded: {label_encoder.transform(LABELS)}")

    return LABELS, y_train_1h, y_test_1h

def getFeatures(filepaths: List[str]) -> np.array:
    """
    Load images from file paths and convert them to numpy arrays.
    Assumes images are already resized to 64x64 during preprocessing.
    """
    images = []
    for imagePath in filepaths:
        image = Image.open(imagePath).convert("RGB")  # Ensure 3-channel RGB
        image = np.array(image)
        images.append(image)
    return np.array(images)

def buildModel(inputShape: tuple, classes: int) -> Sequential:
    """
    Build a CNN model for Alzheimer's image classification.
    """
    model = Sequential()
    height, width, depth = inputShape
    inputShape = (height, width, depth)
    chanDim = -1

    # CONV => RELU => POOL layer set
    model.add(Conv2D(32, (3, 3), padding="same", name='conv_32_1', input_shape=inputShape))
    model.add(Activation("relu"))
    model.add(BatchNormalization(axis=chanDim))
    model.add(MaxPooling2D(pool_size=(2, 2)))
    model.add(Dropout(0.25))

    # (CONV => RELU) * 2 => POOL layer set
    model.add(Conv2D(64, (3, 3), padding="same", name='conv_64_1'))
    model.add(Activation("relu"))
    model.add(BatchNormalization(axis=chanDim))
    model.add(Conv2D(64, (3, 3), padding="same", name='conv_64_2'))
    model.add(Activation("relu"))
    model.add(BatchNormalization(axis=chanDim))
    model.add(MaxPooling2D(pool_size=(2, 2)))
    model.add(Dropout(0.25))

    # (CONV => RELU) * 3 => POOL layer set
    model.add(Conv2D(128, (3, 3), padding="same", name='conv_128_1'))
    model.add(Activation("relu"))
    model.add(BatchNormalization(axis=chanDim))
    model.add(Conv2D(128, (3, 3), padding="same", name='conv_128_2'))
    model.add(Activation("relu"))
    model.add(BatchNormalization(axis=chanDim))
    model.add(Conv2D(128, (3, 3), padding="same", name='conv_128_3'))
    model.add(Activation("relu"))
    model.add(BatchNormalization(axis=chanDim))
    model.add(MaxPooling2D(pool_size=(2, 2)))
    model.add(Dropout(0.25))

    # Fully connected layer
    model.add(Flatten())
    model.add(Dense(512, name='fc_1'))
    model.add(Activation("relu"))
    model.add(BatchNormalization())
    model.add(Dropout(0.5))

    # Output layer with softmax activation
    model.add(Dense(classes, name='output'))
    model.add(Activation("softmax"))

    return model