import streamlit as st
import numpy as np
import matplotlib.pyplot as plt
from sklearn.datasets import make_moons, make_circles, make_blobs
import torch
import torch.nn as nn
import torch.optim as optim

st.set_page_config(page_title="NN Playground", layout="wide")

st.title("🧠 Neural Network Playground (Hugging Face Style)")

# --- Sidebar: Dataset & Hyperparameters ---
st.sidebar.header("Data & Features")
dataset_choice = st.sidebar.selectbox("Dataset", ["Blobs", "Moons", "Circles", "Spiral"])
train_ratio = st.sidebar.slider("Training data ratio", 0.1, 0.9, 0.5)
noise = st.sidebar.slider("Noise", 0.0, 1.0, 0.1)
batch_size = st.sidebar.slider("Batch size", 1, 50, 10)
learning_rate = st.sidebar.slider("Learning rate", 0.001, 0.1, 0.03)
activation_choice = st.sidebar.selectbox("Activation", ["Tanh", "ReLU", "Sigmoid"])
regularization = st.sidebar.selectbox("Regularization", ["None", "L2"])
reg_rate = st.sidebar.slider("Regularization rate", 0.0, 0.1, 0.0)

# --- Generate dataset ---
if dataset_choice == "Blobs":
    X, y = make_blobs(n_samples=200, centers=2, cluster_std=noise)
elif dataset_choice == "Moons":
    X, y = make_moons(n_samples=200, noise=noise)
elif dataset_choice == "Circles":
    X, y = make_circles(n_samples=200, noise=noise, factor=0.5)
else:  # Spiral (custom)
    n_points = 100
    theta = np.sqrt(np.random.rand(n_points)) * 4 * np.pi
    r_a = 2*theta + np.pi
    data_a = np.array([np.cos(theta)*theta, np.sin(theta)*theta]).T
    data_b = np.array([np.cos(theta + np.pi)*theta, np.sin(theta + np.pi)*theta]).T
    X = np.vstack([data_a, data_b])
    y = np.array([0]*n_points + [1]*n_points)
    X += noise * np.random.randn(*X.shape)

# Convert to torch tensors
X = torch.tensor(X, dtype=torch.float32)
y = torch.tensor(y, dtype=torch.long)

# --- Sidebar: Network configuration ---
st.sidebar.header("Network")
n_layers = st.sidebar.slider("Hidden Layers", 1, 4, 2)
neurons = []
for i in range(n_layers):
    neurons.append(st.sidebar.slider(f"Neurons in layer {i+1}", 1, 10, 4))

# --- Build model dynamically ---
layers = []
input_dim = X.shape[1]
for n in neurons:
    layers.append(nn.Linear(input_dim, n))
    if activation_choice == "Tanh":
        layers.append(nn.Tanh())
    elif activation_choice == "ReLU":
        layers.append(nn.ReLU())
    else:
        layers.append(nn.Sigmoid())
    input_dim = n
layers.append(nn.Linear(input_dim, 2))  # output layer
model = nn.Sequential(*layers)

# --- Optimizer & Loss ---
if regularization == "L2":
    optimizer = optim.SGD(model.parameters(), lr=learning_rate, weight_decay=reg_rate)
else:
    optimizer = optim.SGD(model.parameters(), lr=learning_rate)
criterion = nn.CrossEntropyLoss()

# --- Train the model ---
epochs = 200
for epoch in range(epochs):
    optimizer.zero_grad()
    outputs = model(X)
    loss = criterion(outputs, y)
    loss.backward()
    optimizer.step()

# --- Decision boundary ---
xx, yy = np.meshgrid(np.linspace(X[:,0].min()-1, X[:,0].max()+1, 200),
                     np.linspace(X[:,1].min()-1, X[:,1].max()+1, 200))
grid = torch.tensor(np.c_[xx.ravel(), yy.ravel()], dtype=torch.float32)
with torch.no_grad():
    preds = model(grid).argmax(dim=1).numpy().reshape(xx.shape)

# --- Plot ---
plt.figure(figsize=(8,6))
plt.contourf(xx, yy, preds, alpha=0.3, cmap="bwr")
plt.scatter(X[:,0], X[:,1], c=y, edgecolor='k', cmap="bwr")
plt.title("Neural Network Decision Boundary")
st.pyplot(plt)