import streamlit as st
import numpy as np
import matplotlib.pyplot as plt
from sklearn.datasets import make_moons, make_circles, make_blobs
from sklearn.model_selection import train_test_split
from sklearn.preprocessing import StandardScaler
import tensorflow as tf
from tensorflow.keras import models, layers, optimizers

# Function to generate datasets
def get_dataset(name, n_samples=1000, noise=0.2):
    if name == "Moons":
        return make_moons(n_samples=n_samples, noise=noise)
    elif name == "Circles":
        return make_circles(n_samples=n_samples, noise=noise)
    else:
        return make_blobs(n_samples=n_samples, centers=2, cluster_std=1.5)

# Function to build the model
def build_model(activation, lr, num_layers, num_neurons):
    model = models.Sequential()
    model.add(layers.Input(shape=(2,)))
    for _ in range(num_layers):
        model.add(layers.Dense(num_neurons, activation=activation))
    model.add(layers.Dense(1, activation='sigmoid'))

    opt = optimizers.Adam(learning_rate=lr)
    model.compile(optimizer=opt, loss='binary_crossentropy', metrics=['accuracy'])
    return model

# Visualize decision boundary
def visualize_decision_boundary(model, X, y):
    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.linspace(x_min, x_max, 300),
                         np.linspace(y_min, y_max, 300))
    grid = np.c_[xx.ravel(), yy.ravel()]
    preds = model.predict(grid).reshape(xx.shape)
    fig, ax = plt.subplots()
    ax.contourf(xx, yy, preds, cmap='RdBu', alpha=0.6)
    ax.scatter(X[:, 0], X[:, 1], c=y, cmap='RdBu', edgecolor='white')
    ax.set_title("Decision Boundary")
    return fig

# Plot loss
def plot_loss(history):
    fig, ax = plt.subplots()
    ax.plot(history.history['loss'], label='Train Loss')
    ax.plot(history.history['val_loss'], label='Test Loss')
    ax.set_title("Training vs Testing Loss")
    ax.set_xlabel("Epoch")
    ax.set_ylabel("Loss")
    ax.legend()
    return fig

# Streamlit App
st.title("🧠 TensorFlow Playground (Simplified)")
st.write("Explore how hyperparameters affect neural network decision boundaries.")

# Sidebar Inputs
st.sidebar.header("Hyperparameters")

dataset = st.sidebar.radio("Select Dataset", ["Moons", "Circles", "Blobs"])
noise = st.sidebar.slider("Dataset Noise", 0.0, 1.0, value=0.2, step=0.01)
activation = st.sidebar.selectbox("Activation Function", ["relu", "sigmoid", "tanh"])
lr = st.sidebar.slider("Learning Rate", 0.001, 0.1, value=0.01, step=0.001)
split = st.sidebar.slider("Train-Test Split Ratio", 0.1, 0.9, value=0.7)
batch = st.sidebar.select_slider("Batch Size", options=list(range(8, 129, 8)), value=32)
epochs = st.sidebar.slider("Epochs", 10, 200, value=100, step=10)
num_neurons = st.sidebar.slider("Neurons per Hidden Layer", 1, 100, value=10)
num_layers = st.sidebar.slider("Number of Hidden Layers", 1, 5, value=2)

if st.button("Train Model"):
    with st.spinner("Training model..."):
        X, y = get_dataset(dataset, noise=noise)
        X = StandardScaler().fit_transform(X)
        X_train, X_test, y_train, y_test = train_test_split(X, y, train_size=split)

        model = build_model(activation, lr, num_layers, num_neurons)
        history = model.fit(X_train, y_train, validation_data=(X_test, y_test),
                            batch_size=batch, epochs=epochs, verbose=0)

        st.subheader("📊 Decision Boundary")
        fig1 = visualize_decision_boundary(model, X, y)
        st.pyplot(fig1)

        st.subheader("📈 Training vs Testing Loss")
        fig2 = plot_loss(history)
        st.pyplot(fig2)