from matplotlib.pyplot import *
from math import *
from time import *

while True:  # Loop for restarting the script
    # The x-axis minimum is set to -3, change this value to change the window
    x_axis_min = -3

    # Automatically set the x&y axes based on the x-min 
    y_axis_min = x_axis_min
    x_axis_max = abs(x_axis_min)
    y_axis_max = abs(x_axis_min)

    # Create a list of x and corresponding y lattice point coordinates using axes max/min
    x = list(range(x_axis_min, x_axis_max + 1, 1))  # range(start, stop, step)
    y = list(range(x_axis_min, x_axis_max + 1, 1))  # range(start, stop, step)

    print("Input your dy/dx")
    print("-Explicitly define every")
    print("operation")
    print("Ex: 2x/y input as 2*x/y")

    # Ask user for the dydx function
    dydx_expression = input("dy/dx= ")

    def dydx(_x, _y):
        try:
            return eval(dydx_expression, {"x": _x, "y": _y, "__builtins__": {}})
        except Exception as e:
            print("Error in dydx function:", e)
            return 0

    # Create list of x, y pairs
    x_y = []
    xm = []
    ym = []

    # Create nested list of lists of all ordered pairs then extract xm and ym lists
    for i in range(len(x)):
        for j in range(len(y)):
            x_y.append((x[i], y[j]))

    for i in range(len(x_y)):
        xm.append(x_y[i][0])  # Select all x midpoint values

    for i in range(len(x_y)):
        ym.append(x_y[i][1])  # Select all y midpoint values   

    # Create x and y axes, extend by 1 in each direction so the slope lines are visible
    axis((x_axis_min-1, x_axis_max+1, y_axis_min-1, y_axis_max+1))   # axis(Xmin, Xmax, Ymin, Ymax)
    axis("on")                  # matplotlib: turn axes "on"
    grid(False)                 # matplotlib: turn grid "off"

    # matplotlib: draw arrows(x, y, dx, dy) where (x, y) is start and (x + dx, y + dy) is end
    d = 0.250  # Initial value of half the length of each line segment
    xs = []
    ys = []
    m = []
    dx = []
    dy = []

    for i in range(len(xm)):
        m.append(dydx(xm[i], ym[i]))  # Call dydx to calculate m for each (xm, ym) ordered pair
    for i in range(len(xm)):
        dx.append((d) / (sqrt(1 + pow(m[i], 2))))  # Calculate dx for each xs
    for i in range(len(xm)):
        dy.append(m[i] * dx[i])      # Calculate dy for each xs
    for i in range(len(xm)):
        xs.append(xm[i] - dx[i])     # xs = Starting x values of arrows
        ys.append(ym[i] - dy[i])     # ys = Starting y values of arrows

    # Draw arrows starting at each (xs, ys) ordered pair and moving by 2dx and 2dy
    for i in range(len(xs)):
        arrow(xs[i], ys[i], 2*dx[i], 2*dy[i], color="red", head_width=0.05)

    show()  # Ensure the graph is displayed and halts execution
    sleep(30)
    
    # Pause until the graph window is closed, then continue
    input("Press Enter to close the graph and proceed...")

    # Ask the user if they want to restart or exit
    restart = input("Do you want to re-execute and plot another? (yes/no): ").strip().lower()
    if restart != 'yes':  # Exit loop if the user doesn't input "yes"
        print("Goodbye!")
        break