import os
os.environ['TF_CPP_MIN_LOG_LEVEL'] = '2'  # Optional: Suppress TensorFlow info/warnings

import pandas as pd
import numpy as np
import matplotlib.pyplot as plt
from tensorflow.keras.models import Sequential
from tensorflow.keras.layers import Dense
from tensorflow.keras.optimizers import Adam

# Load CSV data
df = pd.read_csv('data_2d_classification.csv')
if df.empty:
    raise RuntimeError("Data file is empty or failed to load.")

X = df[['x1', 'x2']].values.astype(np.float32)
y = df['label'].values.astype(np.float32)

# Build the neural network model explicitly
model = Sequential()
model.add(Dense(10, input_dim=2, activation='relu'))
model.add(Dense(1, activation='sigmoid'))

model.compile(optimizer=Adam(learning_rate=0.01), loss='binary_crossentropy')

# Train the model with progress output every 500 epochs
model.fit(X, y, epochs=3000, verbose=1, batch_size=32)

# Prepare grid for boundary visualization
h = 0.02
x_min, x_max = X[:, 0].min() - 1, X[:, 0].max() + 1
y_min, y_max = X[:, 1].min() - 1, X[:, 1].max() + 1
xx, yy = np.meshgrid(np.arange(x_min, x_max, h), np.arange(y_min, y_max, h))
grid = np.c_[xx.ravel(), yy.ravel()]
Z = model.predict(grid, verbose=0).reshape(xx.shape)

# Plot results
plt.figure(figsize=(8, 6))
plt.contourf(xx, yy, Z, alpha=0.3, levels=np.linspace(0, 1, 11), cmap=plt.cm.RdBu)
plt.scatter(X[y == 0][:, 0], X[y == 0][:, 1], color='blue', label='Group 0', edgecolor='k')
plt.scatter(X[y == 1][:, 0], X[y == 1][:, 1], color='red', label='Group 1', edgecolor='k')
plt.contour(xx, yy, Z, levels=[0.5], linewidths=2, colors='k')
plt.legend()
plt.title('Decision Boundary from Neural Network Binary Classification (Keras)')
plt.xlabel('x1')
plt.ylabel('x2')
plt.tight_layout()
plt.show()  # Ensure the plot appears