text:
# check_hyperbolic_integrals.jl

# module check_hyperbolic_integral
# export hyperbolic_integral_one, hyperbolic_integral_two, hyperbolic_integral_three, hyperbolic_integral_four, hyperbolic_integralfive, hyperbolic_integral_six

using QuadGK
using Plots
using SpecialFunctions
plotly()
Plots.PlotlyBackend()

"""
This file tests a mathematical expansion of the integral:

    ∫_0^{β/2} (1 - cosh(Ω * x)) * cosh(x - β / 2) / (a^2 - β^2 / 4 + x * (β - x) - b * cosh(v * (x - β / 2)))^(3 / 2) dx

using some change of variables, two binomial expansions, a product of coshines expansion, and an identity that relates simpler resultant integrals to 1F2(a, {b1, b2}; z) hypergeometric functions.

Code is far from tidy or efficient. It is more a proof of implementation.
"""

"""
Original integral:

    ∫_0^{β/2} (1 - cosh(Ω * x)) * cosh(x - β / 2) / (a^2 - β^2 / 4 + x * (β - x) - b * cosh(v * (x - β / 2)))^(3 / 2) dx
"""
function hyperbolic_integral_one(Ω, β, α, v, w)

    # Initialise constants.
    R = (v^2 - w^2) / (w^2 * v)
    a = sqrt(β^2 / 4 + R * β * coth(β * v / 2))
    b = R * β / sinh(β * v / 2)

    coefficient = 2 * α * β^(3 / 2) * v^3 / (3 * sqrt(π) * sinh(β / 2) * w^3)

    integrand(x) = (1 - cosh(Ω * x)) * cosh(x - β / 2) / (a^2 - β^2 / 4 + x * (β - x) - b * cosh(v * (x - β / 2)))^(3 / 2)

    integral = quadgk(x -> integrand(x), 0, β / 2)

    return coefficient * integral[1]
end

"""
Transform integral to limits (0, 1):

    β/2 ∫_0^1 (1 - cosh(Ωβ(1-x)/2)) cosh(xβ/2) / (a^2 - β^2 x^2/4 - bcosh(vβx/2))^{3/2} dx
"""
function hyperbolic_integral_two(Ω, β, α, v, w)

    # Initialise constants.
    R = (v^2 - w^2) / (w^2 * v)
    a = sqrt(β^2 / 4 + R * β * coth(β * v / 2))
    b = R * β / sinh(β * v / 2)

    coefficient = 2 * α * β^(3 / 2) * v^3 / (3 * sqrt(π) * sinh(β / 2) * w^3)

    integrand(x) = (1 - cosh(Ω * β * (1 - x) / 2)) * cosh(x * β / 2) / (a^2 - β^2 * x^2 / 4 - b * cosh(v * β * x / 2))^(3 / 2)

    integral = quadgk(x -> integrand(x), 0, 1)

    return coefficient * β * integral[1] / 2
end

"""
First binomial expansion of denominator in integrand:

    ∑_{n=0}^∞ C(-3/2, n) (2/β)^{2n + 2} (-b)^n ∫_0^1 (1 - cosh(Ωβ(1-x)/2)) cosh(βx/2) cosh^n(βvx/2) / (4a^2/β^2 - x^2)^{n + 3/2} dx.

C(n, k) are binomial coefficients.
"""
function hyperbolic_integral_three(Ω, β, α, v, w)

    # Initialise constants.
    R = (v^2 - w^2) / (w^2 * v)
    a = sqrt(β^2 / 4 + R * β * coth(β * v / 2))
    b = R * β / sinh(β * v / 2)

    coefficient = 2 * α * β^(3 / 2) * v^3 / (3 * sqrt(π) * sinh(β / 2) * w^3)

    integrand(x, n) = (1 - cosh(Ω * β * (1 - x) / 2)) * cosh(x * β / 2) * cosh(β * v * x / 2)^n / (1 - x^2 + 4 * R * coth(β * v / 2) / β)^(n + 3 / 2)

    n_binomial_coeff(n) = -2 * √π / (gamma(-n - 1 / 2) * gamma(n + 1))

    n_coefficient(n) = n_binomial_coeff(n) * (2 / β)^(2 * n + 2) * (-b)^n

    total_sum = 0.0

    for n in 0:8 # Seems to converge very quickly.
        integral = quadgk(x -> integrand(x, n), 0, 1)[1]
        total_sum += n_coefficient(n) * integral
    end

    return coefficient * total_sum
end

"""
Second binomial expansion of the denominator in the integrand:

    ∑_{n=0}^∞ C(-3/2, n) (2/β)^{2n + 2} (-b)^n ∑_{m=0}^∞ C(-n-3/2, m) (-1)^m (β/2a)^{2n+2m+3} ∫_0^1 (1 - cosh(Ωβ(1-x)/2)) cosh(βx/2) cosh^n(βvx/2) x^{2m} dx
"""
function hyperbolic_integral_four(Ω, β, α, v, w)

    # Initialise constants.
    R = (v^2 - w^2) / (w^2 * v)
    a = sqrt(β^2 / 4 + R * β * coth(β * v / 2))
    b = R * β / sinh(β * v / 2)

    coefficient = 2 * α * β^(3 / 2) * v^3 / (3 * sqrt(π) * sinh(β / 2) * w^3)

    integrand(x, n, m) = (1 - cosh(Ω * β * (1 - x) / 2)) * cosh(x * β / 2) * cosh(β * v * x / 2)^n * (x^2)^m

    n_binomial_coeff(n) = -2 * √π / (gamma(-n - 1 / 2) * gamma(n + 1))

    m_binomial_coeff(n, m) = gamma(-n - 1 / 2) / (gamma(m + 1) * gamma(-m - n - 1 / 2))

    n_coefficient(n) = n_binomial_coeff(n) * (2 / β)^(2 * n + 2) * (-b)^n

    m_coefficient(n, m) = m_binomial_coeff(n, m) * (-1)^m * (β / (2 * a))^(2 * n + 2 * m + 3)

    total_sum_n = 0.0

    for n in 0:8

        total_sum_m = 0.0

        for m in 0:10
            integral = quadgk(x -> integrand(x, n, m), 0, 1)[1]
            total_sum_m += m_coefficient(n, m) * integral
        end
        total_sum_n += n_coefficient(n) * total_sum_m
    end

    return coefficient * total_sum_n
end

"""
Expand product of coshines in the integrand:

    ∑_{n=0}^∞ C(-3/2, n) (2/β)^{2n + 2} (-b/2)^n ∑_{m=0}^∞ C(-n-3/2, m) (-1)^m (β/2a)^{2n+2m+3} {

    C(n, n/2)(1 - nmod2)[∫_0^1 cosh(βx/2) x^{2m} dx - (1/2) ∑_{z_2} (cosh(Ωβ/2) ∫_0^1 cosh(βx(z_2)/2) x^{2m} dx - sinh(Ωβ/2) ∫_0^1 sinh(βx(z_2)/2) x^{2m} dx)]

    + ∑_{k=0}^{floor(n/2-1/2)} C(n, k) [∑_{z_3} ∫_0^1 cosh(βx(z_3)/2) x^{2m} dx - (1/2) ∑_{z_4} (cosh(Ωβ/2) ∫_0^1 cosh(βx(z_4)/2) x^{2m} dx - sinh(Ωβ/2) ∫_0^1 sinh(βx(z_4)/2) x^{2m} dx)]
    }

where:

    z_2 ∈ {Ω + 1, Ω - 1}
    z_3 ∈ {1 + v(n - 2k), 1 - v(n - 2k)}
    z_4 ∈ {Ω + 1 + v(n - 2k), Ω - 1 + v(n - 2k), Ω + 1 - v(n - 2k), Ω - 1 - v(n - 2k)}.
"""
function hyperbolic_integral_five(Ω, β, α, v, w)

    # Initialise constants.
    R = (v^2 - w^2) / (w^2 * v)
    a = sqrt(β^2 / 4 + R * β * coth(β * v / 2))
    b = R * β / sinh(β * v / 2)

    coefficient = 2 * α * β^(3 / 2) * v^3 / (3 * sqrt(π) * sinh(β / 2) * w^3)

    n_binomial_coeff_one(n) = -2 * √π / (gamma(-n - 1 / 2) * gamma(n + 1))

    m_binomial_coeff(n, m) = gamma(-n - 1 / 2) / (gamma(m + 1) * gamma(-m - n - 1 / 2))

    n_coefficient(n) = n_binomial_coeff_one(n) * (2 / β)^(2 * n + 2) * (-b)^n

    m_coefficient(n, m) = m_binomial_coeff(n, m) * (-1)^m * (β / (2 * a))^(2 * n + 2 * m + 3)

    ############################################################################

    cosh_integral(z, m) = quadgk(x -> cosh(β * x * z / 2) * (x^2)^m, 0, 1)[1]

    cosh_sinh_integral(z, m) = cosh(Ω * β / 2) * quadgk(x -> cosh(β * x * z / 2) * (x^2)^m, 0, 1)[1] - sinh(Ω * β / 2) * quadgk(x -> sinh(β * x * z / 2) * (x^2)^m, 0, 1)[1]

    k_binomial_coeff(n, k) = gamma(n + 1) / (gamma(k + 1) * gamma(n - k + 1))

    n_binomial_coeff_two(n) = 2^n * gamma((n + 1) / 2) / (√π * gamma(n / 2 + 1))

    z_1_args = 1
    z_2_args = [Ω + 1, Ω - 1]
    z_3_args(n, k) = [1 + v * (n - 2 * k), 1 - v * (n - 2 * k)]
    z_4_args(n, k) = [Ω + 1 + v * (n - 2 * k), Ω - 1 + v * (n - 2 * k), Ω + 1 - v * (n - 2 * k), Ω - 1 - v * (n - 2 * k)]

    ############################################################################

    total_sum = 0.0
    for n in 0:8

        n_coeff = n_coefficient(n)

        for k in 0:Int64(floor((n - 1) / 2))

            k_coeff = k_binomial_coeff(n, k)
            z_3, z_4 = z_3_args(n, k), z_4_args(n, k)

            for m in 0:10

                integral_3 = sum([cosh_integral(z, m) for z in z_3])
                integral_4 = sum([cosh_sinh_integral(z, m) for z in z_4])

                total_sum += n_coeff * m_coefficient(n, m) * k_coeff * (integral_3 - integral_4 / 2) / 2^n
            end
        end

        n_bin_coeff_2 = n_binomial_coeff_two(n)

        if mod(n, 2) == 0

            for m in 0:10

                integral_1 = cosh_integral(z_1_args, m)
                integral_2 = sum([cosh_sinh_integral(z, m) for z in z_2_args])

                total_sum += n_coeff * m_coefficient(n, m) * n_bin_coeff_2 * (integral_1 - integral_2 / 2) / 2^n
            end
        end
    end

    return coefficient * total_sum
end

"""
Substitute remaining integrals for equivalent summations and extract exponential term:

    e^{Ωβ/2}/2 ∑_{n=0}^∞ C(-3/2, n) (2/β)^{2n + 2} (-b/2)^n ∑_{m=0}^∞ C(-n-3/2, m) (-1)^m (β/2a)^{2n+2m+3} {

    C(n, n/2)(1 - nmod2)[2e^{-Ωβ/2} J(1, 2t, 0) - (1/2) ∑_{z_2} (J(z_2, 2t, 1) - J(z_2, 2t+1, -1))]

    + ∑_{k=0}^{floor(n/2-1/2)} C(n, k) [2e^{-Ωβ/2} ∑_{z_3} J(z_3, 2t, 0) - (1/2) ∑_{z_4} (J(z_4, 2t, 1) - J(z_4, 2t+1, -1))]
    }

where:

    z_2 ∈ {Ω + 1, Ω - 1}
    z_3 ∈ {1 + v(n - 2k), 1 - v(n - 2k)}
    z_4 ∈ {Ω + 1 + v(n - 2k), Ω - 1 + v(n - 2k), Ω + 1 - v(n - 2k), Ω - 1 - v(n - 2k)}

and:

    J(z, t, c) ≡ ∑_{t=0}^∞ (βz/2)^t (1-c⋅e^{-Ωβ}) / ((2m + t + 1)⋅(t!)).

The integral identities are:

    ∫_0^1 cosh(zx) x^{2m} dx = 1F2(m+1/2; {1/2, m+3/2}; z^2/4)/(2m+1)
                             = ∑_{t=0}^∞ z^{2t} / ((2t+2m+1)⋅(2t)!)

    ∫_0^1 sinh(zx) x^{2m} dx = z ⋅ 1F2(m+1; {3/2, m+2}; z^2/4)/(2m+2)
                             = ∑_{t=0}^∞ z^{2t+1} / ((2t+2m+2)⋅(2t+1)!)

where 1F2 is the hypergeometric function pFq with p = 1 and q = 2.
"""
function hyperbolic_integral_six(Ω, β, α, v, w; prec = 64)

    setprecision(BigFloat, prec)

    Ω = BigFloat("$Ω")
    β = BigFloat("$β")
    α = BigFloat("$α")
    v = BigFloat("$v")
    w = BigFloat("$w")

    # Initialise constants.
    R = (v^2 - w^2) / (w^2 * v)
    a = sqrt(β^2 / 4 + R * β * coth(β * v / 2))
    b = R * β / sinh(β * v / 2)

    coefficient = 2 * α * β^(3 / 2) * v^3 / (3 * sqrt(π) * sinh(β / 2) * w^3)

    n_binomial_coeff_one(n) = -2 * √π / (gamma(-n - 1 / 2) * gamma(n + 1))

    m_binomial_coeff(n, m) = gamma(-n - 1 / 2) / (gamma(m + 1) * gamma(-m - n - 1 / 2))

    n_coefficient(n) = n_binomial_coeff_one(n) * (2 / β)^(2 * n + 2) * (-b)^n

    m_coefficient(n, m) = m_binomial_coeff(n, m) * (-1)^m * (β / (2 * a))^(2 * n + 2 * m + 3)

    k_binomial_coeff(n, k) = gamma(n + 1) / (gamma(k + 1) * gamma(n - k + 1))

    n_binomial_coeff_two(n) = 2^n * gamma((n + 1) / 2) / (√π * gamma(n / 2 + 1))

    z_1_args = 1
    z_2_args = [Ω + 1, Ω - 1]
    z_3_args(n, k) = [1 + v * (n - 2 * k), 1 - v * (n - 2 * k)]
    z_4_args(n, k) = [Ω + 1 + v * (n - 2 * k), Ω - 1 + v * (n - 2 * k), Ω + 1 - v * (n - 2 * k), Ω - 1 - v * (n - 2 * k)]

    J(z, t, c, m) = (β * z / 2)^t * (1 + c * exp(-Ω * β)) / ((2 * m + t + 1) * gamma(t + 1))

    total_sum = 0.0
    for n in 0:3, m in 0:3, t in 0:100

        n_coeff = n_coefficient(n)
        m_coeff = m_coefficient(n, m)

        for k in 0:Int64(floor((n - 1) / 2))

            k_coeff = k_binomial_coeff(n, k)
            z_3, z_4 = z_3_args(n, k), z_4_args(n, k)

            integral_3 = sum([2 * J(z, 2 * t, 0, m) * exp(-Ω * β / 2) for z in z_3])
            integral_4 = sum([J(z, 2 * t, 1, m) - J(z, 2 * t + 1, -1, m) for z in z_4])
            total_sum += n_coeff * m_coeff * k_coeff * (integral_3 - integral_4 / 2) / 2^n
        end

        n_bin_coeff_2 = n_binomial_coeff_two(n)

        if mod(n, 2) == 0

            integral_1 = 2 * J(z_1_args, 2 * t, 0, m) * exp(-Ω * β / 2)
            integral_2 = sum([J(z, 2 * t, 1, m) - J(z, 2 * t + 1, -1, m) for z in z_2_args])
            total_sum += n_coeff * m_coeff * n_bin_coeff_2 * (integral_1 - integral_2 / 2) / 2^n
        end
    end

    return coefficient * total_sum * exp(Ω * β / 2) / 2
