import numpy as np
import matplotlib.pyplot as plt

# Constants
k = 8.99e9  # Coulomb's constant in N*m^2/C^2

# Define point charges and their positions
charges = [(1e-9, (0, 0)), (-1e-9, (1, 0))]  # (charge, (x, y))

# Create a grid of points
x = np.linspace(-2, 2, 400)
y = np.linspace(-2, 2, 400)
X, Y = np.meshgrid(x, y)

# Initialize the electric field components and potential
Ex, Ey = np.zeros_like(X), np.zeros_like(Y)
V = np.zeros_like(X)

# Calculate the electric field components and potential at each point in the grid
for charge, (cx, cy) in charges:
    Rx = X - cx
    Ry = Y - cy
    R = np.sqrt(Rx**2 + Ry**2)
    Ex += k * charge * Rx / R**3
    Ey += k * charge * Ry / R**3
    V += k * charge / R


# Avoid division by zero by setting very small values to a finite number
V[np.isinf(V)] = np.nan
V = np.nan_to_num(V)

# Plot the electric field vectors
fig, ax = plt.subplots()
ax.streamplot(X, Y, Ex, Ey, color='b', linewidth=1, density=1.5, arrowstyle='->')

# Plot the equipotential lines
contour_levels = np.linspace(-20, 20, 50)
contour = ax.contour(X, Y, V, levels=contour_levels, alpha=0.75)
fig.colorbar(contour, ax=ax, label='Potencial (V)')

# Plot the point charges
for charge, (cx, cy) in charges:
    ax.plot(cx, cy, 'ro' if charge > 0 else 'bo', markersize=10)

# Set plot limits and labels
ax.set_xlim(-2, 2)
ax.set_ylim(-2, 2)
ax.set_xlabel('x')
ax.set_ylabel('y')
ax.set_title('Siločiary a ekvipotenciály elektrického poľa bodových nábojov')
plt.show()