import tensorflow as tf
from tensorflow.keras.layers import Input, Conv2D, MaxPooling2D, Flatten, Dense
from tensorflow.keras.models import Model, Sequential
import os

def create_functional_cnn():
    """Creates a two-layer CNN model using the Functional API for Grad-CAM compatibility."""
    
    # Input layer expects 224x224 3-channel (RGB) images
    inputs = Input(shape=(224, 224, 3), name='input_layer')
    
    # First Convolutional Layer
    x = Conv2D(32, (3, 3), activation='relu', name='conv1')(inputs)
    x = MaxPooling2D((2, 2), name='pool1')(x)
    
    # Second, Final Convolutional Layer (This layer's name is used by Grad-CAM)
    x = Conv2D(64, (3, 3), activation='relu', name='conv2')(x)
    x = MaxPooling2D((2, 2), name='pool2')(x)
    
    x = Flatten(name='flatten')(x)
    
    # Output layer with 1 neuron for binary prediction
    outputs = Dense(1, activation='sigmoid', name='output')(x) 
    
    model = Model(inputs, outputs)
    
    # Compile the model
    model.compile(optimizer='adam', loss='binary_crossentropy')
    return model

# 1. Create the model
cnn_model = create_functional_cnn()

# 2. Save the model in the required format
model_filename = 'cnn_model.h5'
cnn_model.save(model_filename)

print(f"\n✅ Functional CNN Model saved successfully as '{model_filename}'")
print("This model is now fully compatible with the Grad-CAM implementation.")