import Plots as plt
using NLsolve

function odeEE(f, y0, t)
    ns = length(t)-1
    y = zeros(ns+1)
    y[1] = y0
    
    for k in 1:ns
        h = t[k+1]-t[k]
        y[k+1] = y[k] + h*f(y[k],t[k])
    end
    
    return y
end

#---------------------------

function odeIE(f, y0, t) 
    ns = length(t) - 1
    y = zeros(ns+1)
    y[1] = y0
    
    function F!(Fval, ynext)
        Fval[1] = ynext[1] - y[kk] - Δt*f(ynext[1], t[kk+1])
    end
    
    kk=undef
    Δt=undef
    for k in 1:ns
        kk = k
        Δt = t[k+1]-t[k]
        results = nlsolve(F!, [y[k]])
        y[k+1] = results.zero[1]
    end
    
    return y
end;

f(y,t) = (α=1; -α*y)

#---------------------------

y0 = 1
tend = 10
Δt = 0.1
t = LinRange(0.0, tend+Δt, trunc(Int, tend/Δt+1))

yex = y0*exp.(-t)
yee = odeEE(f, y0, t)
yie = odeIE(f, y0, t)

#---------------------------

plt.resetfontsizes(); plt.scalefontsizes(1.5)
plt.plot( t,yex, lw=1, label="exact", color="black")
plt.plot!(t,yee, lw=1, label="EE")
plt.plot!(t,yie, lw=1, label="IE")
plt.plot!(xlabel="t", ylabel="y")
plt.plot!(foreground_color_legend=nothing)

t = LinRange(0,5,10000)
y = -2*exp.(-1000*t) + t .+ 2;

plt.resetfontsizes()
plots=[]
plt.plot(t,y)
plt.plot!(xlim=[0,5])
p=plt.plot!(ylabel="y", xlabel="t", legend=nothing)
plt.plot!(title="full x range")
push!(plots,p)
plt.plot(t,y)
plt.plot!(xlim=[0,0.005], ylim=[0,2.5])
plt.plot!(title="Zoom in")
p=plt.plot!(ylabel="y", xlabel="t", legend=nothing)
push!(plots,p)
plt.plot(plots..., layout=(2,1))

f(y,t) = (α=100; -α*(y-t-2.0) + 1.0)

y0   = 0            # try with 0 and 2
tend = 4.5

tEX  = LinRange(0,tend,10000)
yEX = -2*exp.(-100*tEX) + tEX .+ 2;

fac = 0.9                # try values: 0.1, 0.5, 0.9, 1.1
dt  = 2.0/100*fac        # 2/alpha is the stability limit, * a factor to adjust up or down
t   = LinRange(0,tend,trunc(Int,tend/dt+1))

yEE = odeEE(f, y0, t)

#----------------------

println("n steps = $(length(t)-1)")

plt.plot(tEX,yEX); 
p1 = plt.plot!(t,yEE, legend=nothing, title="full x range")

plt.plot(tEX,yEX); 
plt.plot!(t,yEE)
p2 = plt.plot!(xlim=[0, 0.05], legend=nothing, title="zoomed in ")

plt.plot([p1,p2]..., layout=(2,1))

fac = 10                 # try values: 0.1, 0.5, 0.9, 1.1, 10
dt = 2.0/100*fac        
t   = LinRange(0,tend,trunc(Int,tend/dt+1))

yIE = odeIE(f, y0, t)

#----------------------

println("n steps = $(length(t)-1)")

plt.plot(tEX,yEX)
p1 = plt.plot!(t,yIE, legend=nothing, title="full x range")

plt.plot(tEX,yEX)
plt.plot!(t,yIE)
p2 = plt.plot!(xlim=[0, 0.05], legend=nothing, title="zoomed in ")

plt.plot([p1,p2]..., layout=(2,1))