import pandas as pd
import tensorflow as tf
from sklearn.preprocessing import StandardScaler, OneHotEncoder

data_train = pd.read_csv('dane/water_train.csv')

X_train = data_train.drop('is_safe', axis=1)
y_train = data_train['is_safe']

categorical_cols = ['bacteria', 'viruses']
encoder = OneHotEncoder(sparse=False, handle_unknown='ignore')
X_train_encoded = pd.DataFrame(encoder.fit_transform(X_train[categorical_cols]))
X_train_processed = pd.concat([X_train.drop(categorical_cols, axis=1), X_train_encoded], axis=1)
X_train_processed.columns = X_train_processed.columns.astype(str)

scaler = StandardScaler()
X_train_scaled = scaler.fit_transform(X_train_processed)

model = tf.keras.Sequential([
    tf.keras.layers.Dense(128, activation='relu'),
    tf.keras.layers.Dense(256, activation='relu'),
    tf.keras.layers.Dense(256, activation='relu'),
    tf.keras.layers.Dense(1, activation='sigmoid')
])

model.compile(
    loss=tf.keras.losses.binary_crossentropy,
    optimizer=tf.keras.optimizers.Adam(lr=0.03),
    metrics=[
        tf.keras.metrics.BinaryAccuracy(name='accuracy'),
        tf.keras.metrics.Precision(name='precision'),
        tf.keras.metrics.Recall(name='recall')
    ]
)

model.fit(X_train_scaled, y_train, batch_size=32, epochs=5, verbose=2)
model.save("savedmodel")