end

# end # end of module

"""
Plotting of the above six functions to compare them. They should all produce the same graph (which they do).
"""
Ω_range = 0.01:0.1:20
@time hyper_int_1 = [hyperbolic_integral_one(Ω, 10, 7, 5.8, 1.6) for Ω in Ω_range]
@time hyper_int_2 = [hyperbolic_integral_two(Ω, 10, 7, 5.8, 1.6) for Ω in Ω_range]
@time hyper_int_3 = [hyperbolic_integral_three(Ω, 10, 7, 5.8, 1.6) for Ω in Ω_range]
@time hyper_int_4 = [hyperbolic_integral_four(Ω, 10, 7, 5.8, 1.6) for Ω in Ω_range]
@time hyper_int_5 = [hyperbolic_integral_five(Ω, 10, 7, 5.8, 1.6) for Ω in Ω_range]
@time hyper_int_6 = [hyperbolic_integral_six(Ω, 10, 7, 5.8, 1.6) for Ω in Ω_range]

p = plot(Ω_range, abs.(hyper_int_1), yaxis=:log, label = "hyper_int_1")
plot!(Ω_range,  abs.(hyper_int_2), label = "hyper_int_2")
plot!(Ω_range,  abs.(hyper_int_3), label = "hyper_int_3")
plot!(Ω_range,  abs.(hyper_int_4), label = "hyper_int_4")
plot!(Ω_range,  abs.(hyper_int_5), label = "hyper_int_5")
plot!(Ω_range,  abs.(hyper_int_6), label = "hyper_int_6")
display(p)

"""
Mini tests of some individual parts of the above six functions. Such as the coshine product expansion, the 1F2 hypergeometric function, etc.
"""

# function inner_hyp_integral(x, n, β, v, w)
#
#     # Initialise constants.
#     R = (v^2 - w^2) / (w^2 * v)
#     a = sqrt(β^2 / 4 + R * β * coth(β * v / 2))
#     b = R * β / sinh(β * v / 2)
#
#     m_binomial_coeff(n, m) = gamma(-n - 1 / 2) / (gamma(m + 1) * gamma(-m - n - 1 / 2))
#
#     normal_term = (4 * a^2 / β^2 - x^2)^(-n - 3 / 2)
#
#     total_sum_m = 0.0
#     for m in 0:120
#         m_coefficient = m_binomial_coeff(n, m) * (-1)^m * (β / (2 * a))^(2 * n + 2 * m + 3)
#         total_sum_m += x^(2 * m) * m_coefficient
#     end
#
#     return normal_term, total_sum_m
# end
#
# # m = inner_hyp_integral(0.99, 0, 200, 5.8, 1.6)
# # @show(m)
#
# function cosh_expansion(x, n)
#
#     normal = cosh(x)^n
#
#     k_binomial_coeff(n, k) = gamma(n + 1) / (gamma(k + 1) * gamma(n - k + 1))
#
#     n_binomial_coeff_two(n) = 2^n * gamma((n + 1) / 2) / (√π * gamma(n / 2 + 1))
#
#     total_sum_k = 0.0
#     for k in 0:Int64(floor((n - 1) / 2))
#
#         total_sum_k += n_binomial_coeff_two(n) * (1 - mod(n, 2)) / 2^n + k_binomial_coeff(n, k) * cosh(x * (n - 2 * k)) / 2^(n - 1)
#     end
#
#     return normal, total_sum_k
# end
#
# # x_range = 0.01:0.01:1.0
# # cosh_exp = [cosh_expansion(x, 2) for x in x_range]
# # cosh_exp_norm = [x[1] for x in cosh_exp]
# # cosh_exp_sum = [x[2] for x in cosh_exp]
# # p = plot(x_range, cosh_exp_norm, yaxis = :log, label = "cosh_exp_norm")
# # plot!(x_range, cosh_exp_sum, label = "cosh_exp_sum")
# # display(p)
#
# function hyperbolic_expansion_small(β, v, n, m)
#
#     integrand(x, n, m) = cosh(x * β / 2) * cosh(β * v * x / 2)^n * (x^2)^m
#     integral = quadgk(x -> integrand(x, n, m), 0, 1)[1]
#
#     ##########################################################################
#
#     z_3(n, k) = [1 + v * (n - 2 * k), 1 - v * (n - 2 * k)]
#
#     cosh_integral(z, m) = quadgk(x -> cosh(β * x * z / 2) * (x^2)^m, 0, 1)[1]
#
#     k_binomial_coeff(n, k) = gamma(n + 1) / (gamma(k + 1) * gamma(n - k + 1))
#
#     n_binomial_coeff_two(n) = 2^n * gamma((n + 1) / 2) / (√π * gamma(n / 2 + 1))
#
#     total_sum_k = 0.0
#     for k in 0:Int64(floor((n - 1) / 2))
#
#         integral_3 = sum([cosh_integral(z, m) for z in z_3(n, k)])
#
#         total_sum_k += k_binomial_coeff(n, k) * integral_3
#     end
#
#     if mod(n, 2) == 0
#
#         integral_1 = cosh_integral(1, m)
#
#         total_sum_k += n_binomial_coeff_two(n) * integral_1
#     end
#
#     return integral, total_sum_k / 2^n
# end
#
# # β_range = 0.01:0.1:1
# # hyper_exp_s = [hyperbolic_expansion_small(β, 5.8, 1, 1) for β in β_range]
# # hyper_exp_s_int = [x[1] for x in hyper_exp_s]
# # hyper_exp_s_exp = [x[2] for x in hyper_exp_s]
# # p = plot(β_range, hyper_exp_s_int, yaxis = :log, label = "hyper_exp_s_int")
# # plot!(β_range, hyper_exp_s_exp, label = "hyper_exp_s_exp")
# # display(p)
#
# function hyperbolic_expansion_large(Ω, β, v, n, m)
#
#     integrand(x, n, m) = cosh(Ω * β * (1 - x) / 2) * cosh(x * β / 2) * cosh(β * v * x / 2)^n * (x^2)^m
#     integral = quadgk(x -> integrand(x, n, m), 0, 1)[1]
#
#     ##########################################################################
#
#     z_2 = [1 + Ω, -1 + Ω]
#     z_4(n, k) = [1 + Ω + v * (n - 2 * k), -1 + Ω - v * (n - 2 * k), 1 + Ω - v * (n - 2 * k), -1 + Ω + v * (n - 2 * k)]
#
#     cosh_sinh_integral(z, m) = cosh(Ω * β / 2) * quadgk(x -> cosh(β * x * z / 2) * (x^2)^m, 0, 1)[1] - sinh(Ω * β / 2) * quadgk(x -> sinh(β * x * z / 2) * (x^2)^m, 0, 1)[1]
#
#     k_binomial_coeff(n, k) = gamma(n + 1) / (gamma(k + 1) * gamma(n - k + 1))
#
#     n_binomial_coeff_two(n) = 2^n * gamma((n + 1) / 2) / (√π * gamma(n / 2 + 1))
#
#     total_sum_k = 0.0
#     for k in 0:Int64(floor((n - 1) / 2))
#
#         integral_4 = sum([cosh_sinh_integral(z, m) for z in z_4(n, k)])
#
#         total_sum_k += k_binomial_coeff(n, k) * integral_4
#     end
#
#     if mod(n, 2) == 0
#
#         integral_2 = sum([cosh_sinh_integral(z, m) for z in z_2])
#
#         total_sum_k += n_binomial_coeff_two(n) * integral_2
#     end
#
#     return integral, total_sum_k / 2^(n + 1)
# end
#
# # Ω_range = 0.01:0.1:1
# # hyper_exp_l = [hyperbolic_expansion_large(Ω, 2, 5.8, 1, 1) for Ω in Ω_range]
# # hyper_exp_l_int = [x[1] for x in hyper_exp_l]
# # hyper_exp_l_exp = [x[2] for x in hyper_exp_l]
# # p = plot(Ω_range, hyper_exp_l_int, label = "hyper_exp_l_int")
# # plot!(Ω_range, hyper_exp_l_exp, label = "hyper_exp_l_exp")
# # display(p)
#
# function final_cosh_expansion(z, β, m)
#
#     cosh_integral = quadgk(x -> cosh(β * x * z / 2) * (x^2)^m, 0, 1)[1]
#
#     ############################################################################
#
#     J(z, t, c, m) = (β * z / 2)^t / ((2 * m + t + 1) * gamma(t + 1))
#
#     total_sum = 0.0
#     for t in 0:100
#         total_sum += J(z, 2 * t, 0, m)
#     end
#
#     return cosh_integral, total_sum
# end
#
# # z_range = 0.01:0.1:10
# # f_cosh_exp = [final_cosh_expansion(z, 6, 0) for z in z_range]
# # f_cosh_exp_int = [x[1] for x in f_cosh_exp]
# # f_cosh_exp_exp = [x[2] for x in f_cosh_exp]
# # p = plot(z_range, f_cosh_exp_int, yaxis=:log, label = "f_int_exp_int")
# # plot!(z_range, f_cosh_exp_exp, label = "f_int_exp_exp")
# # display(p)
#
# function final_hyperbolic_expansion(Ω, z, β, m)
#
#     cosh_sinh_integral = cosh(Ω * β / 2) * quadgk(x -> cosh(β * x * z / 2) * (x^2)^m, 0, 1)[1] - sinh(Ω * β / 2) * quadgk(x -> sinh(β * x * z / 2) * (x^2)^m, 0, 1)[1]
#
#     ############################################################################
#
#     J(z, t, c, m) = (β * z / 2)^t * (1 + c * exp(-Ω * β)) / ((2 * m + t + 1) * gamma(t + 1))
#
#     total_sum = 0.0
#     for t in 0:100
#         total_sum += J(z, 2 * t, 1, m) - J(z, 2 * t + 1, -1, m)
#     end
#
#     return cosh_sinh_integral, total_sum * exp(Ω * β / 2) / 2
# end
#
# # Ω_range = 0.01:0.1:10
# # f_hyp_exp = [final_hyperbolic_expansion(Ω, 6, 6, 10) for Ω in Ω_range]
# # f_hyp_exp_int = [x[1] for x in f_hyp_exp]
# # f_hyp_exp_exp = [x[2] for x in f_hyp_exp]
# # p = plot(Ω_range, f_hyp_exp_int, yaxis=:log, label = "f_int_exp_int")
# # plot!(Ω_range, f_hyp_exp_exp, label = "f_int_exp_exp")
# # display(p)


text:
using simulationQP934
using Test

@testset "simulationQP934.jl" begin
    # Write your tests here.
    data = Data();
    box = initial_point(data);
    nc, fatm, natm = linkedlist(box,data);
    @test utotal(box, data, fatm, natm, nc) \approx utotal(box,data)

    upair,f = fpair(box[1],box[2],r,data)
    @test upair \approx  -0.2955814601750326

end

ext: jl


text:
abstract type AbstractController end

@inline function stepsize_controller!(integrator, alg)
  stepsize_controller!(integrator, integrator.opts.controller, alg)
end

# checks whether the controller should accept a step based on the error estimate
@inline function accept_step_controller(integrator, controller::AbstractController)
  return integrator.EEst <= 1
end

@inline function step_accept_controller!(integrator, alg, q)
  step_accept_controller!(integrator, integrator.opts.controller, alg, q)
end

@inline function step_reject_controller!(integrator, alg)
  step_reject_controller!(integrator, integrator.opts.controller, alg)
end

reset_alg_dependent_opts!(controller::AbstractController, alg1, alg2) = nothing

DiffEqBase.reinit!(integrator::ODEIntegrator, controller::AbstractController) = nothing

# Standard integral (I) step size controller
"""
    IController()

The standard (integral) controller is the most basic step size controller.
This controller is usually the first one introduced in numerical analysis classes
but should only be used rarely in practice because of efficiency problems for
many problems/algorithms.

Construct an integral (I) step size controller adapting the time step
based on the formula
```
Δtₙ₊₁ = εₙ₊₁^(1/k) * Δtₙ
```
where `k = get_current_adaptive_order(alg, integrator.cache) + 1` and `εᵢ` is the
inverse of the error estimate `integrator.EEst` scaled by the tolerance
(Hairer, Nørsett, Wanner, 2008, Section II.4).
The step size factor is multiplied by the safety factor `gamma` and clipped to
the interval `[qmin, qmax]`.
A step will be accepted whenever the estimated error `integrator.EEst` is
less than or equal to unity. Otherwise, the step is rejected and re-tried with
the predicted step size.

## References
- Hairer, Nørsett, Wanner (2008)
  Solving Ordinary Differential Equations I Nonstiff Problems
  [DOI: 10.1007/978-3-540-78862-1](https://doi.org/10.1007/978-3-540-78862-1)
"""
struct IController <: AbstractController
end

@inline function stepsize_controller!(integrator, controller::IController, alg)
  @unpack qmin, qmax, gamma = integrator.opts
  EEst = DiffEqBase.value(integrator.EEst)

  if iszero(EEst)
    q = inv(qmax)
  else
    expo = 1 / (get_current_adaptive_order(alg, integrator.cache) + 1)
    qtmp = DiffEqBase.fastpow(EEst, expo) / gamma
    @fastmath q = DiffEqBase.value(max(inv(qmax), min(inv(qmin), qtmp)))
    # TODO: Shouldn't this be in `step_accept_controller!` as for the PI controller?
    integrator.qold = DiffEqBase.value(integrator.dt) / q
  end
  q
end

function step_accept_controller!(integrator, controller::IController, alg, q)
  @unpack qsteady_min, qsteady_max = integrator.opts

  if qsteady_min <= q <= qsteady_max
    q = one(q)
  end
  integrator.dt / q # new dt
