import numpy as np
import tensorflow as tf
import matplotlib.pyplot as plt

# Definicija sistema nelinearnih jednačina
def system_equations(x, y):
    eq1 = x**2 + y**2 - 25
    eq2 = x*y - 9
    return eq1, eq2

# Generisanje podataka za obuku
np.random.seed(42)
num_samples = 1000
x_train = np.random.uniform(low=-5, high=5, size=(num_samples, 2))
y_train = np.array([system_equations(x, y) for x, y in x_train])

# Definisanje neuronske mreže
model = tf.keras.Sequential([
    tf.keras.layers.Dense(16, activation='relu', input_shape=(2,)),
    tf.keras.layers.Dense(16, activation='relu'),
    tf.keras.layers.Dense(2)  # Izlaz sadrži dva neurona, po jedan za svaku jednačinu
])

# Kompilacija modela
model.compile(optimizer='adam', loss='mean_squared_error')

# Obuka modela
model.fit(x_train, y_train, epochs=50, batch_size=32, verbose=0)

# Generisanje test podataka za vizualizaciju
x_vals = np.linspace(-5, 5, 100)
y_vals = np.linspace(-5, 5, 100)
X, Y = np.meshgrid(x_vals, y_vals)
test_data = np.column_stack((X.ravel(), Y.ravel()))

# Predviđanje pomoću obučenog modela
predictions = model.predict(test_data)

# Vizualizacija rezultata
Z1, Z2 = predictions[:, 0].reshape(X.shape), predictions[:, 1].reshape(X.shape)

plt.contour(X, Y, Z1, levels=[0], colors='r', label='Equation 1 (approximation)')
plt.contour(X, Y, Z2, levels=[0], colors='b', label='Equation 2 (approximation)')
plt.xlabel('x')
plt.ylabel('y')
plt.legend()
plt.title('Solving systems of nonlinear equations using a neural network')
plt.show()