import itertools
import numpy as np
import matplotlib.pyplot as plt


def euler(f, t_start, x_start, t_stop, delta_t):
    ts = [t_start]
    xs = [x_start]
    while ts[-1] < t_stop:
        t = ts[-1]
        x = xs[-1]
        ts.append(t + delta_t)
        xs.append(x + f(x, t) * delta_t)
    return ts, xs


def apply(f, t):
    return f(t)


def fun4(g, k):
    return lambda t: (-g*np.e**(-k*t) + g) / k


def main():
    ############ 3 ############
    t = np.linspace(0, 3)
    plt.subplot(1, 3, 1)
    plt.plot(t, apply(lambda t: t, t))
    plt.plot(*euler(lambda x, t: 1, 0, 0, 3, 1))
    plt.plot(*euler(lambda x, t: 1, 0, 0, 3, 0.1))
    plt.plot(*euler(lambda x, t: 1, 0, 0, 3, 0.01))
    plt.legend(["Analytical", "Δt = 1", "Δt = 0.1", "Δt = 0.01"])
    plt.gca().set_title("dx/dt = 1")

    plt.subplot(1, 3, 2)
    plt.plot(t, apply(lambda t: t**2 - 4, t))
    plt.plot(*euler(lambda x, t: 2*t, 0, -4, 3, 1))
    plt.plot(*euler(lambda x, t: 2*t, 0, -4, 3, 0.1))
    plt.plot(*euler(lambda x, t: 2*t, 0, -4, 3, 0.01))
    plt.legend(["Analytical", "Δt = 1", "Δt = 0.1", "Δt = 0.01"])
    plt.gca().set_title("dx/dt = 2t")

    plt.subplot(1, 3, 3)
    plt.plot(t, apply(lambda t: 4 * np.e**-t, t))
    plt.plot(*euler(lambda x, t: -x, 0, 4, 3, 1))
    plt.plot(*euler(lambda x, t: -x, 0, 4, 3, 0.1))
    plt.plot(*euler(lambda x, t: -x, 0, 4, 3, 0.01))
    plt.legend(["Analytical", "Δt = 1", "Δt = 0.1", "Δt = 0.01"])
    plt.gca().set_title("dx/dt = -x")

    plt.xlabel("t")
    plt.ylabel("x(t)")
    plt.savefig("3.png")
    plt.clf()


    ############ 4e ############
    t = np.linspace(0, 5)

    plt.plot(t, apply(fun4(2, 3), t))
    plt.plot(t, apply(fun4(1, 1.5), t))
    plt.plot(t, apply(fun4(2, 2), t))
    plt.plot(t, apply(fun4(1, 1), t))

    plt.legend(["g = 2, k = 3", "g = 1, k = 1.5", "g = 2, k = 2", "g = 1, k = 1"])

    plt.xlabel("t")
    plt.ylabel("x(t)")
    plt.savefig("4e.png")
    plt.clf()

    ############ 4f ############
    t = np.linspace(0, 5)

    plt.subplot(2, 2, 1)
    plt.plot(t, apply(fun4(2, 3), t))
    plt.plot(*euler(lambda x, t: 2 - 3*x, 0, 0, 5, 0.1/3*2))
    plt.legend(["Analytical", "Euler (Δt = 0.1/2k)"])
    plt.gca().set_title("g = 2, k = 3")

    plt.subplot(2, 2, 2)
    plt.plot(t, apply(fun4(1, 1.5), t))
    plt.plot(*euler(lambda x, t: 1 - 1.5*x, 0, 0, 5, 0.1/1.5*2))
    plt.legend(["Analytical", "Euler (Δt = 0.1/2k)"])
    plt.gca().set_title("g = 1, k = 1.5")

    plt.subplot(2, 2, 3)
    plt.plot(t, apply(fun4(2, 2), t))
    plt.plot(*euler(lambda x, t: 2 - 2*x, 0, 0, 5, 0.1/2*2))
    plt.legend(["Analytical", "Euler (Δt = 0.1/2k)"])
    plt.gca().set_title("g = 2, k = 2")

    plt.subplot(2, 2, 4)
    plt.plot(t, apply(fun4(1, 1), t))
    plt.plot(*euler(lambda x, t: 1 - 1*x, 0, 0, 5, 0.1/2))
    plt.legend(["Analytical", "Euler (Δt = 0.1/2k)"])
    plt.gca().set_title("g = 1, k = 1")

    plt.xlabel("t")
    plt.ylabel("x(t)")
    plt.tight_layout()
    plt.savefig("4f.png")
    plt.clf()


    ############ 4h ############
    t = np.linspace(0, 20)
    gk = list(itertools.product(np.arange(1, 5, 1), np.arange(1, 5, 1)))
    for g, k in gk:
        x = fun4(g, k)(t)
        plt.plot(x, apply(lambda x: g - k*x, x))

    plt.xlabel("x(t)")
    plt.ylabel("dx/dt")
    plt.legend([f"g = {g}, k = {k}" for g, k in gk])
    plt.title("dx/dt vs. x(t) for t in range [0, 20]")
    plt.tight_layout()

    plt.savefig("4h.png")
    plt.clf()


    ############ 5d ############
    rk = [(1, 2), (1, 1), (2, 1)]
    for r, k in rk:
        plt.plot(*euler(lambda x, t: r*x - k*x, 0, 4, 2, 0.001))

    plt.legend([f"r = {r}, k = {k}" for r, k in rk])
    plt.tight_layout()

    plt.savefig("5d.png")
    plt.clf()


    ############ 5h ############
    for x0 in [0.5, 1, 1.1]:
        plt.plot(*euler(lambda x, t: x*x - x, 0, x0, 2, 0.001))

    plt.xlabel("t")
    plt.ylabel("x(t)")
    plt.legend(["x0 = 0.5", "x0 = 1", "x0 = 1.1"])
    plt.savefig("5h.png")
    plt.clf()


    ############ 5i ############
    xx = [(10, 50), (10, 20), (50, 30), (50, 100)]
    for x0, xmax in xx:
        plt.plot(*euler(lambda x, t: x*(1-x/xmax), 0, x0, 5, 0.001))

    plt.xlabel("t")
    plt.ylabel("x(t)")
    plt.legend([f"x0 = {x0}, xmax = {xmax}" for x0, xmax in xx])
    plt.tight_layout()
    plt.savefig("5i.png")
    plt.clf()


    ############ 5i ############
    xmax = 30
    x = np.linspace(-1, xmax+1)
    plt.plot(x, apply(lambda x: x*(1-x/xmax), x))

    plt.xlabel("x")
    plt.ylabel("dx/dt")
    plt.tight_layout()
    plt.savefig("5l.png")
    plt.clf()


    ############ 50 ############
    plt.plot(*euler(lambda x, t: x*(1-x/20) - 2*x, 0, 5, 5, 1))

    plt.xlabel("t")
    plt.ylabel("x(t)")
    plt.tight_layout()
    plt.savefig("5o.png")
    plt.clf()

if __name__ == "__main__":
    main()