end

function step_reject_controller!(integrator, controller::IController, alg)
  @unpack qold = integrator
  integrator.dt = qold
end

# PI step size controller
"""
    PIController(beta1, beta2)

The proportional-integral (PI) controller is a widespread step size controller
with improved stability properties compared to the [`IController`](@ref).
This controller is the default for most algorithms in OrdinaryDiffEq.jl.

Construct a PI step size controller adapting the time step based on the formula
```
Δtₙ₊₁ = εₙ₊₁^β₁ * εₙ^β₂ * Δtₙ
```
where `εᵢ` are inverses of the error estimates scaled by the tolerance
(Hairer, Nørsett, Wanner, 2010, Section IV.2).
The step size factor is multiplied by the safety factor `gamma` and clipped to
the interval `[qmin, qmax]`.
A step will be accepted whenever the estimated error `integrator.EEst` is
less than or equal to unity. Otherwise, the step is rejected and re-tried with
the predicted step size.

!!! note
    The coefficients `beta1, beta2` are not scaled by the order of the method,
    in contrast to the [`PIDController`](@ref). For the `PIController`, this
    scaling by the order must be done when the controller is constructed.

## References
- Hairer, Nørsett, Wanner (2010)
  Solving Ordinary Differential Equations II Stiff and Differential-Algebraic Problems
  [DOI: 10.1007/978-3-642-05221-7](https://doi.org/10.1007/978-3-642-05221-7)
- Hairer, Nørsett, Wanner (2008)
  Solving Ordinary Differential Equations I Nonstiff Problems
  [DOI: 10.1007/978-3-540-78862-1](https://doi.org/10.1007/978-3-540-78862-1)
"""
mutable struct PIController{QT} <: AbstractController
  beta1::QT
  beta2::QT
end

@inline function stepsize_controller!(integrator, controller::PIController, alg)
  @unpack qold = integrator
  @unpack qmin, qmax, gamma = integrator.opts
  @unpack beta1, beta2 = controller
  EEst = DiffEqBase.value(integrator.EEst)

  if iszero(EEst)
    q = inv(qmax)
  else
    q11 = DiffEqBase.fastpow(EEst, float(beta1))
    q = q11 / DiffEqBase.fastpow(qold, float(beta2))
    integrator.q11 = q11
    @fastmath q = max(inv(qmax), min(inv(qmin), q / gamma))
  end
  q
end

function step_accept_controller!(integrator, controller::PIController, alg, q)
  @unpack qsteady_min, qsteady_max, qoldinit = integrator.opts
  EEst = DiffEqBase.value(integrator.EEst)

  if qsteady_min <= q <= qsteady_max
    q = one(q)
  end
  integrator.qold = max(EEst, qoldinit)
  return integrator.dt / q # new dt
end

function step_reject_controller!(integrator, controller::PIController, alg)
  @unpack q11 = integrator
  @unpack qmin, gamma = integrator.opts
  integrator.dt /= min(inv(qmin), q11 / gamma)
end

function reset_alg_dependent_opts!(controller::PIController, alg1, alg2)
  if controller.beta2 == beta2_default(alg1)
    controller.beta2 = beta2_default(alg2)
  end
  if controller.beta1 == beta1_default(alg1, controller.beta2)
    controller.beta1 = beta1_default(alg2, controller.beta2)
  end
end

# PID step size controller
"""
    PIDController(beta1, beta2, beta3=zero(beta1);
                  limiter=default_dt_factor_limiter,
                  accept_safety=0.81)

The proportional-integral-derivative (PID) controller is a generalization of the
[`PIController`](@ref) and can have improved stability and efficiency properties.

Construct a PID step size controller adapting the time step based on the formula
```
Δtₙ₊₁ = εₙ₊₁^(β₁/k) * εₙ^(β₂/k) * εₙ₋₁^(β₃/ k) * Δtₙ
```
where `k = min(alg_order, alg_adaptive_order) + 1` and `εᵢ` are inverses of
the error estimates scaled by the tolerance (Söderlind, 2003).
The step size factor is limited by the `limiter` with default value
```
limiter(x) = one(x) + atan(x - one(x))
```
as proposed by Söderlind and Wang (2006). A step will be accepted whenever the
predicted step size change is bigger than `accept_safety`. Otherwise, the step
is rejected and re-tried with the predicted step size.

Some standard controller parameters suggested in the literature are

| Controller | `beta1` | `beta2` | `beta3` |
|:-----------|--------:|--------:|:-------:|
|    basic   |  `1.00` |  `0.00` |  `0`    |
|    PI42    |  `0.60` | `-0.20` |  `0`    |
|    PI33    |  `2//3` | `-1//3` |  `0`    |
|    PI34    |  `0.70` | `-0.40` |  `0`    |
|   H211PI   |  `1//6` |  `1//6` |  `0`    |
|   H312PID  | `1//18` |  `1//9` | `1//18` |

!!! note
    In contrast to the [`PIController`](@ref), the coefficients `beta1, beta2, beta3`
    are scaled by the order of the method. Thus, standard controllers such as PI42
    can use the same coefficients `beta1, beta2, beta3` for different algorithms.

!!! note
    In contrast to other controllers, the `PIDController` does not use the keyword
    arguments `qmin, qmax` to limit the step size change or the safety factor `gamma`.
    These common keyword arguments are replaced by the `limiter` and `accept_safety`
    to guarantee a smooth behavior (Söderlind and Wang, 2006).
    Because of this, a `PIDController` behaves different from a [`PIController`](@ref),
    even if `beta1, beta2` are adapted accordingly and `iszero(beta3)`.

## References
- Söderlind (2003)
  Digital Filters in Adaptive Time-Stepping
  [DOI: 10.1145/641876.641877](https://doi.org/10.1145/641876.641877)
- Söderlind, Wang (2006)
  Adaptive time-stepping and computational stability
  [DOI: 10.1016/j.cam.2005.03.008](https://doi.org/10.1016/j.cam.2005.03.008)
- Ranocha, Dalcin, Parsani, Ketcheson (2021)
  Optimized Runge-Kutta Methods with Automatic Step Size Control for
  Compressible Computational Fluid Dynamics
  [arXiv:2104.06836](https://arxiv.org/abs/2104.06836)
"""
struct PIDController{QT, Limiter} <: AbstractController
  beta::MVector{3,QT} # controller coefficients
  err ::MVector{3,QT} # history of the error estimates
  accept_safety::QT   # accept a step if the predicted change of the step size
                      # is bigger than this parameter
  limiter::Limiter    # limiter of the dt factor (before clipping)
end

function PIDController(beta1, beta2, beta3=zero(beta1); limiter=default_dt_factor_limiter,
                                                        accept_safety=0.81)
  beta = MVector(map(float, promote(beta1, beta2, beta3))...)
  QT = eltype(beta)
  err = MVector{3,QT}(true, true, true)
  return PIDController(beta, err, convert(QT, accept_safety), limiter)
end

function Base.show(io::IO, controller::PIDController)
  print(io, "PIDController(beta=", controller.beta,
            ", accept_safety=", controller.accept_safety,
            ", limiter=", controller.limiter,
            ")")
end

@inline default_dt_factor_limiter(x) = one(x) + atan(x - one(x))

@inline function stepsize_controller!(integrator, controller::PIDController, alg)
  @unpack qmax = integrator.opts
  beta1, beta2, beta3 = controller.beta
  controller.err[1] = inv(DiffEqBase.value(integrator.EEst))
  err1, err2, err3 = controller.err

  iszero(err1) && return qmax

  k = min(alg_order(alg), alg_adaptive_order(alg)) + 1
  dt_factor = err1^(beta1 / k) * err2^(beta2 / k) * err3^(beta3 / k)
  dt_factor = controller.limiter(dt_factor)
  # Note: No additional limiting of the form
  #   dt_factor = max(qmin, min(qmax, dt_factor))
  # is necessary since the `limiter` should take care of that. The default limiter
  # ensures
  #   0.21 ≈ limiter(0) <= dt_factor <= limiter(Inf) ≈ 2.57
  # See Söderlind, Wang (2006), Section 6.
  integrator.qold = dt_factor
  return dt_factor
end

@inline function accept_step_controller(integrator, controller::PIDController)
  return integrator.qold >= controller.accept_safety
end

function step_accept_controller!(integrator, controller::PIDController, alg, dt_factor)
  @unpack qsteady_min, qsteady_max = integrator.opts

  if qsteady_min <= inv(dt_factor) <= qsteady_max
    dt_factor = one(dt_factor)
  end
  @inbounds begin
    controller.err[3] = controller.err[2]
    controller.err[2] = controller.err[1]
  end
  return integrator.dt * dt_factor # new dt
end

function step_reject_controller!(integrator, controller::PIDController, alg)
  integrator.dt *= integrator.qold
end

# Gustafsson predictive step size controller
"""
    PredictiveController()

The Gustafsson acceleration algorithm accelerates changes so that way algorithms
can more swiftly change to handle quick transients. This algorithm is thus
well-suited for stiff solvers where this can be expected, and is the default
for algorithms like the (E)SDIRK methods.

```julia
gamma = integrator.opts.gamma
niters = integrator.cache.newton_iters
fac = min(gamma,(1+2*integrator.alg.max_newton_iter)*gamma/(niters+2*integrator.alg.max_newton_iter))
expo = 1/(alg_order(integrator.alg)+1)
qtmp = (integrator.EEst^expo)/fac
@fastmath q = max(inv(integrator.opts.qmax),min(inv(integrator.opts.qmin),qtmp))
if q <= integrator.opts.qsteady_max && q >= integrator.opts.qsteady_min
  q = one(q)
end
integrator.qold = q
q
```
In this case, `niters` is the number of Newton iterations which was required
in the most recent step of the algorithm. Note that these values are used
differently depending on acceptance and rejectance. When the step is accepted,
the following logic is applied:
```julia
if integrator.success_iter > 0
  expo = 1/(alg_adaptive_order(integrator.alg)+1)
  qgus=(integrator.dtacc/integrator.dt)*(((integrator.EEst^2)/integrator.erracc)^expo)
  qgus = max(inv(integrator.opts.qmax),min(inv(integrator.opts.qmin),qgus/integrator.opts.gamma))
  qacc=max(q,qgus)
else
  qacc = q
end
integrator.dtacc = integrator.dt
integrator.erracc = max(1e-2,integrator.EEst)
integrator.dt/qacc
```
When it rejects, it's the same as the [`IController`](@ref):
```julia
if integrator.success_iter == 0
  integrator.dt *= 0.1
else
  integrator.dt = integrator.dt/integrator.qold
end
```
"""
struct PredictiveController <: AbstractController
end

@inline function stepsize_controller!(integrator, controller::PredictiveController, alg)
  @unpack qmin, qmax, gamma = integrator.opts
  EEst = DiffEqBase.value(integrator.EEst)

  if iszero(EEst)
    q = inv(qmax)
  else
    if alg isa Union{RKC,IRKC,SERK2}
      fac = gamma
    else
      if alg isa Union{RadauIIA3, RadauIIA5}
        @unpack iter = integrator.cache
        @unpack maxiters = alg
      else
        @unpack iter, maxiters = integrator.cache.nlsolver
      end
      fac = min(gamma, ( 1 + 2 * maxiters) * gamma / (iter + 2 * maxiters))
    end
    expo = 1 / (get_current_adaptive_order(alg, integrator.cache) + 1)
    qtmp = DiffEqBase.fastpow(EEst, expo) / fac
    @fastmath q = DiffEqBase.value(max(inv(qmax), min(inv(qmin), qtmp)))
    integrator.qold = q
  end
  q
end

function step_accept_controller!(integrator, controller::PredictiveController, alg, q)
  @unpack qmin, qmax, gamma, qsteady_min, qsteady_max = integrator.opts
  EEst = DiffEqBase.value(integrator.EEst)

  if integrator.success_iter > 0
    expo = 1 / (get_current_adaptive_order(alg, integrator.cache) + 1)
    qgus = (integrator.dtacc / integrator.dt) * DiffEqBase.fastpow((EEst^2) / integrator.erracc, expo)
    qgus = max(inv(qmax), min(inv(qmin), qgus / gamma))
    qacc = max(q, qgus)
  else
    qacc = q
  end
  if qsteady_min <= qacc <= qsteady_max
    qacc = one(qacc)
  end
  integrator.dtacc = integrator.dt
  integrator.erracc = max(1e-2, EEst)
  return integrator.dt / qacc
end

function step_reject_controller!(integrator, controller::PredictiveController, alg)
  @unpack dt, success_iter, qold = integrator
  integrator.dt = success_iter == 0 ? 0.1 * dt : dt / qold
end

# Dummy controller without any method implementations.
# This is used to transfer the special controllers associated to certain
# algorithms to the new controller infrastructure with
struct DummyController <: AbstractController
end

# JVODE
function stepsize_controller!(integrator, alg::JVODE)
  if iszero(integrator.EEst)
    η = integrator.opts.qmax
  else
    η = integrator.cache.η
    integrator.qold = η
  end
  η
end

function step_accept_controller!(integrator,alg::JVODE,η)
  q = inv(η)
  if q <= integrator.opts.qsteady_max && q >= integrator.opts.qsteady_min
    q = one(q)
  end
  return integrator.dt/q  # dtnew
end

function step_reject_controller!(integrator,alg::JVODE)
  integrator.dt *= integrator.qold
end

# QNBDF
stepsize_controller!(integrator, alg::QNDF) = nothing

# this stepsize and order controller is taken from
# Implementation of an Adaptive BDF2 Formula and Comparison with the MATLAB Ode15s paper
# E. Alberdi Celaya, J. J. Anza Aguirrezabala, and P. Chatzipantelidis

