import numpy as np
import matplotlib.pyplot as plt
import tensorflow as tf
from tensorflow.keras.models import Sequential
from tensorflow.keras.layers import Dense
import streamlit as st

# Function to generate synthetic data
def generate_data(dataset_type, noise, n_samples=500):
    np.random.seed(0)
    if dataset_type == 'moons':
        from sklearn.datasets import make_moons
        X, y = make_moons(n_samples=n_samples, noise=noise)
    elif dataset_type == 'circles':
        from sklearn.datasets import make_circles
        X, y = make_circles(n_samples=n_samples, noise=noise, factor=0.5)
    elif dataset_type == 'linear':
        X = np.random.randn(n_samples, 2)
        y = (X[:, 0] > X[:, 1]).astype(int)
    else:
        X = np.random.randn(n_samples, 2)
        y = np.random.randint(0, 2, n_samples)
    return X, y

# Function to create model
def create_model(input_shape, hidden_layers, activation, learning_rate, regularization_rate):
    model = Sequential()
    model.add(Dense(hidden_layers[0], input_shape=input_shape, activation=activation,
                    kernel_regularizer=tf.keras.regularizers.l2(regularization_rate)))
    for units in hidden_layers[1:]:
        model.add(Dense(units, activation=activation,
                        kernel_regularizer=tf.keras.regularizers.l2(regularization_rate)))
    model.add(Dense(1, activation='sigmoid'))
    model.compile(optimizer=tf.keras.optimizers.Adam(learning_rate),
                  loss='binary_crossentropy',
                  metrics=['accuracy'])
    return model

# Streamlit UI
st.title('Interactive Neural Network Visualization')
st.sidebar.header('Model Parameters')

# Dataset selection
dataset_type = st.sidebar.selectbox('Select dataset', ['moons', 'circles', 'linear'])
noise = st.sidebar.slider('Noise level', 0.0, 1.0, 0.2)
X, y = generate_data(dataset_type, noise)
split = st.sidebar.slider('Train/Test split ratio', 0.1, 0.9, 0.5)
split_idx = int(split * len(X))
X_train, X_test = X[:split_idx], X[split_idx:]
y_train, y_test = y[:split_idx], y[split_idx:]

# Model parameters
learning_rate = st.sidebar.slider('Learning rate', 0.001, 0.1, 0.01)
activation = st.sidebar.selectbox('Activation function', ['relu', 'tanh', 'sigmoid'])
regularization_rate = st.sidebar.slider('Regularization rate', 0.0, 0.1, 0.01)
hidden_layers = [st.sidebar.slider('Layer 1 units', 1, 10, 4),
                 st.sidebar.slider('Layer 2 units', 1, 10, 2)]

# Create and train model
model = create_model((2,), hidden_layers, activation, learning_rate, regularization_rate)
history = model.fit(X_train, y_train, epochs=100, verbose=0, validation_split=0.1)

# Evaluation
train_loss, train_acc = model.evaluate(X_train, y_train, verbose=0)
test_loss, test_acc = model.evaluate(X_test, y_test, verbose=0)
st.write(f'Training loss: {train_loss:.4f}, Training accuracy: {train_acc:.4f}')
st.write(f'Test loss: {test_loss:.4f}, Test accuracy: {test_acc:.4f}')

# Plot data and decision boundary
fig, ax = plt.subplots()
ax.scatter(X_test[:, 0], X_test[:, 1], c=y_test, cmap='viridis', marker='o', edgecolor='k', s=50)
xx, yy = np.meshgrid(np.linspace(X_test[:, 0].min(), X_test[:, 0].max(), 100),
                     np.linspace(X_test[:, 1].min(), X_test[:, 1].max(), 100))
Z = model.predict(np.c_[xx.ravel(), yy.ravel()])
Z = Z.reshape(xx.shape)
ax.contourf(xx, yy, Z, alpha=0.5, cmap='viridis')
ax.set_title('Data and Model Decision Boundary')
st.pyplot(fig)