import numpy as np
import matplotlib.pyplot as plt
from scipy.integrate import odeint

def plot_direction_field(ax, f, x_range, y_range, nx=20, ny=20):
    x = np.linspace(x_range[0], x_range[1], nx)
    y = np.linspace(y_range[0], y_range[1], ny)
    X, Y = np.meshgrid(x, y)
    
    dx = np.ones_like(X)
    dy = f(X, Y)
    
    length = np.sqrt(dx**2 + dy**2)
    dx /= length
    dy /= length
    
    ax.quiver(X, Y, dx, dy, color='b', alpha=0.3)

def solve_ode(f, y0, x_range, num_points=1000):
    x = np.linspace(x_range[0], x_range[1], num_points)
    y = odeint(f, y0, x)
    return x, y.flatten()

# A) y dy/dx = -x; y(1) = 4
def f_a(y, x):
    return -x / y

# B) dy/dx = 1/y; y(0) = 1
def f_b(y, x):
    return 1 / y

# C) dy/dx = 0.2x^2 + y; y(0) = 1/2
def f_c(y, x):
    return 0.2 * x**2 + y

# D) y' = y - cos(π/2 * x); y(2) = 2
def f_d(y, x):
    return y - np.cos(np.pi/2 * x)

# Configuración de la figura
fig, axs = plt.subplots(2, 2, figsize=(15, 15))
axs = axs.flatten()

# A) y dy/dx = -x; y(1) = 4
plot_direction_field(axs[0], lambda x, y: -x/y, [-5, 5], [-5, 5])
x_a, y_a = solve_ode(f_a, 4, [1, 5])
axs[0].plot(x_a, y_a, 'r-', linewidth=2)
axs[0].set_title('A) y dy/dx = -x; y(1) = 4')

# B) dy/dx = 1/y; y(0) = 1
plot_direction_field(axs[1], lambda x, y: 1/y, [-5, 5], [0.1, 5])
x_b, y_b = solve_ode(f_b, 1, [0, 5])
axs[1].plot(x_b, y_b, 'r-', linewidth=2)
axs[1].set_title('B) dy/dx = 1/y; y(0) = 1')

# C) dy/dx = 0.2x^2 + y; y(0) = 1/2
plot_direction_field(axs[2], lambda x, y: 0.2*x**2 + y, [-5, 5], [-5, 5])
x_c, y_c = solve_ode(f_c, 0.5, [0, 5])
axs[2].plot(x_c, y_c, 'r-', linewidth=2)
axs[2].set_title('C) dy/dx = 0.2x^2 + y; y(0) = 1/2')

# D) y' = y - cos(π/2 * x); y(2) = 2
plot_direction_field(axs[3], lambda x, y: y - np.cos(np.pi/2 * x), [-5, 5], [-5, 5])
x_d, y_d = solve_ode(f_d, 2, [2, 7])
axs[3].plot(x_d, y_d, 'r-', linewidth=2)
axs[3].set_title('D) y\' = y - cos(π/2 * x); y(2) = 2')

for ax in axs:
    ax.grid(True)
    ax.axhline(y=0, color='k', linestyle='--')
    ax.axvline(x=0, color='k', linestyle='--')

plt.tight_layout()
plt.show()