function step_accept_controller!(integrator,alg::QNDF{max_order},q) where max_order
  #step is accepted, reset count of consecutive failed steps
  integrator.cache.consfailcnt = 0
  integrator.cache.nconsteps += 1
  if iszero(integrator.EEst)
    q = inv(integrator.opts.qmax)
  else
    est = integrator.EEst
    estₖ₋₁ = integrator.cache.EEst1
    estₖ₊₁ = integrator.cache.EEst2
    h = integrator.dt
    k = integrator.cache.order
    cache = integrator.cache
    if integrator.cache.nconsteps < integrator.cache.order + 2
      q = one(integrator.qold) #quasi-contsant steps
    else
      zₛ = 1.2 # equivalent to integrator.opts.gamma
      zᵤ = 0.1
      Fᵤ = 10
      expo = 1/(k+1)
      z = zₛ * ((est)^expo)
      F = inv(z)
      hₙ = h
      kₙ = k
      if z <= zᵤ
        hₖ = Fᵤ * h
      elseif zᵤ < z
        hₖ = F * h
      end
      hₖ₋₁ = 0.0
      hₖ₊₁ = 0.0

      if k > 1
        expo = 1/k
        zₖ₋₁ = 1.3 * ((estₖ₋₁)^expo)
        Fₖ₋₁ = inv(zₖ₋₁)
        if zₖ₋₁ <= 0.1
          hₖ₋₁ =  10* h
        elseif 1/10 < zₖ₋₁ <= 1.3
          hₖ₋₁ = Fₖ₋₁ * h
        end
        if hₖ₋₁ > hₖ
          hₙ = hₖ₋₁
          kₙ = k-1
        else
          hₙ = hₖ
          kₙ = k
        end
      end

      if k < max_order
        expo = 1/(k+2)
        zₖ₊₁ = 1.4 * ((estₖ₊₁)^expo)
        Fₖ₊₁ = inv(zₖ₊₁)

        if zₖ₊₁<= 0.1
          hₖ₊₁ = 10 * h
        elseif 0.1 < zₖ₊₁ <= 1.4
          hₖ₊₁ = Fₖ₊₁ * h
        end
        if hₖ₊₁ > hₙ
          hₙ = hₖ₊₁
          kₙ = k+1
        end
      end
      if hₙ <= h
        hₙ = h
        kₙ = k
      end
      cache.order = kₙ
      q = integrator.dt/hₙ
    end
  end
  if q <= integrator.opts.qsteady_max && q >= integrator.opts.qsteady_min
    q = one(q)
  end
  return integrator.dt/q
end

function step_reject_controller!(integrator,alg::QNDF)
  k = integrator.cache.order
  h = integrator.dt
  integrator.cache.consfailcnt += 1
  integrator.cache.nconsteps = 0
  if integrator.cache.consfailcnt > 1
    h = h/2
  end
  zₛ = 1.2  # equivalent to integrator.opts.gamma
  expo = 1/(k+1)
  z = zₛ * ((integrator.EEst)^expo)
  F = inv(z)
  if z <= 10
    hₖ = F * h
  elseif z > 10
    hₖ = 0.1 * h
  end
  hₙ = hₖ
  kₙ = k
  if k > 1
    expo = 1/k
    zₖ₋₁ = 1.3 * ((integrator.cache.EEst1)^expo)
    Fₖ₋₁ = inv(zₖ₋₁)
    if zₖ₋₁ <= 10
      hₖ₋₁ = Fₖ₋₁ * h
    elseif zₖ₋₁ > 10
      hₖ₋₁ = 0.1 * h
    end
    if hₖ₋₁ > hₖ
      hₙ = min(h,hₖ₋₁)
      kₙ = k-1
    end
  end
  integrator.dt = hₙ
  integrator.cache.order = kₙ
end

function choose_order!(alg::Union{FBDF,DFBDF}, integrator, cache::OrdinaryDiffEqMutableCache, max_order)
  @unpack t,dt,u,cache,uprev = integrator
  @unpack atmp, ts_tmp, terkm2, terkm1, terk, terkp1, terk_tmp, u_history = cache
  k = cache.order
  # only when the order of amount of terk follows the order of step size, and achieve enough constant step size, the order could be increased.
  if k < max_order && integrator.cache.nconsteps >= integrator.cache.order + 2 && ((k == 1 && terk > terkp1) ||
    (k == 2 && terkm1 > terk > terkp1) ||
    (k > 2 && terkm2 > terkm1 > terk > terkp1))
    k += 1
    terk = terkp1
  else
    while !(terkm2 > terkm1 > terk > terkp1) && k > 2
      terkp1 = terk
      terk = terkm1
      terkm1 = terkm2
      fd_weights = calc_finite_difference_weights(ts_tmp,t+dt,k-2,Val(max_order))
      terk_tmp = @.. fd_weights[k-2,1] * u
      vc = _vec(terk_tmp)
      for i in 2:k-2
        @.. @views vc += fd_weights[i,k-2] * u_history[:,i-1]
      end
      @.. terk_tmp *= abs(dt^(k-2))
      calculate_residuals!(atmp,_vec(terk_tmp), _vec(uprev), _vec(u), integrator.opts.abstol, integrator.opts.reltol, integrator.opts.internalnorm, t)
      terkm2 = integrator.opts.internalnorm(atmp,t)
      k -= 1
    end
    if !(terkm1 > terk > terkp1) && k == 2
      k -= 1
      terk = terkm1
    end
  end
  return k, terk
end

function choose_order!(alg::Union{FBDF,DFBDF}, integrator, cache::OrdinaryDiffEqConstantCache, max_order)
  @unpack t,dt,u,cache,uprev = integrator
  @unpack ts_tmp,terkm2, terkm1, terk, terkp1,u_history = cache
  k = cache.order
  if k < max_order && integrator.cache.nconsteps >= integrator.cache.order + 2 && ((k == 1 && terk > terkp1) ||
    (k == 2 && terkm1 > terk > terkp1) ||
    (k > 2 && terkm2 > terkm1 > terk > terkp1))
    k += 1
    terk = terkp1
  else
    while !(terkm2 > terkm1 > terk > terkp1) && k > 2
      terkp1 = terk
      terk = terkm1
      terkm1 = terkm2
      fd_weights = calc_finite_difference_weights(ts_tmp,t+dt,k-2,Val(max_order))
      terk_tmp = @.. fd_weights[k-2,1] * u
      if typeof(u) <: Number
        for i in 2:k-2
          terk_tmp += fd_weights[i,k-2] * u_history[i-1]
        end
        #@show fd_weights,u_history,u,terk
        terk_tmp *= abs(dt^(k-2))
      else
        vc = _vec(terk_tmp)
        for i in 2:k-2
          @.. @views vc += fd_weights[i,k-2] * u_history[:,i-1]
        end
        @.. terk_tmp *= abs(dt^(k-2))
      end
      atmp = calculate_residuals(_vec(terk_tmp), _vec(uprev), _vec(u), integrator.opts.abstol, integrator.opts.reltol, integrator.opts.internalnorm, t)
      terkm2 = integrator.opts.internalnorm(atmp,t)
      k -= 1
    end
    if !(terkm1 > terk > terkp1) && k == 2
      k -= 1
      terk = terkm1
    end
  end
  return k, terk
end

function stepsize_controller!(integrator, alg::Union{FBDF{max_order},DFBDF{max_order}}) where max_order
  @unpack cache = integrator
  cache.prev_order = cache.order
  k, terk = choose_order!(alg, integrator, cache, max_order)
  if k != cache.order
    integrator.cache.nconsteps = 0
    cache.order = k
  end
  if iszero(terk)
    q = inv(integrator.opts.qmax)
  else
    q = ((2*terk/(k+1))^(1/(k+1)))
  end
  integrator.qold = q
  q
end

function step_accept_controller!(integrator,alg::Union{FBDF{max_order},DFBDF{max_order}},q) where max_order
  integrator.cache.consfailcnt = 0
  if q <= integrator.opts.qsteady_max && q >= integrator.opts.qsteady_min
    q = one(q)
  end
  integrator.cache.nconsteps += 1
  integrator.cache.iters_from_event += 1
  return integrator.dt/q
end

function step_reject_controller!(integrator,alg::Union{FBDF,DFBDF})
  integrator.cache.consfailcnt += 1
  integrator.cache.nconsteps = 0
  dt = integrator.dt
  if integrator.cache.consfailcnt == 1
    dt *= 0.5
  else
    dt *= 0.25
  end
  integrator.dt = dt
end

# Extrapolation methods
mutable struct ExtrapolationController{QT} <: AbstractController
  beta1::QT
end

function reset_alg_dependent_opts!(controller::ExtrapolationController, alg1, alg2)
  if controller.beta1 == beta1_default(alg1, beta2_default(alg1))
    controller.beta1 = beta1_default(alg2, beta2_default(alg2))
  end
end

@inline function stepsize_controller!(integrator,alg::Union{ExtrapolationMidpointDeuflhard,ImplicitDeuflhardExtrapolation})
  # Dummy function
  # ExtrapolationMidpointDeuflhard's stepsize scaling is stored in the cache;
  # it is computed by  stepsize_controller_internal! (in perfom_step!) resp. stepsize_predictor!
  # (in step_accept_controller! and step_reject_controller!)
  zero(typeof(integrator.opts.qmax))
end

