using DifferentialEquations
using Plots
gr()
function nil!(dw,w,tp,t)
    dw[1] = 0
    dw[2] = 0
    dw[3] = 0
    dw[4] = 0
    dw[5] = 0
    dw[6] = 0
end

function finite_depth_u_crossing_nondim(x, y, z, t, p)
    (St, theta, W, K) = p
    k1 = [cos(theta/2),  sin(theta/2)]
    k2 = [cos(theta/2), -sin(theta/2)]
    k1x = k1[1]*x + k1[2]*y
    k2x = k2[1]*x + k2[2]*y
    coshfac = cosh(z + K) / sinh(K)
    sinhfac = sinh(z + K) / sinh(K)
    u_x = -coshfac * (sin(k1x - t)*k1[1] + sin(k2x - t)*k2[1])
    u_y = -coshfac * (sin(k1x - t)*k1[2] + sin(k2x - t)*k2[2])
    u_z = sinhfac * (cos(k1x - t) + cos(k2x - t))
    return [u_x, u_y, u_z]
end

function drag_fd_3d_crossing_nondim!(dw, w, p, t)
    (St, theta, W, K) = p
    u = finite_depth_u_crossing_nondim(w[1], w[2], w[3], t, p)
    dw[1] = w[4]
    dw[2] = w[5]
    dw[3] = w[6]
    dw[4] = 1/St * (u[1] - w[4])
    dw[5] = 1/St * (u[2] - w[5])
    dw[6] = 1/St * (u[3] - w[6] - W)
end


St=0.1
theta = 0
tspan=(0, 10000)
u0=[0,0,-2,0,0,0]
W=0.0072
K=10
p = (St, theta, W, K)

condition_bottom(u,t,integrator) = u[3] + K
global cad = ContinuousCallback(condition_bottom, terminate!)

prob1 = SDEProblem(drag_fd_3d_crossing_nondim!, nil!, u0, tspan, p)
prob2 = SDEProblem(drag_fd_3d_crossing_nondim!, nil!, u0, tspan, p)
abstol = 1e-9
sol1 = solve(prob1, callback=cad, abstol = abstol, maxiters = 1e7, alg_hints = [:stiff])
sol2 = solve(prob2, callback=cad, abstol = abstol, maxiters = 1e7)
plot(sol1[1, :], sol1[2, :], sol1[3, :],  label="Path 1")
plot!(sol2[1, :], sol2[2, :], sol2[3, :], label="Path 2")
d1 = (sqrt(last(sol1.u)[1]^2 + last(sol1.u)[2]^2))
d2 = (sqrt(last(sol2.u)[1]^2 + last(sol2.u)[2]^2))
println("Advection distance of sol1: $d1")
println("Advection distance of sol2: $d2")
println("Time steps for solution 1: $(length(sol1.t))")
println("Time steps for solution 2: $(length(sol2.t))")
maxlength = max(length(sol1.t), length(sol2.t))
println("Error bound: abstol * length = $abstol * $maxlength = $(abstol * maxlength)")
println("Actual error: $(abs(d1-d2))")
println("Algorithm used for sol1 (hints = [:stiff]): $(sol1.alg)")
println("Algorithm used for sol2 (no hints): $(sol2.alg)")

Advection distance of sol1: 4.48163592520678
Advection distance of sol2: 4.301763934710236
Time steps for solution 1: 874828
Time steps for solution 2: 2970126
Error bound: abstol * length = 1.0e-9 * 2970126 = 0.002970126
Actual error: 0.17987199049654468
Algorithm used for sol1 (hints = [:stiff]): ImplicitRKMil{0, false, Nothing, typeof(OrdinaryDiffEqCore.DEFAULT_PRECS), Val{:central}, true, nothing, NLNewton{Rational{Int64}, Rational{Int64}, Rational{Int64}, Rational{Int64}}, Float64, :Predictive, SciMLBase.AlgorithmInterpretation.Ito}(nothing, NLNewton{Rational{Int64}, Rational{Int64}, Rational{Int64}, Rational{Int64}}(1//100, 10, 1//5, 1//5, false, true, 0//1), OrdinaryDiffEqCore.DEFAULT_PRECS, 1.0, :constant, 0.001, false)
Algorithm used for sol2 (no hints): SOSRI()