import streamlit as st
import numpy as np
import matplotlib.pyplot as plt
import sympy as sp

# --- Streamlit App ---
st.title("🔽 Gradient Descent Visualizer")

# Function input
st.subheader("Define Function")
func_str = st.text_input("Enter a function in terms of x:", "x**2 + x")

# Sympy setup
x = sp.symbols("x")
try:
    func = sp.sympify(func_str)
    derivative = sp.diff(func, x)
except Exception as e:
    st.error(f"Invalid function: {e}")
    st.stop()

# Parameters
st.subheader("Parameters")
start_point = st.number_input("Starting Point", value=5.0)
learning_rate = st.number_input("Learning Rate", value=0.25, step=0.01)
iterations = st.slider("Number of Iterations", 1, 50, 16)

# Convert sympy to numpy function
f_np = sp.lambdify(x, func, "numpy")
fprime_np = sp.lambdify(x, derivative, "numpy")

# Gradient descent iterations
points = [start_point]
for i in range(iterations):
    grad = fprime_np(points[-1])
    new_point = points[-1] - learning_rate * grad
    points.append(new_point)

# Final point
current_point = points[-1]

# Plot function and descent path
st.subheader(f"Iteration {iterations}")
x_vals = np.linspace(-6, 6, 400)
y_vals = f_np(x_vals)

fig, ax = plt.subplots()
ax.plot(x_vals, y_vals, label=str(func))
ax.axhline(0, color="brown", linewidth=1)
ax.axvline(0, color="gray", linewidth=1)

# Plot descent points
y_points = f_np(np.array(points))
ax.plot(points, y_points, "ro-")

ax.set_xlabel("x - axis")
ax.set_ylabel("y - axis")
ax.legend()
st.pyplot(fig)

# Show current point
st.success(f"📍 Current Point: {current_point}")