function stepsize_controller_internal!(integrator,alg::Union{ExtrapolationMidpointDeuflhard,ImplicitDeuflhardExtrapolation})
  # Standard step size controller
  # Compute and save the stepsize scaling based on the latest error estimate of the current order
  @unpack controller = integrator.opts

  if iszero(integrator.EEst)
    q = inv(integrator.opts.qmax)
  else
    # Update gamma and beta1
    controller.beta1 = typeof(controller.beta1)(1 // (2integrator.cache.n_curr + 1))
    integrator.opts.gamma = DiffEqBase.fastpow(typeof(integrator.opts.gamma)(1 // 4),controller.beta1)
    # Compute new stepsize scaling
    qtmp = DiffEqBase.fastpow(integrator.EEst,controller.beta1) / integrator.opts.gamma
    @fastmath q = max(inv(integrator.opts.qmax), min(inv(integrator.opts.qmin), qtmp))
  end
  integrator.cache.Q[integrator.cache.n_curr - alg.n_min + 1] = q
end

function stepsize_predictor!(integrator,alg::Union{ExtrapolationMidpointDeuflhard,ImplicitDeuflhardExtrapolation},n_new::Int)
  # Compute and save the stepsize scaling for order n_new based on the latest error estimate of the current order.
  @unpack controller = integrator.opts

  if iszero(integrator.EEst)
    q = inv(integrator.opts.qmax)
  else
    # Initialize
    @unpack t,EEst = integrator
    @unpack stage_number = integrator.cache
    tol = integrator.opts.internalnorm(integrator.opts.reltol,t) # Deuflhard's approach relies on EEstD ≈ ||relTol||
    s_curr = stage_number[integrator.cache.n_curr - alg.n_min + 1]
    s_new = stage_number[n_new - alg.n_min + 1]
    # Update gamma and beta1
    controller.beta1 = typeof(controller.beta1)(1 // (2integrator.cache.n_curr + 1))
    integrator.opts.gamma = DiffEqBase.fastpow(typeof(integrator.opts.gamma)(1 // 4),controller.beta1)
    # Compute new stepsize scaling
    qtmp = EEst * DiffEqBase.fastpow(DiffEqBase.fastpow(tol,(1.0 - s_curr / s_new)),controller.beta1) / integrator.opts.gamma
    @fastmath q = max(inv(integrator.opts.qmax),min(inv(integrator.opts.qmin),qtmp))
  end
  integrator.cache.Q[n_new - alg.n_min + 1] = q
end

function step_accept_controller!(integrator,alg::Union{ExtrapolationMidpointDeuflhard,ImplicitDeuflhardExtrapolation},q)
  # Compute new order and stepsize, return new stepsize
  @unpack n_min, n_max = alg
  @unpack n_curr, n_old, Q = integrator.cache
  s = integrator.cache.stage_number

  # Compute new order based on available quantities
  tmp = (n_min:n_curr) .- n_min .+ 1 # Index range of quantities computed so far
  dt_new = Vector{eltype(Q)}(undef,length(tmp)+1)
  dt_new[1:end-1] = integrator.dt ./ Q[tmp] # Store for the possible new stepsizes
  dtmin = timedepentdtmin(integrator)
  dt_new[1:end-1] = max.(dtmin, min.(abs(integrator.opts.dtmax), abs.(dt_new[1:end-1]))) # Safety scaling

  # n_new is the most efficient order of the last step
  work = s[tmp] ./ dt_new[1:end-1]
  n_new = argmin(work) + n_min - 1

  # Check if n_new may be increased
  if n_new == n_curr < min(n_max, n_old + 1) # cf. win_max in perfom_step! of the last step
    # Predict stepsize scaling for order (n_new + 1)
    stepsize_predictor!(integrator, alg, n_new+1) # Update cache.Q

    # Compute and scale the corresponding stepsize
    dt_new[end] = integrator.dt ./ Q[tmp[end]+1]
    dt_new[end] = max(dtmin, min(abs(integrator.opts.dtmax), abs.(dt_new[end])))

    # Check if (n_new  + 1) would have been more efficient than n_new
    if work[end] > s[tmp[end]+1] / dt_new[end]
      n_new = n_new + 1
    end
  end

  integrator.cache.n_curr = n_new
  dt_new[n_new - n_min + 1]
end

function step_reject_controller!(integrator, alg::Union{ExtrapolationMidpointDeuflhard,ImplicitDeuflhardExtrapolation})
  # Compute and save reduced stepsize dt_red of order n_old
  # Use the latest error estimate to predict dt_red if an estimate of order n_old is not available
  if integrator.cache.n_curr < integrator.cache.n_old
      stepsize_predictor!(integrator,alg,integrator.cache.n_old) # Update cache.Q
  end
  integrator.cache.n_curr = integrator.cache.n_old # Reset order for redoing the rejected step
  dt_red = integrator.dt / integrator.cache.Q[integrator.cache.n_old - integrator.alg.n_min + 1]
  dtmin = timedepentdtmin(integrator)
  dt_red = integrator.tdir*max(dtmin, min(abs(integrator.opts.dtmax), abs(dt_red))) # Safety scaling
  integrator.dt = dt_red
end

@inline function stepsize_controller!(integrator,alg::Union{ExtrapolationMidpointHairerWanner, ImplicitHairerWannerExtrapolation, ImplicitEulerExtrapolation, ImplicitEulerBarycentricExtrapolation})
  # Dummy function
  # ExtrapolationMidpointHairerWanner's stepsize scaling is stored in the cache;
  # it is computed by  stepsize_controller_internal! (in perfom_step!), step_accept_controller! or step_reject_controller!
  zero(typeof(integrator.opts.qmax))
end

function stepsize_controller_internal!(integrator,alg::Union{ExtrapolationMidpointHairerWanner, ImplicitHairerWannerExtrapolation, ImplicitEulerExtrapolation, ImplicitEulerBarycentricExtrapolation})
  # Standard step size controller
  # Compute and save the stepsize scaling based on the latest error estimate of the current order
  @unpack controller = integrator.opts

  if typeof(alg) <: Union{ImplicitEulerExtrapolation,ImplicitEulerBarycentricExtrapolation,ImplicitHairerWannerExtrapolation}
    if iszero(integrator.EEst)
      q = inv(integrator.opts.qmax)
    else
      # Update gamma and beta1
      if typeof(alg) <: ImplicitHairerWannerExtrapolation
        controller.beta1 = typeof(controller.beta1)(1 // (2integrator.cache.n_curr + 1))
      elseif typeof(alg) <: ImplicitEulerExtrapolation
        controller.beta1 = typeof(controller.beta1)(1 // (integrator.cache.n_curr))
      else
        controller.beta1 = typeof(controller.beta1)(1 // (integrator.cache.n_curr - 1))
      end
      integrator.opts.gamma = DiffEqBase.fastpow(typeof(integrator.opts.gamma)(65 // 100),controller.beta1)
      # Compute new stepsize scaling
      qtmp = DiffEqBase.fastpow(integrator.EEst,controller.beta1) / (integrator.opts.gamma)
      @fastmath q = max(inv(integrator.opts.qmax), min(inv(integrator.opts.qmin), qtmp))
    end
    integrator.cache.Q[integrator.cache.n_curr + 1] = q
  else
    if iszero(integrator.EEst)
      q = inv(integrator.opts.qmax)
    else
      # Update gamma and beta1
      controller.beta1 = typeof(controller.beta1)(1 // (2integrator.cache.n_curr + 1))
      integrator.opts.gamma = DiffEqBase.fastpow(typeof(integrator.opts.gamma)(65 // 100),controller.beta1)
      # Compute new stepsize scaling
      qtmp = DiffEqBase.fastpow(integrator.EEst,controller.beta1) / integrator.opts.gamma
      @fastmath q = max(inv(integrator.opts.qmax), min(inv(integrator.opts.qmin), qtmp))
    end
    integrator.cache.Q[integrator.cache.n_curr + 1] = q
  end
end

function step_accept_controller!(integrator,alg::Union{ExtrapolationMidpointHairerWanner, ImplicitHairerWannerExtrapolation, ImplicitEulerExtrapolation, ImplicitEulerBarycentricExtrapolation},q)
  # Compute new order and stepsize, return new stepsize
  @unpack n_min, n_max = alg
  @unpack n_curr, n_old, Q, sigma, work, dt_new = integrator.cache
  s = integrator.cache.stage_number

  # Compute new order based on available quantities
  win_min_old = min(n_old, n_curr) - 1 # cf. win_min in perfom_step! of the last step
  tmp = win_min_old:(max(n_curr, n_old) + 1) # Index range for the new order
  #@show size(dt_new)
  fill!(dt_new, zero(eltype(dt_new)))
  @.. Q = integrator.dt/Q
  copyto!(dt_new,win_min_old,Q,win_min_old,(max(n_curr, n_old) + 1) - win_min_old + 1)
  @.. Q = integrator.dt/Q
  dtmin = timedepentdtmin(integrator)
  fill!(work,zero(eltype(work))) # work[n] is the work for order (n-1)
  for i in tmp
    work[i] = s[i]/dt_new[i]
  end
  # Order selection
  n_new = n_old
  if n_curr == n_min # Enforce n_min + 1 ≦ n_new
    n_new = n_min + 1
  else
    if n_curr <= n_old
      if work[n_curr-1] < sigma * work[n_curr]
        n_new = max(n_curr-1,n_old-1,n_min+1) # Enforce n_min + 1≦ n_new
      elseif work[n_curr] < sigma * work[n_curr-1]
        n_new = min(n_curr+1,n_max-1) # Enforce n_new ≦ n_max - 1
      else
        n_new = n_curr # n_min + 1 ≦ n_curr
      end
    else
      if work[n_old] < sigma *  work[n_old+1]
        n_new = max(n_old-1,n_min+1)  # Enforce n_min + 1 ≦ n_new
      end
      if work[n_curr+1] <  sigma * work[n_new+1]
        n_new = min(n_new+1,n_max-1) # Enforce n_new ≦ n_max - 1
      end
    end
  end
  integrator.cache.n_curr = n_new

  # Stepsize selection
  if n_new == n_curr + 1
    # Compute the new stepsize of order n_new based on the optimal stepsize of order n_curr
    dt_new[n_new+1] = s[n_curr + 2]/s[n_curr + 1 ] * dt_new[n_curr+1]
    dt_new[n_new+1] = max(dtmin, min(abs(integrator.opts.dtmax), abs(dt_new[n_new+1])))
  end
  dt_new[n_new + 1]
end

function step_reject_controller!(integrator, alg::Union{ExtrapolationMidpointHairerWanner, ImplicitHairerWannerExtrapolation, ImplicitEulerExtrapolation, ImplicitEulerBarycentricExtrapolation})
  # Compute and save order and stepsize for redoing the current step
  @unpack n_old, n_curr, Q = integrator.cache

  # Order selection
  n_red = n_old
  if n_curr == n_old - 1
    n_red = max(alg.n_min+1,n_old-1) # Enforce n_min + 1 ≦ n_red
  end
  integrator.cache.n_curr = n_red

  # Stepsize selection
  dt_red = integrator.dt / Q[n_red + 1]
  dtmin = timedepentdtmin(integrator)
  dt_red = integrator.tdir*max(dtmin, min(abs(integrator.opts.dtmax), abs(dt_red))) # Safety scaling
  integrator.dt = dt_red
end


text:
# Constants hard wired everywhere

const pkgDir = normpath(@__DIR__, "..");
const dataDir = joinpath(pkgDir, "data");

## -------------  Data types

# const Double = Float64
# # For no of courses
# const ncInt = UInt8
# # For time variables
# const TimeInt = UInt8
# # For indexing types
# # const TypeInt = UInt16
# # For indexing colleges
# const CollInt = UInt8
# # For indexing school levels
# const SchoolInt = UInt8
# For indexing grid points
# const GridInt = UInt16
# const CaseNameType = Union{Symbol, Vector{Symbol}}

const MtMean = :mean;
const MtStd = :std;
const MtCount = :count;

## -------------  Normalizations and Bounds

# Number of data courses per model course
# const dataCoursesPerCourse = 1

## ------------  Data

# Comment string in CSV files
const commentStr = "#"

## --------------  Debugging

# const dbgLow = true
# const dbgMedium = true
# const dbgHigh = true

# # Required Julia version (as a version string)
# const minVersion = "1.4"

# ----------


text:
# The original license on the t-tests from HypothesisTests.jl
# is as follows. We're including it here because we're mostly copying the
# documentation of the functions.
#
# jarque_bera.jl
# Jarque-Bera goodness-of-fit test
#
# Copyright (C) 2017   Benjamin Born
#
# Permission is hereby granted, free of charge, to any person obtaining
# a copy of this software and associated documentation files (the
# "Software"), to deal in the Software without restriction, including
# without limitation the rights to use, copy, modify, merge, publish,
# distribute, sublicense, and/or sell copies of the Software, and to
# permit persons to whom the Software is furnished to do so, subject to
# the following conditions:
#
# The above copyright notice and this permission notice shall be
# included in all copies or substantial portions of the Software.
#
# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND,
# EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF
# MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND
# NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR COPYRIGHT HOLDERS BE
# LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION
# OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION
# WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE.

import HypothesisTests.JarqueBeraTest

"""
    JarqueBeraTest(d::AbstractUncertainValue, n::Int = 1000) -> JarqueBeraTest

Compute the Jarque-Bera statistic to test the null hypothesis that an uncertain
value is normally distributed.
"""
function JarqueBeraTest(d::AbstractUncertainValue, n::Int = 1000)
    x = resample(d, n)
    JarqueBeraTest(x)
end

"""
    JarqueBeraTestPooled(ud::UncertainDataset, n::Int = 1000) -> JarqueBeraTest

First, draw `n` realisations of each uncertain value in `ud` and pool them
together. Then, compute the Jarque-Bera statistic to test the null hypothesis
that the values of the pool are normally distributed.
"""
function JarqueBeraTestPooled(ud::UncertainDataset, n::Int = 1000)
    x = vcat(resample(ud, n)...,)
    JarqueBeraTest(x)
end

"""
    OneSampleADTestElementWise(ud::UncertainDataset,
        n::Int = 1000) -> Vector{JarqueBeraTest}

First, draw `n` realisations of each uncertain value in `ud`, keeping one
pool of values for each uncertain value.

Then, compute the Jarque-Bera statistic to test the null hypothesis that each
value pool is normally distributed.
"""
function JarqueBeraTestElementWise(ud::UncertainDataset, n::Int = 1000)
    [JarqueBeraTest(resample(ud[i], n)) for i = 1:length(ud)]
end

export
JarqueBeraTest,
JarqueBeraTestPooled,
JarqueBeraTestElementWise


text:
T_memoize = Dict{Expr,Expr}()

@testset "@T" begin
    T = collect(1:100)

    @T(A = t)
    @test A == collect(1:100)

    @T(B = t)
    @test B == collect(1:100)

    @T begin D = t + T[t] end
    @test D == collect(1:100)*2

    @T(C = ( t==1  ? 51 : C[t-1]+1 ))
    @test C == collect(51:50+100)

    q = 10
    @T(D = C + B + 7 + q)
    @test D == @.(C + B + 7 + q)

    @T begin sum_assured_t0 =
        10000
    end
    @test sum_assured_t0 == 10000

    @T H = t==1 ? 1 : A[t-1] + B
    @test H == [t==1 ? 1 : A[t-1] + B[t] for t ∈ T]

    @test_throws CalculationSequenceError @T a = a[t]

    T = collect(1:600)
    @T begin
        age_t0 = 30
        duration_t0 = 6
        premium_t0 = 100
        sum_assured_t0 = 10000
        age = floor(Int, age_t0 + t/12)
        duration = duration_t0 + t
        qx = age / 1000
        lapse_rate_m = 1 / qx / 1000
        no_pols_som = if t==1 1 else no_pols_eom[t-1] end
        no_deaths = no_pols_som*qx/12*(1-0.5*lapse_rate_m)
        no_surrenders = no_pols_som*lapse_rate_m*(1-0.5*qx)
        no_pols_eom = if t==1 1 else no_pols_eom[t-1] end - no_deaths - no_surrenders
        premium_income = premium_t0*no_pols_eom/12
        claims_outgo = sum_assured_t0*no_deaths
        profit = premium_income-claims_outgo
        yield_curve = 0.1
        profit_discounted  = profit*(1+yield_curve)^(-t)
        liability = -profit_discounted+ if t<600 liability[t+1] else 0 end
    end

end

@testset "cluster_topology" begin
    formulas = OrderedDict(
        :A => :((t == 1 ? 1 : A[t - 1]- C) - D),
        :B => :(t == 1 ? 1 : A[t - 1]),
        :C => :(((B * E) / 12) * (1 - 0.5*F)),
        :D => :(B * F * (1 - 0.5*E))
      )
      @test cluster_topology(formulas) == (-1, [:B, :C, :D, :A])
end

"
For a given cluster of formulae, derive the order in which these formulae
should be computed within `for t ∈ T`, and return this order along with
how t is referenced: [-1, 0, +1, nothing]
"
function cluster_topology(formulas::OrderedDict)
    refs_dict = Dict{Symbol, Any}()
    Δ_dict = Dict{Symbol, Any}()
    all_Δ = Set()

    # Get all the t references and convert them to offsets Δ
    for (key, value) ∈ formulas
        refs = (Set(var=>ex for (var, ex) ∈ get_t_refs(value) if haskey(formulas, var)) ∪
                Set(var=>0 for var ∈ vars_without_t_index(value) if haskey(formulas, var)))

        Δ = Set(var=>   if t_linearcombinations(expr) == true
                            eval(assign_variable(expr, :t, 0))
                        else
                            0 # TODO: more inspection needed
                        end for (var, expr) ∈ refs)

        refs_dict[key] = refs
        Δ_dict[key] = Δ
        all_Δ = all_Δ ∪ [δ for (_, δ) ∈ Δ]
    end

    # Continue only if all the references are either backwards or forwards looking and not mixed
    rtype = all(all_Δ .<= 0) ? -1 : all(all_Δ .>= 0) ? +1 : nothing
    if rtype === nothing
        return rtype, keys(formulas)
    end

    edges = Vector{Tuple{Symbol, Symbol}}()
    for (key, Δ) in Δ_dict
        for (var, δ) ∈ Δ
            if δ==0
                push!(edges, (var, key))
            end
        end
    end

    # Get sequence... flatten... append lost keys
    sequence = RowSheets.traversalsequence!(RowSheets.DiGraph(edges))
    sequence = [var for cluster in sequence for var in cluster]
    sequence = [sequence; setdiff!(OrderedSet(keys(formulas)), sequence)...]

    return rtype, sequence
end

function transform_expr_T(expr; line::Union{LineNumberNode, Nothing}=nothing)
    empty!(T_memoize) # make empty to suppress memoization
    expr′ = striplines(expr)
    expr isa Expr && haskey(T_memoize, expr′) && return T_memoize[expr′]

    formulas = expr_to_assignment_dict(expr; line)
    ret = transform_expr_all(formulas)

    # Add the last entry as the return value
    (lastvar, _) = last(formulas)
    push!(ret.args, lastvar)

    T_memoize[expr′] = ret
    return ret
end

"
Transform all formulae/equations in a given variable library `:(B = t==1 ? 1 : A[t-1]; A = B[t]; C=A+100)`
into a coherent forloop `:(for t ∈ T A[t] = B[t]; B[t] = t==1 ? 1 : A[t-1] end; @.(C = A+100))`.

This will traverse the full graph dependencies
"
function transform_expr_all(formulas::OrderedDict{Symbol, CodePoint})
    # Convert the dictionary into a sequence
    symbol_links = Dict(key => Vector{Symbol}() for key ∈ keys(formulas))
    for (varⱼ, codepoint) ∈ formulas
        ex = codepoint.expr
        for varᵢ ∈ get_variables(ex)
            haskey(formulas, varᵢ) && push!(symbol_links[varᵢ], varⱼ)
        end
    end
    sequence = RowSheets.traversalsequence!(RowSheets.DiGraph(symbol_links))
    sequence_bias = OrderedDict(key=>i for (i,key) ∈ enumerate(keys(formulas)))

    # Keep order in cluster same as input order (to allow user intervension)
    expr′_list = []
    for cluster ∈ sequence
        formulas′ = OrderedDict(i=>formulas[i] for i ∈ sort(cluster, by=x->(sequence_bias[x])))
        push!(expr′_list, transform_expr_cluster(formulas′))
    end

    return Expr(:block, expr′_list...)
end

"
Transform a cluster of formulae/equations that is possibly cyclic `:(B = t==1 ? 1 : A[t-1]; A = B[t])`
into a coherent forloop `:(for t ∈ T A[t] = B[t]; B[t] = t==1 ? 1 : A[t-1] end)`.

This can be used in each traversal step when traversing the full formuale dependency graph.
"
function transform_expr_cluster(formulas::OrderedDict{Symbol, CodePoint})

    # Single definition might have some good shortcuts
    ret = nothing
    if length(formulas) == 1
        (var, codepoint) = first(formulas)
        e = codepoint.expr
        if !has_symbol(e, :t)
            # Do `$var = @.($expr)`
            ret = Expr(:block,
                        codepoint.line,
                        Expr(:(=), var, Expr(:macrocall, Symbol("@__dot__"), codepoint.line, e)))

        elseif !has_symbol(e, var)
            # Do  [$var[t] = $expr′[t] for t ∈ T]
            ret =  Expr(:block,
                        codepoint.line,
                        Expr(:(=), var, Expr(:comprehension, Expr(:generator, wrap_t_getter(e), Expr(:(=), :t, :T)))))
        end
    end
    if ret === nothing
        # Long way around
        rtype, seq = cluster_topology(OrderedDict(i=>j.expr for (i,j) ∈ formulas))

        definitions = Expr(:block)
        for var ∈ seq
            push!(definitions.args, formulas[var].line)
            push!(definitions.args, :($(var) = Vector(undef, length(T))))
        end

        loopover =  rtype > 0 ? :(reverse(T)) : (:T)

        assignments = Expr(:block)
        for var ∈ seq
            push!(assignments.args, formulas[var].line)
            push!(assignments.args, :($(var)[t] = $(wrap_t_getter(formulas[var].expr))))
        end

        @gensymx e
        ret = (quote
            $definitions
            try
                for t in $loopover
                    $assignments
                end
            catch $e
                $e isa UndefRefError && throw($CalculationSequenceError())
                rethrow($e)
            end
            nothing
        end)
    end

    return ret
end

macro T(expr)
    esc(transform_expr_T(expr; line=__source__))
end

macro no_op(x...)
end


text:
push!(LOAD_PATH, string(pwd(), "/src/"))
using WRIdual
using LinearAlgebra, SparseArrays
using PyPlot
using JLD2

## Set slowness model and discretization step

n = (201, 201)
h = (10e0, 10e0)
val = 0.5e0
x = reshape((collect(1:n[1]).-1f0).*h[1], :, 1)
z = reshape((collect(1:n[2]).-1f0).*h[2], 1, :)
v = 1000e0.*(2e0.-val*exp.(-((x./1000e0.-1e0)/0.25e0).^2e0.-((z./1000e0.-1e0)/0.5e0).^2e0))
m = 1e0./v.^2e0

## Create data

# Set frequencies
freqs = [6e0]

# Set source/receiver geometry
ix_src = 3:4:199
Ns = length(ix_src)
iz_src = 2
x_src = (Float64.(ix_src).-1e0)*h[1]
z_src = (Float64.(iz_src).-1e0)*h[2]
ix_rcv = 3:199
iz_rcv = 200
x_rcv = (Float64.(ix_rcv).-1e0)*h[1]
z_rcv = (Float64.(iz_rcv).-1e0)*h[2]

# Setting source wavelet
src_enc = C.(Matrix(sparse(I, length(x_src), length(x_src))))

# Absorbing boundary condition padding
nabc = (33, 33, 33, 33)

# Solving PDE for each frequency
dat = gendata(m, h, freqs, x_rcv, z_rcv, x_src, z_src, src_enc; nabc = nabc, verbose = true)

## Save data & related parameters

@save "./data/Gausslens/Gausslens_data.jld" h m freqs x_rcv z_rcv x_src z_src src_enc nabc dat


text:
# Protein secondary structure

# SS symbols
#   Helix
#     H: alpha helix
#     G: 3_10 helix
#     I: pi helix
#     T: turn
#   Sheet
#     E: strand
#     B: bridge
#   Loop
#     C: coil
#     S: bend

module SecondaryStructure

using ..Dorothy
using ..Dorothy.Geometry
using ..Dorothy.Hierarchies
using ..Dorothy.PDB
using ..Dorothy.Selectors
using ..Dorothy.Utils

export SSInferenceStrategy, inferss, inferss!, SSByStride

abstract type SSInferenceStrategy end

inferss(model::ParticleCollection) = inferss!(fill("", length(model)), model)

inferss!(model::ParticleCollection) = inferss!(get!(model, :SS, ""), model)

inferss!(model::ParticleCollection, strategy::SSInferenceStrategy) =
		inferss!(get!(model, :SS, ""), model, strategy)

struct SSByStride <: SSInferenceStrategy
	stridepath::String

	SSByStride(; stridepath::AbstractString = "stride") = new(stridepath)
end

inferss!(SS::AbstractVector{<:AbstractString}, model::ParticleCollection) =
		inferss!(SS, model, SSByStride())

function inferss!(SS::AbstractVector{<:AbstractString},
		model::ParticleCollection, strategy::SSByStride)
	@boundscheck length(SS) == length(model) ||
			error("size mismatch between model and output array")
	I = findall(map(Protein, model))
	protein = view(model, I)
	chainids = get(protein, :chainids, Repeated("", length(protein)))
	residues = mcrp(protein).flath2.tree
	stride!(view(SS, I), protein.names, protein.resids, protein.resnames,
			chainids, protein.R, residues, strategy.stridepath)
	SS
end

function stride!(SS::AbstractVector{<:AbstractString},
		names::AbstractVector{<:AbstractString},
		resids::AbstractVector{<:Integer},
		resnames::AbstractVector{<:AbstractString},
		chainids::AbstractVector{<:AbstractString},
		R::AbstractVector{Vector3D},
		residues::AbstractVector{<:AbstractVector{<:Integer}},
		stridepath::AbstractString)
	n = length(SS)
	@boundscheck length(names) == length(resids) == length(resnames) ==
			length(chainids) == length(R) == n ||
			error("size mismatch between property arrays")
	striderecords = Tuple{Int,String,String}[]
	mktemp() do pdbpath, io
		writepdb(io, nothing, nothing, 1:n, names, resids, resnames, chainids,
				R, Repeated(1.0, n), Repeated(0.0, n), Repeated("", n),
				nothing, true, Ref(0))
		close(io)
		prevss = ""
		for line in readlines(open(`$(stridepath) $(pdbpath)`))
			if startswith(line, "ASG")
				tokens = split(line)
				resname = tokens[2]
				resid = parse(Int, tokens[4])
				if tokens[6] == prevss
					ss = prevss
				else
					ss = tokens[6]
				end
				push!(striderecords, (resid, resname, ss))
				prevss = ss
			end
		end
	end
	nstride = length(striderecords)
	nres = length(residues)
	nstride == nres ||
			error("cannot assign $(nstride) stride records to $(nres) residues")
	for (I, (resid, resname, ss)) in zip(residues, striderecords)
		resids[I[1]] == resid ||
				error("stride resid $(resid) does not match $(resids[I[1]])")
		resnames[I[1]] == resname ||
				error("stride resname mismatch at resid $(resid[I[1]])")
		for i in I
			SS[i] = ss
		end
	end
	SS
end

end


text:
plot_dataset(args...; kwargs...) = plot_dataset!(Scene(), args...; kwargs...)
function plot_dataset!(scene, data::Dataset{2}, color = :black; kwargs...)
    makiedata = [Point2f0(d) for d in data]
    scatter!(scene, makiedata; color = color, markersize = 0.01, kwargs...)
    return scene
end
function plot_dataset!(scene, data::Matrix, color = :black; kwargs...)
    size(data, 2) != 2 && error("Matrix dimension 2 must be 2")
    makiedata = [Point2f0(data[i, 1], data[i, 2]) for i in 1:size(data, 1)]
    scatter!(scene, makiedata; color = color, markersize = 0.01, kwargs...)
    return scene
end
function plot_dataset!(scene, data::Dataset{3}, color = :black; kwargs...)
    makiedata = [Point3f0(d) for d in data]
    lines!(scene, makiedata; color = color, transparency = true, linewidth = 2.0, kwargs...)
    return scene
end

plot_datasets(args...; kwargs...) = plot_datasets!(Scene(), args...; kwargs...)
function plot_datasets!(scene, datasets, colors; kwargs...)
    length(datasets) == length(colors) || error("Datasets and colors must "*
    "have equal length")

    for i in eachindex(datasets)
        plot_dataset!(scene, datasets[i], colors[i]; kwargs...)
    end
    return scene
end


text:
#################### Sampler ####################

const samplerfxargs = [(:model, Mamba.Model), (:block, Integer)]

#################### Types and Constructors ####################

type NullFunction end

type SamplingBlock
  model::Model
  index::Int
  transform::Bool

  SamplingBlock(model::Model, index::Integer=0, transform::Bool=false) =
    new(model, index, transform)
end

Sampler(param::Symbol, args...) = Sampler([param], args...)

function Sampler(params::Vector{Symbol}, f::Function, tune::Any=Dict())
  Sampler(params, modelfx(samplerfxargs, f), tune, Symbol[])
end

function SamplerVariate{T<:SamplerTune, U<:Real}(x::AbstractVector{U}, tune::T)
  SamplerVariate{T}(x, tune)
end

function SamplerVariate(block::SamplingBlock, pargs...; kargs...)
  m = block.model
  SamplerVariate(unlist(block), m.samplers[block.index], m.iter, pargs...;
                 kargs...)
end

function SamplerVariate{T<:SamplerTune, U<:Real}(x::AbstractVector{U},
                                                 s::Sampler{T}, iter::Integer,
                                                 pargs...; kargs...)
  if iter == 1
    v = SamplerVariate{T}(x, pargs...; kargs...)
    s.tune = v.tune
  else
    v = SamplerVariate(x, s.tune)
  end
  v
end

#################### Base Methods ####################

function Base.show(io::IO, s::Sampler)
  print(io, "An object of type \"$(summary(s))\"\n")
  print(io, "Sampling Block Nodes:\n")
  show(io, s.params)
  print(io, "\n\n")
  show(io, "text/plain", first(code_typed(s.eval)))
  println(io)
end

function Base.showall(io::IO, s::Sampler)
  show(io, s)
  print(io, "\nTuning Parameters:\n")
  show(io, s.tune)
  print(io, "\n\nTarget Nodes:\n")
  show(io, s.targets)
end

#################### Variate Validators ####################

validate(v::SamplerVariate) = v

function validatebinary(v::SamplerVariate)
  all(insupport(Bernoulli, v)) ||
    throw(ArgumentError("variate is not a binary vector"))
  v
end

function validatesimplex(v::SamplerVariate)
  isprobvec(v) || throw(ArgumentError("variate is not a probability vector"))
  v
end

#################### sample! Generics ####################

function sample!(v::SamplerVariate, density::Nullable; args...)
  isnull(density) && error("must specify a target density in $(typeof(v))",
                           " constructor or sample! method")
  sample!(v, get(density); args...)
end

#################### Simulation Methods ####################

function gradlogpdf!{T<:Real}(block::SamplingBlock, x::AbstractArray{T},
                              dtype::Symbol=:forward)
  gradlogpdf!(block.model, x, block.index, block.transform, dtype=dtype)
end

function logpdf!{T<:Real}(block::SamplingBlock, x::AbstractArray{T})
  logpdf!(block.model, x, block.index, block.transform)
end

function logpdfgrad!{T<:Real}(block::SamplingBlock, x::AbstractVector{T},
                              dtype::Symbol)
  grad = gradlogpdf!(block, x, dtype)
  logf = logpdf!(block, x)
  (logf, ifelse(isfinite(grad), grad, 0.0))
end

function unlist(block::SamplingBlock)
  unlist(block.model, block.index, block.transform)
end

function relist{T<:Real}(block::SamplingBlock, x::AbstractArray{T})
  relist(block.model, x, block.index, block.transform)
end

#################### Auxiliary Functions ####################

asvec(x::Union{Number, Symbol}) = [x]
asvec(x::Vector) = x


text:
" Linear solver that uses PETSc "
type PetscImplicitSolver <: DenseImplicitSolver
    " The System the solve will be done for "
    system::System

    " The 'stiffness matrix' "
    mat::PetscMat

    " The RHS (force vector) "
    rhs::PetscVec

    " The solution vector "
    solution::PetscVec

    " Whether or not this Solver has been initialized "
    initialized::Bool

    " The ghosted solution vector.  This is last because we're going to leave it unitialized in the beginning. "
    ghosted_solution::GhostedPetscVec

    PetscImplicitSolver(system::System) = new(system, PetscMat(), PetscVec("rhs"), PetscVec("solution"), false)
end

" Initializes the matrix and rhs to the correct size "
function initialize!(solver::PetscImplicitSolver)
    @assert solver.system.initialized

    n_dofs = solver.system.n_dofs

    local_n_dofs = solver.system.local_n_dofs

    setSize!(solver.mat, m_local=(Int32)(local_n_dofs), n_local=(Int32)(local_n_dofs))
    generateSparsity!(solver.system)
    setPreallocation!(solver.mat, solver.system.local_dof_sparsity, solver.system.off_proc_dof_sparsity)

    setSize!(solver.rhs, n_local=(Int32)(local_n_dofs))
    setSize!(solver.solution, n_local=(Int32)(local_n_dofs))

    findGhostedDofs!(solver.system)
    solver.ghosted_solution = GhostedPetscVec(Int32[(Int32)(dof) for dof in solver.system.ghosted_dofs], n_local=(Int32)(local_n_dofs))

    solver.initialized = true
end

" Update the ghosted_solution vector from the actual solution "
function updateGhostedSolution!(solver::PetscImplicitSolver)
    copy!(solver.ghosted_solution, solver.solution)
end

"""
    Solve using the built-in dense matrix/vector types

    `assemble` controls whether or not the system will be assembled automatically
"""
function solve!(solver::PetscImplicitSolver; assemble=true)
    startLog(main_perf_log, "solve()")

    if !solver.initialized
        initialize!(solver)
    end

    if assemble
        assembleResidualAndJacobian(solver, solver.system)
    end

    assemble!(solver.solution)

    # Solve the system
    ksp = PetscKSP()
    setOperators(ksp, solver.mat)
    scale!(solver.rhs, -1.0)

    startLog(main_perf_log, "linear_solve")
    solve!(ksp, solver.rhs, solver.solution)
    stopLog(main_perf_log, "linear_solve")

    stopLog(main_perf_log, "solve()")
end


text:
# *****************************************************************************
# Written by Ritchie Lee, ritchie.lee@sv.cmu.edu
# *****************************************************************************
# Copyright ã 2015, United States Government, as represented by the
# Administrator of the National Aeronautics and Space Administration. All
# rights reserved.  The Reinforcement Learning Encounter Simulator (RLES)
# platform is licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License. You
# may obtain a copy of the License at
# http://www.apache.org/licenses/LICENSE-2.0. Unless required by applicable
# law or agreed to in writing, software distributed under the License is
# distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
# KIND, either express or implied. See the License for the specific language
# governing permissions and limitations under the License.
# _____________________________________________________________________________
# Reinforcement Learning Encounter Simulator (RLES) includes the following
# third party software. The SISLES.jl package is licensed under the MIT Expat
# License: Copyright (c) 2014: Youngjun Kim.
# Permission is hereby granted, free of charge, to any person obtaining a copy
# of this software and associated documentation files (the "Software"), to
# deal in the Software without restriction, including without limitation the
# rights to use, copy, modify, merge, publish, distribute, sublicense, and/or
# sell copies of the Software, and to permit persons to whom the Software is
# furnished to do so, subject to the following conditions:
# The above copyright notice and this permission notice shall be included in
# all copies or substantial portions of the Software. THE SOFTWARE IS PROVIDED
# "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR IMPLIED, INCLUDING BUT
# NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, FITNESS FOR A PARTICULAR
# PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR COPYRIGHT
# HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER IN AN
# ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN
# CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE.
# *****************************************************************************

module SaveHelpers

using ..Config_ACASX_GM, ..ConfigAST, ..ConfigMCTS, ..ConfigMCBest
using ..DefineSave
using ..SaveTypes
using RLESUtils, DataFrameUtils, Obj2DataFrames, Loggers, RNGWrapper
using DataFrames
using AdaptiveStressTesting

export set_sim_params!, set_ast_params!, set_mcts_params!, set_compute_info!,
    set_study_params!, set_run_type!, set_q_values!, set_action_seq!, set_result!,
    set_mcbest_params!, set_study_results!
export get_sim_params, get_ast_params, get_mcts_params, get_compute_info, get_result,
    get_action_seq, get_study_params, get_mcbest_params, get_study_results, get_log,
    get_unit, get_run_type, get_mcts_iterations, get_mcbest_n, get_logprobs,
    get_encounter_id, get_q_values
export get_reward, get_num_aircraft, get_vmd, get_hmd, get_md_time, get_nmac, is_nmac,
    nmacs_only, is_differential

set_sim_params!(d::TrajLog, sim_params) = Loggers.set!(d.log, "sim_params", to_df(sim_params))
set_ast_params!(d::TrajLog, ast_params) = Loggers.set!(d.log, "ast_params", to_df(ast_params))
set_mcts_params!(d::TrajLog, mcts_params) = Loggers.set!(d.log, "mcts_params", to_df(mcts_params))
set_mcbest_params!(d::TrajLog, mcbest_params) = Loggers.set!(d.log, "mcbest_params",
    to_df(mcbest_params))
set_compute_info!(d::TrajLog, compute_info) = Loggers.set!(d.log, "compute_info",
    to_df(compute_info))
set_study_params!(d::TrajLog, study_params) = Loggers.set!(d.log, "study_params",
    to_df(study_params))
set_study_results!(d::TrajLog, study_results) = Loggers.set!(d.log, "study_results",
    to_df(study_results))
set_run_type!(d::TrajLog, run_type) = Loggers.set!(d.log, "run_type",
    DataFrame(Dict(:run_type=>[run_type])))
set_q_values!(d::TrajLog, q_vals) = Loggers.set!(d.log, "q_values",
    DataFrame(Dict(:q_values=>q_vals)))

function set_action_seq!(d::TrajLog, action_seq)
    n = length(action_seq[1].rsg)
    df = DataFrame(fill(UInt32, n), 0)
    for a in action_seq
        push!(df, a.rsg.state)
    end
    Loggers.set!(d.log, "action_seq", df)
end

get_sim_params(d::TrajLog) = Obj2DataFrames.set!(defineSimParams(), ObjDataFrame(d["sim_params"]))
get_ast_params(d::TrajLog) = Obj2DataFrames.set!(defineASTParams(), ObjDataFrame(d["ast_params"]))
get_mcts_params(d::TrajLog) = Obj2DataFrames.set!(defineMCTSParams(),
    ObjDataFrame(d["mcts_params"]))
get_mcbest_params(d::TrajLog) = Obj2DataFrames.set!(defineMCBestParams(),
    ObjDataFrame(d["mcbest_params"]))
get_compute_info(d::TrajLog) = Obj2DataFrames.set!(ComputeInfo(),
    ObjDataFrame(d["compute_info"]))
function get_action_seq(d::TrajLog)
    x = ASTAction[]
    df = d["action_seq"]
    n = ncol(df)
    for row in eachrow(df)
        A = squeeze(convert(Array, row), 1)
        push!(x, ASTAction(RSG(A)))
    end
    x
end
get_run_type(d::TrajLog) = d.log["run_type"][1, :run_type]
get_q_values(d::TrajLog) = d.log["q_values"][:q_values]
get_study_params(d::TrajLog) = get_study_params(d, Val{Symbol(get_run_type(d))})
get_study_results(d::TrajLog) = get_study_results(d, Val{Symbol(get_run_type(d))})
function get_result(d::TrajLog)
    reward = get_reward(d)
    action_seq = get_action_seq(d)
    q_values = get_q_values(d)
    StressTestResults(reward, action_seq, q_values)
end

get_mcts_iterations(d::TrajLog) = d.log["mcts_params"][1, :n]
get_mcbest_n(d::TrajLog) = d.log["mcbest_params"][1, :n]

get_reward(d::TrajLog) = d.log["run_info"][1, :reward]
get_num_aircraft(d::TrajLog) = d.log["sim_params"][1, :num_aircraft]
get_vmd(d::TrajLog) = d.log["run_info"][1, :vmd]
get_hmd(d::TrajLog) = d.log["run_info"][1, :hmd]
get_md_time(d::TrajLog) = d.log["run_info"][1, :md_time]

function get_log(d::TrajLog, logname::AbstractString, aircraft_number::Int64)
    d.log["$(logname)_$(aircraft_number)"]
end

function get_unit(d::TrajLog, logname::AbstractString, aircraft_number::Int64, var::Symbol)
    d.log["$(logname)_units_$(aircraft_number)"][1, var]
end

get_encounter_id(d::TrajLog) = d.log["sim_params"][1, :encounter_number]
get_logprobs(d::TrajLog) = convert(Array, d.log["logProb"][:logprob])

is_nmac(file::AbstractString) = file |> trajLoad |> get_nmac
get_nmac(d::TrajLog) = d.log["run_info"][1, :nmac]
nmacs_only(file::AbstractString) = is_nmac(file)
nmacs_only{T<:AbstractString}(files::Vector{T}) = filter(is_nmac, files)

function is_differential(d::TrajLog)
    d.log["sim_params"][1, :libcas2] != "NONE" && d.log["sim_params"][1, :libcas2_config] != "NONE"
end

end #module

max_stars_repo_name: sisl/RLESCAS

max_issues_repo_name: sisl/RLESCAS

max_forks_repo_name: sisl/RLESCAS


text:
# Thermostats

export
    NoThermostat,
    apply_thermostat!,
    AndersenThermostat,
    velocity,
    maxwellboltzmann,
    temperature

"""
    NoThermostat()

Placeholder thermostat that does nothing.
"""
struct NoThermostat <: Thermostat end

"""
    apply_thermostat!(simulation, thermostat)

Apply a thermostat to modify a simulation.
Custom thermostats should implement this function.
"""
function apply_thermostat!(s::Simulation, ::NoThermostat)
    return s
end

"""
    AndersenThermostat(coupling_const)

Rescale random velocities according to the Andersen thermostat.
"""
struct AndersenThermostat{T} <: Thermostat
    coupling_const::T
end

function apply_thermostat!(s::Simulation, thermostat::AndersenThermostat)
    dims = length(first(s.velocities))
    for i in 1:length(s.velocities)
        if rand() < s.timestep / thermostat.coupling_const
            mass = s.atoms[i].mass
            s.velocities[i] = velocity(mass, s.temperature; dims=dims)
        end
    end
    return s
end

"""
    velocity(mass, temperature; dims=3)
    velocity(T, mass, temperature; dims=3)

Generate a random velocity from the Maxwell-Boltzmann distribution.
"""
function velocity(T::Type, mass::Real, temp::Real; dims::Integer=3)
    return SVector([maxwellboltzmann(T, mass, temp) for i in 1:dims]...)
end

function velocity(mass::Real, temp::Real; dims::Integer=3)
    return velocity(DefaultFloat, mass, temp, dims=dims)
end

"""
    maxwellboltzmann(mass, temperature)
    maxwellboltzmann(T, mass, temperature)

Draw from the Maxwell-Boltzmann distribution.
"""
function maxwellboltzmann(T::Type, mass::Real, temp::Real)
    return rand(Normal(zero(T), sqrt(temp / mass)))
end

function maxwellboltzmann(mass::Real, temp::Real)
    return maxwellboltzmann(DefaultFloat, mass, temp)
end

"""
    temperature(simulation)

Calculate the temperature of a system from the kinetic energy of the atoms.
"""
function temperature(s::Simulation)
    ke = sum([a.mass * dot(s.velocities[i], s.velocities[i]) for (i, a) in enumerate(s.atoms)]) / 2
    df = 3 * length(s.coords) - 3
    return 2 * ke / df
end


text:
# Julia wrapper for header: cutensor.h
# Automatically generated using Clang.jl

@checked function cutensorInit(handle)
    initialize_api()
    ccall((:cutensorInit, libcutensor()), cutensorStatus_t,
                   (Ptr{cutensorHandle_t},),
                   handle)
end

@checked function cutensorInitTensorDescriptor(handle, desc, numModes, extent, stride,
                                               dataType, unaryOp)
    initialize_api()
    ccall((:cutensorInitTensorDescriptor, libcutensor()), cutensorStatus_t,
                   (Ptr{cutensorHandle_t}, Ptr{cutensorTensorDescriptor_t}, UInt32,
                    Ptr{Int64}, Ptr{Int64}, cudaDataType_t, cutensorOperator_t),
                   handle, desc, numModes, extent, stride, dataType, unaryOp)
end

@checked function cutensorElementwiseTrinary(handle, alpha, A, descA, modeA, beta, B,
                                             descB, modeB, gamma, C, descC, modeC, D,
                                             descD, modeD, opAB, opABC, typeScalar, stream)
    initialize_api()
    ccall((:cutensorElementwiseTrinary, libcutensor()), cutensorStatus_t,
                   (Ptr{cutensorHandle_t}, Ptr{Cvoid}, PtrOrCuPtr{Cvoid},
                    Ptr{cutensorTensorDescriptor_t}, Ptr{Int32}, Ptr{Cvoid},
                    PtrOrCuPtr{Cvoid}, Ptr{cutensorTensorDescriptor_t}, Ptr{Int32},
                    Ptr{Cvoid}, PtrOrCuPtr{Cvoid}, Ptr{cutensorTensorDescriptor_t},
                    Ptr{Int32}, PtrOrCuPtr{Cvoid}, Ptr{cutensorTensorDescriptor_t},
                    Ptr{Int32}, cutensorOperator_t, cutensorOperator_t, cudaDataType_t,
                    CUstream),
                   handle, alpha, A, descA, modeA, beta, B, descB, modeB, gamma, C, descC,
                   modeC, D, descD, modeD, opAB, opABC, typeScalar, stream)
end

@checked function cutensorElementwiseBinary(handle, alpha, A, descA, modeA, gamma, C,
                                            descC, modeC, D, descD, modeD, opAC,
                                            typeScalar, stream)
    initialize_api()
    ccall((:cutensorElementwiseBinary, libcutensor()), cutensorStatus_t,
                   (Ptr{cutensorHandle_t}, Ptr{Cvoid}, PtrOrCuPtr{Cvoid},
                    Ptr{cutensorTensorDescriptor_t}, Ptr{Int32}, Ptr{Cvoid},
                    PtrOrCuPtr{Cvoid}, Ptr{cutensorTensorDescriptor_t}, Ptr{Int32},
                    PtrOrCuPtr{Cvoid}, Ptr{cutensorTensorDescriptor_t}, Ptr{Int32},
                    cutensorOperator_t, cudaDataType_t, CUstream),
                   handle, alpha, A, descA, modeA, gamma, C, descC, modeC, D, descD, modeD,
                   opAC, typeScalar, stream)
end

@checked function cutensorPermutation(handle, alpha, A, descA, modeA, B, descB, modeB,
                                      typeScalar, stream)
    initialize_api()
    ccall((:cutensorPermutation, libcutensor()), cutensorStatus_t,
                   (Ptr{cutensorHandle_t}, Ptr{Cvoid}, PtrOrCuPtr{Cvoid},
                    Ptr{cutensorTensorDescriptor_t}, Ptr{Int32}, PtrOrCuPtr{Cvoid},
                    Ptr{cutensorTensorDescriptor_t}, Ptr{Int32}, cudaDataType_t, CUstream),
                   handle, alpha, A, descA, modeA, B, descB, modeB, typeScalar, stream)
end

@checked function cutensorInitContractionDescriptor(handle, desc, descA, modeA,
                                                    alignmentRequirementA, descB, modeB,
                                                    alignmentRequirementB, descC, modeC,
                                                    alignmentRequirementC, descD, modeD,
                                                    alignmentRequirementD, computeType)
    initialize_api()
    ccall((:cutensorInitContractionDescriptor, libcutensor()), cutensorStatus_t,
                   (Ptr{cutensorHandle_t}, Ptr{cutensorContractionDescriptor_t},
                    Ptr{cutensorTensorDescriptor_t}, Ptr{Int32}, UInt32,
                    Ptr{cutensorTensorDescriptor_t}, Ptr{Int32}, UInt32,
                    Ptr{cutensorTensorDescriptor_t}, Ptr{Int32}, UInt32,
                    Ptr{cutensorTensorDescriptor_t}, Ptr{Int32}, UInt32,
                    cutensorComputeType_t),
                   handle, desc, descA, modeA, alignmentRequirementA, descB, modeB,
                   alignmentRequirementB, descC, modeC, alignmentRequirementC, descD,
                   modeD, alignmentRequirementD, computeType)
end

@checked function cutensorInitContractionFind(handle, find, algo)
    initialize_api()
    ccall((:cutensorInitContractionFind, libcutensor()), cutensorStatus_t,
                   (Ptr{cutensorHandle_t}, Ptr{cutensorContractionFind_t}, cutensorAlgo_t),
                   handle, find, algo)
end

@checked function cutensorContractionGetWorkspace(handle, desc, find, pref, workspaceSize)
    initialize_api()
    ccall((:cutensorContractionGetWorkspace, libcutensor()), cutensorStatus_t,
                   (Ptr{cutensorHandle_t}, Ptr{cutensorContractionDescriptor_t},
                    Ptr{cutensorContractionFind_t}, cutensorWorksizePreference_t,
                    Ptr{UInt64}),
                   handle, desc, find, pref, workspaceSize)
end

@checked function cutensorInitContractionPlan(handle, plan, desc, find, workspaceSize)
    initialize_api()
    ccall((:cutensorInitContractionPlan, libcutensor()), cutensorStatus_t,
                   (Ptr{cutensorHandle_t}, Ptr{cutensorContractionPlan_t},
                    Ptr{cutensorContractionDescriptor_t}, Ptr{cutensorContractionFind_t},
                    UInt64),
                   handle, plan, desc, find, workspaceSize)
end

@checked function cutensorContraction(handle, plan, alpha, A, B, beta, C, D, workspace,
                                      workspaceSize, stream)
    initialize_api()
    ccall((:cutensorContraction, libcutensor()), cutensorStatus_t,
                   (Ptr{cutensorHandle_t}, Ptr{cutensorContractionPlan_t}, Ptr{Cvoid},
                    PtrOrCuPtr{Cvoid}, PtrOrCuPtr{Cvoid}, Ptr{Cvoid}, PtrOrCuPtr{Cvoid},
                    PtrOrCuPtr{Cvoid}, CuPtr{Cvoid}, UInt64, CUstream),
                   handle, plan, alpha, A, B, beta, C, D, workspace, workspaceSize, stream)
end

@checked function cutensorContractionMaxAlgos(maxNumAlgos)
    initialize_api()
    ccall((:cutensorContractionMaxAlgos, libcutensor()), cutensorStatus_t,
                   (Ptr{Int32},),
                   maxNumAlgos)
end

@checked function cutensorReduction(handle, alpha, A, descA, modeA, beta, C, descC, modeC,
                                    D, descD, modeD, opReduce, minTypeCompute, workspace,
                                    workspaceSize, stream)
    initialize_api()
    ccall((:cutensorReduction, libcutensor()), cutensorStatus_t,
                   (Ptr{cutensorHandle_t}, Ptr{Cvoid}, PtrOrCuPtr{Cvoid},
                    Ptr{cutensorTensorDescriptor_t}, Ptr{Int32}, Ptr{Cvoid},
                    PtrOrCuPtr{Cvoid}, Ptr{cutensorTensorDescriptor_t}, Ptr{Int32},
                    PtrOrCuPtr{Cvoid}, Ptr{cutensorTensorDescriptor_t}, Ptr{Int32},
                    cutensorOperator_t, cutensorComputeType_t, PtrOrCuPtr{Cvoid}, UInt64,
                    CUstream),
                   handle, alpha, A, descA, modeA, beta, C, descC, modeC, D, descD, modeD,
                   opReduce, minTypeCompute, workspace, workspaceSize, stream)
end

@checked function cutensorReductionGetWorkspace(handle, A, descA_, modeA, C, descC_, modeC,
                                                D, descD_, modeD, opReduce, typeCompute,
                                                workspaceSize)
    initialize_api()
    ccall((:cutensorReductionGetWorkspace, libcutensor()), cutensorStatus_t,
                   (Ptr{cutensorHandle_t}, PtrOrCuPtr{Cvoid},
                    Ptr{cutensorTensorDescriptor_t}, Ptr{Int32}, PtrOrCuPtr{Cvoid},
                    Ptr{cutensorTensorDescriptor_t}, Ptr{Int32}, PtrOrCuPtr{Cvoid},
                    Ptr{cutensorTensorDescriptor_t}, Ptr{Int32}, cutensorOperator_t,
                    cutensorComputeType_t, Ptr{UInt64}),
                   handle, A, descA_, modeA, C, descC_, modeC, D, descD_, modeD, opReduce,
                   typeCompute, workspaceSize)
end

@checked function cutensorGetAlignmentRequirement(handle, ptr, desc, alignmentRequirement)
    initialize_api()
    ccall((:cutensorGetAlignmentRequirement, libcutensor()), cutensorStatus_t,
                   (Ptr{cutensorHandle_t}, PtrOrCuPtr{Cvoid},
                    Ptr{cutensorTensorDescriptor_t}, Ptr{UInt32}),
                   handle, ptr, desc, alignmentRequirement)
end

function cutensorGetErrorString(error)
    ccall((:cutensorGetErrorString, libcutensor()), Cstring,
                   (cutensorStatus_t,),
                   error)
end

function cutensorGetVersion()
    ccall((:cutensorGetVersion, libcutensor()), Csize_t, ())
end

function cutensorGetCudartVersion()
    ccall((:cutensorGetCudartVersion, libcutensor()), Csize_t, ())
end

## added in CUTENSOR 1.2

@checked function cutensorHandleWriteCacheToFile(handle, filename)
    initialize_api()
    ccall((:cutensorHandleWriteCacheToFile, libcutensor()), cutensorStatus_t, (Ptr{cutensorHandle_t}, Ptr{UInt8}), handle, filename)
end

@checked function cutensorContractionFindSetAttribute(handle, find, attr, buf, sizeInBytes)
    initialize_api()
    ccall((:cutensorContractionFindSetAttribute, libcutensor()), cutensorStatus_t, (Ptr{cutensorHandle_t}, Ptr{cutensorContractionFind_t}, cutensorContractionFindAttributes_t, Ptr{Cvoid}, Csize_t), handle, find, attr, buf, sizeInBytes)
end

@checked function cutensorHandleReadCacheFromFile(handle, filename, numCachelinesRead)
    initialize_api()
    ccall((:cutensorHandleReadCacheFromFile, libcutensor()), cutensorStatus_t, (Ptr{cutensorHandle_t}, Ptr{UInt8}, Ref{UInt32}), handle, filename, numCachelinesRead)
end

@checked function cutensorContractionDescriptorSetAttribute(handle, desc, attr, buf, sizeInBytes)
    initialize_api()
    ccall((:cutensorContractionDescriptorSetAttribute, libcutensor()), cutensorStatus_t, (Ptr{cutensorHandle_t}, Ptr{cutensorContractionDescriptor_t}, cutensorContractionDescriptorAttributes_t, Ptr{Cvoid}, Csize_t), handle, desc, attr, buf, sizeInBytes)
end

@checked function cutensorHandleDetachPlanCachelines(handle)
    initialize_api()
    ccall((:cutensorHandleDetachPlanCachelines, libcutensor()), cutensorStatus_t, (Ptr{cutensorHandle_t},), handle)
end

@checked function cutensorHandleAttachPlanCachelines(handle, cachelines, numCachelines)
    initialize_api()
    ccall((:cutensorHandleAttachPlanCachelines, libcutensor()), cutensorStatus_t, (Ptr{cutensorHandle_t}, Ptr{cutensorPlanCacheline_t}, UInt32), handle, cachelines, numCachelines)
end


text:
# Constants for the agent
const BATCH_SIZE = 1000
const LR = 1e-3

# AbstractAgent interface
abstract type AbstractAgent end

function get_action end
function update! end

Base.@kwdef mutable struct SnakeAgent <: AbstractAgent
    n_games::Int = 0
    record::Int = 0
    ϵ::Int = 0
    memory::AgentMemory = CircularBufferMemory()
    model::Flux.Chain = SnakeAI.linear_QNet(11, 256, 3)
    opt::Flux.Optimise.AbstractOptimiser = Flux.ADAM(LR)
    criterion::Function = Flux.Losses.mse
end

function Base.show(io::IO, agent::SnakeAgent)
    println("n_games => ", agent.n_games)
    println("record => ", agent.record)
    println("ϵ => ", agent.ϵ)
    println("memory => ", typeof(agent.memory.data))
    println("model => ", agent.model)
    println("optimizer => ", typeof(agent.opt))
    println("criterion => ", String(Symbol(agent.criterion)))
end


text:
img, imgw = Float64.(Gray.(img)), Float64.(Gray.(imgw))

viewimg(img)
viewimg(imgw)

fhs, window = 15, [31, 31]
u_est = single_lap(img, imgw, fhs, window);

showflow(u_est)
showflow(flow)

# # show the NaN location
# nan_img = map(x -> isnan(x) ? 1.0 : 0.0, u_est);
#
# viewimg(nan_img)

function error_percentage(estim, truth)
    error = abs(estim - truth)
    return error/truth
end

function abs_error(estim, truth)
    return abs(estim - truth)
end

# good_estimation = map(x -> LAP_julia.vec_len(error_percentage(x...)), zip(u_est, flow.*(-1)))
# maximum(filter(!isnan, good_estimation))
# minimum(filter(!isnan, good_estimation))
#
# good_estimation_r = maximum(filter(!isnan, good_estimation)) .- good_estimation
#
#
# maximum(filter(!isnan, good_estimation_r))
# minimum(filter(!isnan, good_estimation_r))
#
# figure()
#
# imgshow(img)
# imshow(good_estimation, cmap = :RdYlGn_r, alpha = 0.3); gcf()
# colorbar(ax=gca()); gcf()

nccs=zeros(100)

for k in 1:100
    img, imgw, flow = gen_init(:lena, :uniform, flow_args=[rand(Complex{Float64}), 5]);
    grad, imge = LAP_julia.gradient_magnitude(img);
    flow_est = single_lap(img, imgw, 10, [21,21]);
    percentage_errors = map(x -> LAP_julia.vec_len(error_percentage(x...)), zip(flow_est, flow.*(-1)));
    nccs[k] = ncc(imge[.!isnan.(flow_est)], percentage_errors[.!isnan.(flow_est)])
end

figure()
plt.hist(nccs, 20); gcf()

mean(nccs)

## with mri data
nccs=zeros(100)

for k in 1:100
    img, imgw, flow = gen_init(:lena, :uniform, flow_args=[rand(Complex{Float64}), 5]);
    grad, imge = LAP_julia.gradient_magnitude(img);
    flow_est = single_lap(img, imgw, 10, [21,21]);
    percentage_errors = map(x -> LAP_julia.vec_len(error_percentage(x...)), zip(flow_est, flow.*(-1)));
    nccs[k] = ncc(imge[.!isnan.(flow_est)], percentage_errors[.!isnan.(flow_est)])
end

figure()
plt.hist(nccs, 20); gcf()
base_path = "/Users/MrTrololord/Downloads/mri_head_homography/"
target_images_path = joinpath(base_path, "images", "target")

function gen_mri(path_folder, flow_function)
    img = load(joinpath(path_folder, "img$(string(rand((collect(1:60))))).png"))
    flow = flow_function(size(img))
    imgw = warp_img(img, -real.(flow), -imag.(flow))
    return img, imgw, flow
end

flow_function(img_size) = gen_uniform_flow(img_size, rand(Complex{Float64}), 5)

img, imgw, flow = gen_mri(target_images_path, flow_function)

# test on mri
nccs=zeros(100)

for k in 1:100
    img, imgw, flow = gen_mri(target_images_path, flow_function);
    grad, imge = LAP_julia.gradient_magnitude(imgw);
    img, imgw = Float64.(Gray.(img)), Float64.(Gray.(imgw))
    flow_est = single_lap(img, imgw, 10, [21,21]);
    percentage_errors = map(x -> LAP_julia.vec_len(error_percentage(x...)), zip(flow_est, flow.*(-1)));
    nccs[k] = ncc(imge[.!isnan.(flow_est)], percentage_errors[.!isnan.(flow_est)])
end

fig, ax = subplots(dpi = 400)
plt.hist(nccs, 20);
title("NCC of Displacement Error and Magnitude of Gradient");
savefig("../plots/ncc_displacement_error_and_grad.png")

gcf()
mean(nccs)

# abs

abs_err = map(x -> LAP_julia.vec_len(abs_error(x...)), zip(u_est, flow.*(-1)))
maximum(filter(!isnan, abs_err))
minimum(filter(!isnan, abs_err))

imgshow(img)
imshow(abs_err, cmap = :RdYlGn_r, alpha = 0.3); gcf()
colorbar(ax=gca()); gcf()

vec, edge = LAP_julia.gradient_magnitude(img)

maximum(filter(!isnan, 1 ./good_estimation))
minimum(filter(!isnan, 1 ./good_estimation))

filter(!isnan, normalize_to_zero_one(1 ./good_estimation))

minimum(filter(!isnan, normalize_to_zero_one(1 ./good_estimation)))

using LAP_julia: normalize_to_zero_one

imgshow(map(x-> isnan(x) ? 0 : x, normalize_to_zero_one(1 ./good_estimation)))

collect(colorview(RGB, normalize_to_zero_one(edge), normalize_to_zero_one(edge), normalize_to_zero_one(1 ./good_estimation)))

collect(colorview(Gray, img.*(1/extrema(img)[2]))) .+ collect(colorview(RGBA, good_estimation, zeros(size(good_estimation)), zeros(size(good_estimation)), fill(0.3, size(good_estimation)...)))

imgoverlay_v2(img, good_estimation)

imgshow(ones(112,112))


text:
using Plots

include("diffusion_2D_expl.jl")
include("diffusion_2D_damp.jl")

save_fig=false

@views function compare_expl_impl(; save_fig=false)
    # run the codes
    do_visu = false
    xc, yc, H_expl = diffusion_2D_expl(; do_visu=do_visu)
    xc, yc, H_damp = diffusion_2D_damp(; do_visu=do_visu)
    # extract ∆
    ∆_expl_impl = 100.0*(H_expl .- H_damp)./H_damp
    ∆max = maximum(abs.(∆_expl_impl))
    # visualise
    fontsize = 12
    opts = (aspect_ratio=1, yaxis=font(fontsize, "Courier"), xaxis=font(fontsize, "Courier"),
            ticks=nothing, framestyle=:box, titlefontsize=fontsize, titlefont="Courier", colorbar_title="",
            xlabel="Lx", ylabel="Ly", xlims=(xc[1],xc[end]), ylims=(yc[1],yc[end]), clims=(-∆max, ∆max))
    display(heatmap(xc, yc, ∆_expl_impl; seriescolor=:balance, title="∆(expl vs impl) in %", opts...))
    if save_fig savefig("diff2D_expl_impl.png") end
    return
end

compare_expl_impl(; save_fig=save_fig)


text:
module Evolution

using QXZoo.Circuit
using QXZoo.GateOps
using QXZoo.DefaultGates

"""
    apply_pauli_string(θ::Number, pauli_gates::Vector{GateCall1})

Create a circuit representating a given exp Pauli string (eg exp(-i0.5*X*X*Y*Y))

# Examples
```julia-repl
```
"""
function apply_pauli_string!(cct::Circuit.Circ, θ::Number, pauli_gates::Vector{GateCall1})
    coupling_map = [i.target for i in pauli_gates]
    gates = [i.gate_symbol for i in pauli_gates]
    @assert issubset(Set(gates), Set([DefaultGates.GateSymbols.x, DefaultGates.GateSymbols.y, DefaultGates.GateSymbols.z]))

    u_t = []
    for g in pauli_gates
        if g.gate_symbol === DefaultGates.GateSymbols.x
            push!(u_t, DefaultGates.h(g.target))
        elseif g.gate_symbol === DefaultGates.GateSymbols.y
            push!(u_t, DefaultGates.r_x(g.target, pi/2))
        end
    end

    #label = GateSymbol(Symbol("exp_"*string([i.label for i in gates]...)))

    for i in u_t
        cct << i
    end

    for i in 1:length(coupling_map)-1
        cct << DefaultGates.c_x(coupling_map[i], coupling_map[i+1])
    end

    cct << DefaultGates.r_z(coupling_map[end], θ)

    for i in reverse(1:length(coupling_map)-1)
        cct << DefaultGates.c_x(coupling_map[i], coupling_map[i+1])
    end

    for i in reverse(u_t)
        cct << adjoint(i)
    end

end

end