from tensorflow.keras.models import Sequential
from tensorflow.keras.layers import Conv2D, MaxPooling2D, Flatten, Dense
from tensorflow.keras.preprocessing.image import ImageDataGenerator
import os

def train_model():
    data_path = 'data/brain_tumor_dataset'
    img_size = (150, 150)

    datagen = ImageDataGenerator(rescale=1./255, validation_split=0.2)

    train_data = datagen.flow_from_directory(
        data_path,
        target_size=img_size,
        batch_size=32,
        class_mode='binary',
        subset='training'
    )

    val_data = datagen.flow_from_directory(
        data_path,
        target_size=img_size,
        batch_size=32,
        class_mode='binary',
        subset='validation'
    )

    model = Sequential([
        Conv2D(32, (3, 3), activation='relu', input_shape=(150, 150, 3)),
        MaxPooling2D(2, 2),
        Flatten(),
        Dense(64, activation='relu'),
        Dense(1, activation='sigmoid')  # Binary classifier: tumor vs no tumor
    ])

    model.compile(optimizer='adam', loss='binary_crossentropy', metrics=['accuracy'])
    model.fit(train_data, epochs=10, validation_data=val_data)

    os.makedirs("model", exist_ok=True)
    model.save("model/tumor_classifier.h5")
    print("✅ Model trained and saved to model/tumor_classifier.h5")

if __name__ == "__main__":
    train_model()