# Many derivative free methods of different orders
#
# TODO: rework Order5 #https://pdfs.semanticscholar.org/ce50/3210d96f653a14b28da96600d5990d2abe97.pdf
# https://content.sciendo.com/view/journals/tmj/10/4/article-p103.xml 7 and 8
# order8: https://www.hindawi.com/journals/ijmms/2012/493456/ref/

##################################################
## Guard against non-robust algorithms
##
## By default, we do this by deciding if we
## should take a secant step instead of the algorithm For example, for
## Steffensen which is quadratically convergent and the Secant method
## which is only superlinear,
## the error, e_{n+1} = x_{n+1} - alpha, may be smaller after a secant
## step than a Steffensen step. (It is only once x_n is close enough
## to alpha that the method is quadratically convergent.
## The Steffensen error is
## Δn+1 = f[x,x+fx, alpha]/f[x, x+fx] * (1 - f[x, alpha]) (x-alpha)^2
##      ≈ f''/(2f') * ( 1 + f') Δn^2
## The Secant error is
## Δn+1 = f[x,x_{-1},alpha] / f[x,x_{-1}] * (x-alpha) * (x_{-1} - alpha)
##      ≈  f''/(2f')  Δn ⋅ Δn-1
## The ratio is ≈ (1 + f')(Δn / Δn-1)
## It seems reasonable, that a Steffensen step is preferred when
## the ratio satisfies -1 < (1+f') ⋅ Δn /Δn-1 < 1
## We could use f' ~ fp = (fx1-fx0)/(x1-x0); but our proxy for
## Δn/Δn-1 is problematic, as we don't know alpha, and using xn-x_{n-1}
## can be an issue when only x1 and not x0 is specified. This needs
## working around.
##
## Instead, as Steffensen is related to Newton as much as
## (f(x+fx) - fx)/fx  ≈ f'(x), we take a Steffensen step if |fx|
## is small enough. For this we use |fx| <= x/1000; which
## seems to work reasonably well over several different test cases.
@inline function do_guarded_step(M::AbstractSecant, o::AbstractUnivariateZeroState{T,S}) where {T, S}
    x, fx = o.xn1, o.fxn1
    1000 * abs(fx) >  max(oneunit(S), abs(x) * oneunit(S) /oneunit(T)) * one(T)
end

# check if we should guard against step for method M; call N if yes, P if not
function update_state_guarded(M::AbstractSecant,N::AbstractUnivariateZeroMethod, P::AbstractUnivariateZeroMethod, fs, o, options, l=NullTracks())
    if do_guarded_step(M, o)
        return update_state(N, fs, o, options, l)
    else
        update_state(P, fs, o, options, l)
    end
end

##################################################
initial_fncalls(::AbstractSecant) = 2

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

## Secant
## https://en.wikipedia.org/wiki/Secant_method
"""
    Secant()
    Order1()
    Orderφ()

The `Order1()` method is an alias for `Secant`. It specifies the
[secant method](https://en.wikipedia.org/wiki/Secant_method).
This method keeps two values in its state, `xₙ` and `xₙ₋₁`. The
updated point is the intersection point of ``x`` axis with the secant line
formed from the two points. The secant method uses ``1`` function
evaluation per step and has order `φ≈ (1+sqrt(5))/2`.

The error, `eᵢ = xᵢ - α`, satisfies
`eᵢ₊₂ = f[xᵢ₊₁,xᵢ,α] / f[xᵢ₊₁,xᵢ] * (xᵢ₊₁-α) * (xᵢ - α)`.

"""
struct Secant <: AbstractSecant end
const Order1 = Secant
const Orderφ = Secant

# init_state(M,F,x) --> call init_state(M,F,x₀,x₁,fx₀, fx₁)
function init_state(M::AbstractSecant, F::Callable_Function, x)
    x₀,x₁ = x₀x₁(x)
    fx₀, fx₁ = first(F(x₀)), first(F(x₁))
    state = init_state(M, F, x₀, x₁, fx₀, fx₁)
end

# initialize from xs, fxs
function init_state(::AbstractSecant, F, x₀, x₁, fx₀, fx₁)
    UnivariateZeroState(x₁, x₀, fx₁, fx₀)
end

function update_state(::Order1, F, o::AbstractUnivariateZeroState, options, l=NullTracks())
    xn0, xn1 = o.xn0, o.xn1
    fxn0, fxn1 = o.fxn0, o.fxn1
    Δ = fxn1 * (xn1 - xn0) / (fxn1 - fxn0)

    if isissue(Δ)
        log_message(l, "Increment `Δx` has issues. ")
        return o, true
     end

    x0,x1 = xn1, xn1 - Δ
    fx0, fx1 = fxn1, F(x1)
    incfn(l)

    @set! o.xn0 = x0
    @set! o.xn1 = x1
    @set! o.fxn0 = fx0
    @set! o.fxn1 = fx1

    return o, false

end

### Steffensen
## https://en.wikipedia.org/wiki/Steffensen's_method#Simple_description

"""
    Steffensen()
    Order2()

The quadratically converging
[Steffensen](https://en.wikipedia.org/wiki/Steffensen's_method#Simple_description)
method is used for the derivative-free `Order2()` algorithm. Unlike
the quadratically converging Newton's method, no derivative is
necessary, though like Newton's method, two function calls per step
are. Steffensen's algorithm is more sensitive than Newton's method to
poor initial guesses when `f(x)` is large, due to how `f'(x)` is
approximated. The `Order2` method replaces a Steffensen step with a secant
step when `f(x)` is large.

The error, `eᵢ - α`, satisfies
`eᵢ₊₁ = f[xᵢ, xᵢ+fᵢ, α] / f[xᵢ,xᵢ+fᵢ] ⋅ (1 - f[xᵢ,α] ⋅ eᵢ²`
"""
struct Steffensen <: AbstractSecant end
struct Order2 <: AbstractSecant end

function update_state(method::Order2, fs, o::UnivariateZeroState{T,S}, options, l=NullTracks())  where {T, S}
     update_state_guarded(method, Secant(), Steffensen(), fs, o, options, l)
end

function update_state(method::Steffensen, F, o::AbstractUnivariateZeroState{T,S}, options, l=NullTracks())  where {T, S}

    x0, x1 = o.xn0, o.xn1
    fx0, fx1 = o.fxn0, o.fxn1

    sgn = sign((fx1 - fx0) / (x1 - x0))
    x2 = x1 - sgn * fx1 / oneunit(S) * oneunit(T)

    f0 = fx1
    f1::S = F(x2)
    incfn(l, 1)

    delta = -sgn * f0 * f0 / (f1 - f0) * oneunit(T) / oneunit(S)

    if isissue(delta)
        log_message(l, "Increment `Δx` has issues. ")
        return o, true
    end

    x0, x1 = x1, x1 - delta
    fx0, fx1 = fx1, F(x1)
    incfn(l)

    @set! o.xn0 = x0
    @set! o.xn1 = x1
    @set! o.fxn0 = fx0
    @set! o.fxn1 = fx1

    return o, false
end

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

"""
    Order5()
    KumarSinghAkanksha()

Implements an order 5 algorithm from *A New Fifth Order Derivative
Free Newton-Type Method for Solving Nonlinear Equations* by Manoj
Kumar, Akhilesh Kumar Singh, and Akanksha, Appl. Math. Inf. Sci. 9,
No. 3, 1507-1513 (2015), DOI: [10.12785/amis/090346](https://doi.org/10.12785/amis/090346). Four function
calls per step are needed.  The `Order5` method replaces a Steffensen step with a secant
step when `f(x)` is large.

The error, `eᵢ = xᵢ - α`, satisfies
`eᵢ₊₁ = K₁ ⋅ K₅ ⋅ M ⋅ eᵢ⁵ + O(eᵢ⁶)`

"""
struct Order5 <: AbstractSecant end
struct KumarSinghAkanksha <: AbstractSecant end

function update_state(method::Order5, fs, o::UnivariateZeroState, options, l=NullTracks())
    update_state_guarded(method, Secant(), KumarSinghAkanksha(), fs, o, options, l)
end

function update_state(M::KumarSinghAkanksha, F, o::AbstractUnivariateZeroState{T,S},
                      options, l=NullTracks()) where {T, S}

    xn = o.xn1
    fxn = o.fxn1

    wn::T = steff_step(M, o.xn1, o.fxn1)

    fwn::S = F(wn)
    incfn(l)

    fp, issue = _fbracket(o.xn1, wn, o.fxn1, fwn)
    if issue
        log_message(l, "Issue with divided difference f[xn, wn]. ")
        @set! o.xn0 = o.xn1
        @set! o.xn1 = wn
        @set! o.fxn0 = o.fxn1
        @set! o.fxn1 = fwn

        return o, true
    end

    yn::T = o.xn1 - o.fxn1 / fp
    fyn::S = F(yn)
    incfn(l)

    zn::T = xn - (fxn + fyn) / fp
    fzn::S = F(zn)
    incfn(l)

    fp, issue = _fbracket_ratio(yn, o.xn1, wn, fyn, o.fxn1, fwn)
    if issue
        log_message(l, "Issue with f[xn,yn]*f[yn,wn] / f[xn, wn]")
        @set! o.xn0 = o.xn1
        @set! o.xn1 = yn
        @set! o.fxn0 = o.fxn1
        @set! o.fxn1 = fyn

        return o, true

    end

    @set! o.xn0 = o.xn1
    @set! o.fxn0 = o.fxn1
    @set! o.xn1 = zn  - fzn  / fp
    @set! o.fxn1 = F(o.xn1)
    incfn(l)

    return o, false

#    nothing
end

struct Order5Derivative <: AbstractSecant end
fn_argout(::Order5Derivative) = 2
function update_state(method::Order5Derivative, f,
                      o::AbstractUnivariateZeroState{T,S}, options, l=NullTracks())  where {T, S}

    xn, fxn = o.xn1, o.fxn1
    a::T, b::S = f(xn)
    fpxn = a/b
    incfn(l)

    if isissue(fpxn)
        return o, true
    end

    yn::T = xn - fxn / fpxn
    fyn::S, Δyn::T = f(yn)
    fpyn = fyn / Δyn
    incfn(l, 2)

    if isissue(fpyn)
        log_message(l, "Issue computing `fpyn`")
        @set! o.xn0 = o.xn1
        @set! o.xn1 = yn
        @set! o.fxn0 = o.fxn1
        @set! o.fxn1 = fyn

        return o, true
    end

    zn::T = xn  - (fxn + fyn) / fpxn
    fzn::S, _ = f(zn)
    incfn(l, 2)

    xn1 = zn - fzn / fpyn
    fxn1,_ = f(xn1)
    incfn(l, 2)

    @set! o.xn0 = xn
    @set! o.xn1 = xn1
    @set! o.fxn0 = fxn
    @set! o.fxn1 = fxn1

    return o

end

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

## cf also: https://doi.org/10.1515/tmj-2017-0049
"""
    Order8()
    Thukral8()

Implements an eighth-order algorithm from *New Eighth-Order
Derivative-Free Methods for Solving Nonlinear Equations* by Rajinder
Thukral, International Journal of Mathematics and Mathematical
Sciences Volume 2012 (2012), Article ID 493456, 12 pages DOI:
[10.1155/2012/493456](https://doi.org/10.1155/2012/493456). Four
function calls per step are required.  The `Order8` method replaces a
Steffensen step with a secant step when `f(x)` is large.

The error, `eᵢ = xᵢ - α`, is expressed as `eᵢ₊₁ = K ⋅ eᵢ⁸` in
(2.25) of the paper for an explicit `K`.

"""
struct Order8 <: AbstractSecant end
struct Thukral8 <: AbstractSecant end

function update_state(method::Order8, fs, o::UnivariateZeroState{T,S}, options, l=NullTracks())  where {T, S}
    update_state_guarded(method, Secant(), Thukral8(), fs, o, options, l)
end

function update_state(M::Thukral8, F, o::AbstractUnivariateZeroState{T,S}, options, l=NullTracks()) where {T, S}

    xn = o.xn1
    fxn = o.fxn1

    wn::T = steff_step(M, xn, fxn)
    fwn::S = F(wn)
    incfn(l)

    if isissue(fwn)
        log_message(l, "issue with Steffensen step fwn")
        @set! o.xn0 = xn
        @set! o.xn1 = wn
        @set! o.fxn0 = fxn
        @set! o.fxn1 = fwn

        return o, true
    end

    fp, issue = _fbracket(xn, wn, fxn, fwn)

    if issue
        log_message(l, "issue with divided difference f[xn, wn]. ")
        return o, true
    end

    yn::T = xn - fxn / fp
    fyn::S = F(yn)
    incfn(l)

    fp, issue = _fbracket(yn, xn, fyn, fxn)
    if issue #fp
        log_message(l, "issue with divided difference f[xn, yn]. ")
        @set! o.xn0 = xn
        @set! o.xn1 = yn
        @set! o.fxn0 = fxn
        @set! o.fxn1 = fyn

        return o, true
    end

    phi = (1 + fyn / fwn)           # pick one of options
    zn =  yn - phi * fyn / fp
    fzn::S = F(zn)
    incfn(l)

    fp, issue =  _fbracket_diff(xn, yn, zn, fxn, fyn, fzn)
    if issue
        log_message(l, "issue with divided difference  f[y,z] - f[x,y] + f[x,z]. ")
        @set! o.xn0 = xn
        @set! o.xn1 = zn
        @set! o.fxn0 = fxn
        @set! o.fxn1 = fzn

        return o, true
    end

    w = 1 / (1 - fzn/fwn)

    xi = (1 - 2fyn*fyn*fyn / (fwn * fwn * fxn))

    xn1 = zn - w * xi * fzn / fp
    fxn1::S = F(xn1)
    incfn(l)

    @set! o.xn0 = xn
    @set! o.xn1 = xn1
    @set! o.fxn0 = fxn
    @set! o.fxn1 = fxn1

    return o, false

end

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

"""
    Order16()
    Thukral16()

Implements the order 16 algorithm from
*New Sixteenth-Order Derivative-Free Methods for Solving Nonlinear Equations*
by R. Thukral,
American Journal of Computational and Applied Mathematics
p-ISSN: 2165-8935;    e-ISSN: 2165-8943; 2012;  2(3): 112-118
DOI: [10.5923/j.ajcam.20120203.08](https://doi.org/10.5923/j.ajcam.20120203.08).

Five function calls per step are required. Though rapidly converging,
this method generally isn't faster (fewer function calls/steps) over
other methods when using `Float64` values, but may be useful for
solving over `BigFloat`.  The `Order16` method replaces a Steffensen step with a secant
step when `f(x)` is large.

The error, `eᵢ = xᵢ - α`, is expressed as `eᵢ₊₁ = K⋅eᵢ¹⁶` for an explicit `K`
in equation (50) of the paper.

"""
struct Order16 <: AbstractSecant end
struct Thukral16 <: AbstractSecant end

function update_state(method::Order16, fs, o::UnivariateZeroState{T,S}, options, l=NullTracks())  where {T, S}
    update_state_guarded(method, Secant(), Thukral16(), fs, o, options, l)
end

function update_state(M::Thukral16, F, o::AbstractUnivariateZeroState{T,S}, options, l=NullTracks()) where {T, S}
    xn = o.xn1
    fxn = o.fxn1

    wn::T = steff_step(M, xn, fxn)
    fwn::S = F(wn)
    incfn(l)

    fp, issue = _fbracket(xn, wn, fxn, fwn)

    if issue
        log_message(l, "issue with f[xn,wn]")
        @set! o.xn0 = o.xn1
        @set! o.xn1 = wn
        @set! o.fxn0 = o.fxn1
        @set! o.fxn1 = fwn
        return o, true
    end

    yn::T = xn - fxn / fp
    fyn::S = F(yn)
    incfn(l)

    fp, issue = _fbracket_ratio(yn, xn, wn, fyn, fxn, fwn)
    if issue
        log_message(l, "issue with f[xn,yn]*f[yn,wn]/f[xn,wn]")
        @set! o.xn0 = o.xn1
        @set! o.xn1 = yn
        @set! o.fxn0 = o.fxn1
        @set! o.fxn1 = fyn

        return o, true
    end

    zn = yn - fyn / fp
    fzn::S = F(zn)
    incfn(l)

    fp, issue = _fbracket_diff(xn, yn, zn, fxn, fyn, fzn)
    u2, u3, u4 = fzn/fwn, fyn/fxn, fyn/fwn

    eta = 1 / (1 + 2*u3*u4^2) / (1 - u2)
    if issue
        log_message(l, "Approximate derivative failed")
        @set! o.xn0 = o.xn1
        @set! o.xn1 = zn
        @set! o.fxn0 = o.fxn1
        @set! o.fxn1 = fzn

        return o, true
    end

    an = zn - eta * fzn / fp
    fan::S = F(an)
    incfn(l)

    fp, issue = _fbracket_ratio(an, yn, zn, fan, fyn, fzn)
    if issue
        log_message(l, "Approximate derivative failed")
        @set! o.xn0 = o.xn1
        @set! o.xn1 = an
        @set! o.fxn0 = o.fxn1
        @set! o.fxn1 = fan

        return o, true
    end

    u1, u5, u6 = fzn/fxn, fan/fxn, fan/fwn

    fp1, issue = _fbracket(xn,yn, fxn, fyn)

    sigma =  1 + u1*u2 - u1*u3*u4^2 + u5 + u6 + u1^2*u4 +
    u2^2*u3 + 3*u1*u4^2*(u3^2 - u4^2)/(fp1/oneunit(fp1))

    xn1 = an - sigma * fan / fp
    fxn1::S = F(xn1)
    incfn(l)

    @set! o.xn0 = xn
    @set! o.xn1 = xn1
    @set! o.fxn0 = fxn
    @set! o.fxn1 = fxn1

    return o, false
end

##################################################
## some means of guarding against large fx when taking a secant step
## TODO: rework this
function steff_step(M::Union{Order5, Order8, Order16}, x::S, fx::T) where {S, T}

    xbar, fxbar = real(x/oneunit(x)), fx/oneunit(fx)
    thresh =  max(1, abs(xbar)) * sqrt(eps(one(xbar))) #^(1/2) # max(1, sqrt(abs(x/fx))) * 1e-6

    out = abs(fxbar) <= thresh ? fxbar  : sign(fx) * thresh
    x + out * oneunit(x)

end


abstract type AbstractPriorExample end

function getproperty(prior_ex::AbstractPriorExample, name::Symbol)
    if name == :distribution
       _distribution(prior_ex)
    end
end

@recipe function f(prior_ex::AbstractPriorExample)
    dbn = prior_ex.distribution
    if isa(dbn, MixtureModel)
        components --> false
    end
   xlab --> L"\mu"
   label --> ""
   title --> _pretty_name(prior_ex)
   dbn
end

"""
    AshSpiky, AshNearnormal, AshFlattop, AshSkew, AshBignormal, AshBimodal

Empirical Bayes priors that are used in the simulations of:

> Stephens, M., 2017. False discovery rates: a new deal. Biostatistics, 18(2), pp.275-294.
"""
AshSpiky, AshNearNormal, AshFlatTop, AshSkew, AshBignormal, AshBimodal

abstract type AshPrior <: AbstractPriorExample end

struct AshSpiky <: AshPrior end
_distribution(::AshSpiky) =  MixtureModel([Normal(0,0.25); Normal(0,0.5); Normal(0,1); Normal(0,2)] ,
                                          [0.4; 0.2; 0.2; 0.2 ])
_pretty_name(::AshSpiky) = "Spiky"

struct AshNearNormal <: AshPrior end
_distribution(::AshNearNormal) = MixtureModel([Normal(0,1); Normal(0,2)], [1/3; 2/3])
_pretty_name(::AshNearNormal) = "Near Normal"

struct AshFlatTop <: AshPrior end
_distribution(::AshFlatTop) = MixtureModel([Normal(x,0.5) for x=-1.5:0.5:1.5])
_pretty_name(::AshFlatTop) = "Flat Top"

struct AshSkew <: AshPrior end
_distribution(::AshSkew) =  MixtureModel([Normal(-2,2), Normal(-1,1.5), Normal(0,1), Normal(1,1)],
                                            [1/4; 1/4; 1/3; 1/6])
_pretty_name(::AshSkew) = "Skew"

struct AshBigNormal <: AshPrior end
_distribution(::AshBigNormal) =  Normal(0,4)
_pretty_name(::AshBigNormal) = "Skew"

struct AshBimodal <: AshPrior end
_distribution(::AshBimodal) =  MixtureModel( [Normal(-2,1); Normal(2,1)],
                            [0.5;0.5] )
_pretty_name(::AshBimodal) = "AshBimodal"

"""
    IWUnimod, IWBimod

Empirical Bayes priors that are used in the simulations of:

> Ignatiadis, N. and Wager, S., 2019. Bias-aware confidence intervals for empirical Bayes analysis. arXiv preprint arXiv:1902.02774.
"""
IWUnimod, IWBimod

abstract type IWPrior <: AbstractPriorExample end

struct IWUnimod <: IWPrior end
_distribution(::IWUnimod) =  MixtureModel([ Normal(-0.2,.2), Normal(0,0.9)],[0.7, 0.3])
_pretty_name(::IWUnimod) = "Unimodal"

struct IWBimod <: IWPrior end
_distribution(::IWBimod) =  MixtureModel([Normal(-1.5,.2), Normal(1.5, .2)])
_pretty_name(::IWBimod) = "Bimodal"

"""
    EfronUnifNormal

Empirical Bayes priors that are used in the simulations of:

> Efron, B., 2016. Empirical Bayes deconvolution estimates. Biometrika, 103(1), pp.1-20.
"""
struct EfronUnifNormal <: AbstractPriorExample end
_distribution(::EfronUnifNormal) =  MixtureModel([Uniform(-3.0,3.0), Normal(0.0, 0.5)], [1/8, 7/8])
_pretty_name(::EfronUnifNormal) = "Uniform / Normal-Mixture"


function partition!(data::Array, p::Int64, r::Int64)
	i::Int64 = rand(p:r)
	data[i], data[r] = data[r], data[i]
	x = data[r]
	i = p - 1
	for j in p : r - 1
		if data[j] <= x
			i += 1
			data[i], data[j] = data[j], data[i]
		end
	end
	data[i + 1], data[r] = data[r], data[i + 1]
	return i + 1
end

function quicksort!(data::Array, start::Int64, stop::Int64)
	if stop - start < 1
		return
	end

	p = partition!(data, start, stop)

	quicksort!(data, start, p - 1)
	quicksort!(data, p + 1, stop)
end

function quicksort!(data::Array)
	quicksort(data, 1, length(data))
end


@testset "K-Means" begin

    Random.seed!(42)

    @testset "Internal functions" begin
        # data inputs
        a = Float64[1, 3, 5, 7, 11, 13]
        b = Float64[2, 4, 6, 8, 12, 14]
        A = Float64[1 3 5; 7 11 13]
        B = Float64[1 3; 5 7; 11 13]

        # euclidean
        @test ClusterAnalysis.euclidean(a, b) == 2.449489742783178

        # squared_error in vector and matrix
        @test ClusterAnalysis.squared_error(a) == 107.33333333333334
        @test ClusterAnalysis.squared_error(A[1, :]) == 8.0

        # totalwithinss
        cluster = [1, 1, 2]
        @test ClusterAnalysis.totalwithinss(B, 2, cluster) == 16.0
    end

    @testset "Kmeans Results" begin
        A = rand(100, 4)
        k = 4

        @testset "test dimensions results" begin
            res1 = kmeans(A, k, init=:kmpp)
            @test length(res1.centroids) == k
            @test length(res1.cluster) == size(A, 1)
            @test length(res1.centroids[1]) == size(A, 2)
        end

        @testset "diferent types of float" begin
            B = rand(Float32, 100, 4)
            res2 = kmeans(B, k, init=:random)
            @test eltype(B) == eltype(res2.centroids[1])
        end

        @testset "convert Int to float" begin
            C = rand(Int64, 100, 4)
            res3 = kmeans(C, k)
            @test Float64 == eltype(res3.centroids[1])
        end

        @testset "test errors" begin
            @test_throws ArgumentError kmeans(A, k, init=:wrongargument)
            @test_throws ArgumentError kmeans([1, 4, 5, 6, 7], k)
            @test_throws MethodError kmeans(A, 1.2)
        end
    end
end


@doc raw"""
# Harmonic Oscillator

"""
module HarmonicOscillator

    using GeometricBase
    using GeometricEquations
    using Parameters

    export harmonic_oscillator_ode, harmonic_oscillator_iode, harmonic_oscillator_pode, harmonic_oscillator_hode, harmonic_oscillator_sode,
           harmonic_oscillator_dae, harmonic_oscillator_idae, harmonic_oscillator_pdae, harmonic_oscillator_hdae

    export hamiltonian, compute_energy_error

    const Δt = 0.1
    const nt = 10

    const k = 0.5
    const ω = √k
    const p = (k=k, ω=ω)

    ϑ₁(t,q) = q[2]
    ϑ₂(t,q) = zero(eltype(q))

    function ϑ(q)
        p = zero(q)
        p[1] = ϑ₁(0,q)
        p[2] = ϑ₂(0,q)
        return p
    end

    function hamiltonian(t, q, params)
        @unpack k = params
        q[2]^2 / 2 + k * q[1]^2 / 2
    end

    function hamiltonian(t, q, p, params)
        @unpack k = params
        p[1]^2 / 2 + k * q[1]^2 / 2
    end

    function lagrangian(t, q, v, params)
        ϑ₁(t,q) * v[1] + ϑ₂(t,q) * v[2] - hamiltonian(t, q, params)
    end

    t₀=0.0
    q₀=[0.5, 0.0]
    z₀=[0.5, 0.0, 0.5]
    p₀=ϑ(q₀)

    A = sqrt(q₀[2]^2 / k + q₀[1]^2)
    ϕ = asin(q₀[1] / A)

    refq = A * sin(ω * Δt * nt + ϕ)
    refp = ω * Δt * nt * A * cos(ω * Δt * nt + ϕ)
    refx = [refq, refp]

    function oscillator_ode_v(t, x, v, params)
        @unpack k = params
        v[1] = x[2]
        v[2] = -k*x[1]
        nothing
    end

    function harmonic_oscillator_ode(x₀=q₀, params=p)
        # @assert size(x₀,1) == 2
        ODE(oscillator_ode_v, x₀; invariants=(h=hamiltonian,), parameters=params)
    end

    function oscillator_pode_v(t, q, p, v, params)
        v[1] = p[1]
        nothing
    end

    function oscillator_pode_f(t, q, p, f, params)
        @unpack k = params
        f[1] = -k*q[1]
        nothing
    end

    function harmonic_oscillator_pode(q₀=[q₀[1]], p₀=[p₀[1]], params=p)
        # @assert length(q₀) == length(p₀)
        # @assert all([length(q) == length(p) == 1 for (q,p) in zip(q₀,p₀)])
        # @assert size(q₀,1) == size(p₀,1) == 1
        PODE(oscillator_pode_v, oscillator_pode_f, q₀, p₀; invariants=(h=hamiltonian,), parameters=params)
    end

    function harmonic_oscillator_hode(q₀=[q₀[1]], p₀=[p₀[1]], params=p)
        # @assert length(q₀) == length(p₀)
        # @assert all([length(q) == length(p) == 1 for (q,p) in zip(q₀,p₀)])
        # @assert size(q₀,1) == size(p₀,1) == 1
        HODE(oscillator_pode_v, oscillator_pode_f, hamiltonian, q₀, p₀; parameters=params)
    end

    function oscillator_sode_v_1(t, q, v, params)
        v[1] = q[2]
        v[2] = 0
        nothing
    end

    function oscillator_sode_v_2(t, q, v, params)
        @unpack k = params
        v[1] = 0
        v[2] = -k*q[1]
        nothing
    end

    function oscillator_sode_q_1(t, q̄, q, h, params)
        q[1] = q̄[1] + h * q̄[2]
        q[2] = q̄[2]
        nothing
    end

    function oscillator_sode_q_2(t, q̄, q, h, params)
        @unpack k = params
        q[1] = q̄[1]
        q[2] = q̄[2] - h * k*q̄[1]
        nothing
    end

    function harmonic_oscillator_sode(q₀=q₀, params=p)
        SODE((oscillator_sode_v_1, oscillator_sode_v_2),
             (oscillator_sode_q_1, oscillator_sode_q_2),
             q₀; parameters=params)
    end

    function oscillator_iode_ϑ(t, q, p, params)
        p[1] = q[2]
        p[2] = 0
        nothing
    end

    function oscillator_iode_ϑ(t, q, v, p, params)
        oscillator_iode_ϑ(t, q, p, params)
    end

    function oscillator_iode_f(t, q, v, f, params)
        @unpack k = params
        f[1] = -k*q[1]
        f[2] = v[1] - q[2]
        nothing
    end

    function oscillator_iode_g(t, q, λ, g, params)
        g[1] = 0
        g[2] = λ[1]
        nothing
    end

    function oscillator_iode_v(t, q, v, params)
        @unpack k = params
        v[1] = q[2]
        v[2] = -k*q[1]
        nothing
    end

    function harmonic_oscillator_iode(q₀=q₀, p₀=ϑ(q₀), params=p)
        # @assert size(q₀,1) == size(p₀,1) == 2
        IODE(oscillator_iode_ϑ, oscillator_iode_f,
             oscillator_iode_g, q₀, p₀;
             invariants=(h=hamiltonian,), parameters=params,
             v̄=oscillator_iode_v)
    end

    function oscillator_dae_v(t, z, v, params)
        @unpack k = params
        v[1] = z[2]
        v[2] = -k*z[1]
        v[3] = z[2] - k*z[1]
        nothing
    end

    function oscillator_dae_u(t, z, λ, u, params)
        u[1] = -λ[1]
        u[2] = -λ[1]
        u[3] = +λ[1]
    end

    function oscillator_dae_ϕ(t, z, ϕ, params)
        ϕ[1] = z[3] - z[1] - z[2]
    end

    function harmonic_oscillator_dae(z₀=z₀, λ₀=[zero(eltype(z₀))], params=p)
        # @assert size(z₀,1) == 3
        # @assert size(λ₀,1) == 1
        # @assert all([length(z) == 3 for z in z₀])
        # @assert all([length(λ) == 1 for λ in λ₀])
        DAE(oscillator_dae_v, oscillator_dae_u, oscillator_dae_ϕ,
            z₀, λ₀; invariants=(h=hamiltonian,), parameters=params, v̄=oscillator_ode_v)
    end

    function oscillator_idae_u(t, q, p, λ, u, params)
        u[1] = λ[1]
        u[2] = λ[2]
        nothing
    end

    function oscillator_idae_g(t, q, p, λ, g, params)
        g[1] = 0
        g[2] = λ[1]
        nothing
    end

    function oscillator_hdae_ū(t, q, p, λ, u, params)
        u[1] = λ[1]
        u[2] = λ[2]
        nothing
    end

    function oscillator_hdae_ḡ(t, q, p, λ, g, params)
        g[1] = 0
        g[2] = λ[1]
        nothing
    end

    function oscillator_idae_ϕ(t, q, p, ϕ, params)
        ϕ[1] = p[1] - q[2]
        ϕ[2] = p[2]
        nothing
    end

    function oscillator_hdae_ψ(t, q, p, v, f, ψ, params)
        ψ[1] = f[1] - v[2]
        ψ[2] = f[2]
        nothing
    end

    function harmonic_oscillator_idae(q₀=q₀, p₀=ϑ(q₀), λ₀=zero(q₀), params=p)
        # @assert size(q₀,1) == size(p₀,1) == size(λ₀,1) == 2
        IDAE(oscillator_iode_ϑ, oscillator_iode_f,
             oscillator_idae_u, oscillator_idae_g, oscillator_idae_ϕ,
             q₀, p₀, λ₀; invariants=(h=hamiltonian,), parameters=params, v̄=oscillator_iode_v)
    end

    function oscillator_pdae_v(t, q, p, v, params)
        @unpack k = params
        v[1] = q[2]
        v[2] = -k*q[1]
        nothing
    end

    function oscillator_pdae_f(t, q, p, f, params)
        @unpack k = params
        f[1] = -k*q[1]
        f[2] = p[1] - q[2]
        nothing
    end

    function harmonic_oscillator_pdae(q₀=q₀, p₀=ϑ(q₀), λ₀=zero(q₀), params=p)
        # @assert size(q₀,1) == size(p₀,1) == 2
        PDAE(oscillator_pdae_v, oscillator_pdae_f,
             oscillator_idae_u,  oscillator_idae_g, oscillator_idae_ϕ,
             q₀, p₀, λ₀; invariants=(h=hamiltonian,), parameters=params)
    end

    function harmonic_oscillator_hdae(q₀=q₀, p₀=ϑ(q₀), λ₀=zero(q₀), params=p)
        # @assert size(q₀,1) == size(p₀,1) == 2
        HDAE(oscillator_pdae_v, oscillator_pdae_f,
             oscillator_idae_u, oscillator_idae_g, oscillator_idae_ϕ,
             oscillator_idae_ū, oscillator_idae_ḡ, oscillator_idae_ψ,
             hamiltonian, q₀, p₀, λ₀; parameters=params)
    end

    function compute_energy_error(t, q::AbstractDataSeries{T}, parameters) where {T}
        h = DataSeries(T, q.nt)
        e = DataSeries(T, q.nt)

        for i in axes(q,2)
            h[i] = hamiltonian(t[i], q[:,i], parameters)
            e[i] = (h[i] - h[0]) / h[0]
        end

        (h, e)
    end

end


using Images

"""
```
lines = hough_transform_standard(
    img_edges::AbstractMatrix;
    stepsize=1,
    angles=range(0,stop=pi,length=minimum(size(img))),
    vote_threshold=minimum(size(img)) / stepsize -1,
    max_linecount=typemax(Int))
```

Returns a vector of tuples corresponding to the tuples of (r,t)
where r and t are parameters for normal form of line:
    `x * cos(t) + y * sin(t) = r`

-   `r` = length of perpendicular from (1,1) to the line
-   `t` = angle between perpendicular from (1,1) to the line and x-axis

The lines are generated by applying hough transform on the image.

Parameters:
-    `img_edges`      = Image to be transformed (eltype should be `Bool`)
-    `stepsize`       = Discrete step size for perpendicular length of line
-    `angles`         = List of angles for which the transform is computed
-    `vote_threshold` = Accumulator threshold for line detection
-    `max_linecount`  = Maximum no of lines to return

# Example
```julia
julia> using ImageFeatures

julia> img = fill(false,5,5); img[3,:] .= true; img
5×5 Array{Bool,2}:
 false  false  false  false  false
 false  false  false  false  false
  true   true   true   true   true
 false  false  false  false  false
 false  false  false  false  false

julia> hough_transform_standard(img)
1-element Array{Tuple{Float64,Float64},1}:
 (3.0, 1.5707963267948966)
```
"""
function hough_transform_standard(
    img_edges::AbstractMatrix{Bool};
    stepsize=1,
    angles=range(0,stop=pi,length=minimum(size(img_edges))),
    vote_threshold=minimum(size(img_edges)) / stepsize -1,
    max_linecount=typemax(Int))

    stepsize > 0 || error("Discrete step size must be positive")
    ρ = stepsize
    θ = angles

    #function to compute local maximum lines with values > threshold and return a vector containing them
    function findlocalmaxima!(validLines::AbstractVector{CartesianIndex{2}}, accumulator_matrix::Array{Int,2}, threshold)
        for val in CartesianIndices(size(accumulator_matrix))
            if  accumulator_matrix[val] >  threshold                             &&
                accumulator_matrix[val] >  accumulator_matrix[val[1],val[2] - 1] &&
                accumulator_matrix[val] >= accumulator_matrix[val[1],val[2] + 1] &&
                accumulator_matrix[val] >  accumulator_matrix[val[1] - 1,val[2]] &&
                accumulator_matrix[val] >= accumulator_matrix[val[1] + 1,val[2]]
                push!(validLines,val)
            end
        end
    end

    indsy, indsx = axes(img_edges)
    ρinv = 1 / ρ
    numangle = length(θ)
    numrho = round(Int,(2(length(indsx) + length(indsy)) + 1)*ρinv)

    accumulator_matrix = zeros(Int, numangle + 2, numrho + 2)

    #Pre-Computed sines and cosines in tables
    sinθ, cosθ = sin.(θ).*ρinv, cos.(θ).*ρinv

    #Hough Transform implementation
    constadd = round(Int,(numrho -1)/2)
    for pix in CartesianIndices(size(img_edges))
        if img_edges[pix]
            for i in 1:numangle
                dist = round(Int, pix[1] * sinθ[i] + pix[2] * cosθ[i])
                dist += constadd
                accumulator_matrix[i + 1, dist + 1] += 1
            end
        end
    end

    #Finding local maximum lines
    validLines = Vector{CartesianIndex{2}}(undef, 0)
    findlocalmaxima!(validLines, accumulator_matrix, vote_threshold)

    #Sorting by value in accumulator_matrix
    @noinline sort_by_votes(validLines, accumulator_matrix) = sort!(validLines, lt = (a,b)-> accumulator_matrix[a]>accumulator_matrix[b])
    sort_by_votes(validLines, accumulator_matrix)

    max_linecount = min(max_linecount, length(validLines))

    lines = Vector{Tuple{Float64,Float64}}(undef, 0)

    #Getting lines with Maximum value in accumulator_matrix && size(lines) < max_linecount
    for l in 1:max_linecount
        lrho = ((validLines[l][2]-1) - (numrho-1)*0.5)*ρ
        langle = θ[validLines[l][1]-1]
        push!(lines,(lrho,langle))
    end

    lines
end

"""
```
circle_centers, circle_radius = hough_circle_gradient(img_edges, img_phase, radii; scale=1, min_dist=minimum(radii), vote_threshold)
```
Returns two vectors, corresponding to circle centers and radius.

The circles are generated using a hough transform variant in which a non-zero point only votes for circle
centers perpendicular to the local gradient. In case of concentric circles, only the largest circle is detected.

Parameters:
-   `img_edges`    = edges of the image
-   `img_phase`    = phase of the gradient image
-   `radii`        = circle radius range
-   `scale`        = relative accumulator resolution factor
-   `min_dist`     = minimum distance between detected circle centers
-   `vote_threshold`   = accumulator threshold for circle detection

[`canny`](@ref) and [`phase`](@ref) can be used for obtaining img_edges and img_phase respectively.

# Example

```julia
julia> using Images, ImageFeatures, FileIO, ImageView

julia> img = load(download("http://docs.opencv.org/3.1.0/water_coins.jpg"));

julia> img = Gray.(img);

julia> img_edges = canny(img, (Percentile(99), Percentile(80)));

julia> dx, dy=imgradients(img, KernelFactors.ando5);

julia> img_phase = phase(dx, dy);

julia> centers, radii = hough_circle_gradient(img_edges, img_phase, 20:30);

julia> img_demo = Float64.(img_edges); for c in centers img_demo[c] = 2; end

julia> imshow(img_demo)
```
"""
function hough_circle_gradient(
        img_edges::AbstractArray{Bool,2},
        img_phase::AbstractArray{<:Number,2},
        radii::AbstractRange{<:Integer};
        scale::Number=1,
        min_dist::Number=minimum(radii),
        vote_threshold::Number=minimum(radii)*min(scale, length(radii)))

    rows,cols=size(img_edges)

    non_zeros=CartesianIndex{2}[]
    centers=CartesianIndex{2}[]
    circle_centers=CartesianIndex{2}[]
    circle_radius=Int[]
    accumulator_matrix=zeros(Int, Int(floor(rows/scale))+1, Int(floor(cols/scale))+1)

    function vote!(accumulator_matrix, x, y)
        fx = Int(floor(x))
        fy = Int(floor(y))

        for i in fx:fx+1
            for j in fy:fy+1
                if checkbounds(Bool, accumulator_matrix, i, j)
                    @inbounds accumulator_matrix[i, j] += 1
                end
            end
        end
    end

    for j in axes(img_edges, 2)
        for i in axes(img_edges, 1)
            if img_edges[i,j]
                sinθ = -cos(img_phase[i,j]);
                cosθ = sin(img_phase[i,j]);

                for r in radii
                    x=(i+r*sinθ)/scale
                    y=(j+r*cosθ)/scale
                    vote!(accumulator_matrix, x, y)

                    x=(i-r*sinθ)/scale
                    y=(j-r*cosθ)/scale
                    vote!(accumulator_matrix, x, y)
                end
                push!(non_zeros, CartesianIndex{2}(i,j));
            end
        end
    end

    for i in findlocalmaxima(accumulator_matrix)
        if accumulator_matrix[i]>vote_threshold
            push!(centers, i);
        end
    end

    @noinline sort_by_votes(centers, accumulator_matrix) = sort!(centers, lt=(a, b) -> accumulator_matrix[a]>accumulator_matrix[b])

    sort_by_votes(centers, accumulator_matrix)

    dist(a, b) = sqrt(sum(abs2, (a-b).I))

    f = CartesianIndex(map(r->first(r), axes(accumulator_matrix)))
    l = CartesianIndex(map(r->last(r), axes(accumulator_matrix)))
    radius_accumulator=Vector{Int}(undef, Int(floor(dist(f,l)/scale)+1))

    for center in centers
        center=(center-1*one(center))*scale
        fill!(radius_accumulator, 0)

        too_close=false
        for circle_center in circle_centers
            if dist(center, circle_center)< min_dist
                too_close=true
                break
            end
        end
        if too_close
            continue;
        end

        for point in non_zeros
            r=Int(floor(dist(center, point)/scale))
            if radii.start/scale<=r<=radii.stop/scale
                radius_accumulator[r+1]+=1
            end
        end

        voters, radius = findmax(radius_accumulator)
        radius=(radius-1)*scale;

        if voters>vote_threshold
            push!(circle_centers, center)
            push!(circle_radius, radius)
        end
    end
    return circle_centers, circle_radius
end

@deprecate hough_circle_gradient(
        img_edges, img_phase,
        scale, min_dist,
        vote_threshold, radii) hough_circle_gradient(img_edges, img_phase, radii; scale=scale, min_dist=min_dist, vote_threshold=vote_threshold)

@deprecate hough_transform_standard(image, ρ, θ, threshold, linesMax) hough_transform_standard(
    image, stepsize=ρ, angles=θ, vote_threshold=threshold, max_linecount=linesMax)


A = sparse([2,4,3],[2,2,4],[0.2,0.4,0.6])
Af = Matrix(A)

let k
k = 0
for I in eachindex(stored(A))
    @test A[I] == A.nzval[k+=1]
end

k = 0
for j = inds(A, 2)
    for I in eachindex(stored(A, :, j))
        @test A[I] == A.nzval[k+=1]
    end
end

k = 0
for v in each(stored(A))
    @test v == A.nzval[k+=1]
end

k = 0
for j = inds(A, 2)
    for v in each(stored(A, :, j))
        @test v == A.nzval[k+=1]
    end
end

k = 0
for I in each(index(A))
    @test A[I] == Af[k+=1]
end

k = 0
for j = inds(A, 2)
    for I in each(index(A, :, j))
        @test A[I] == Af[k+=1]
    end
end

k = 0
for v in each(A)
    @test v == Af[k+=1]
end

k = 0
for j = inds(A, 2)
    for v in each(A, :, j)
        @test v == Af[k+=1]
    end
end
end

# Sparse matrix-vector multiplication
function matvecmul_ind!(b::AbstractVector, A::AbstractMatrix, x::AbstractVector)
    fill!(b, 0)
    inds(A, 2) == inds(x, 1) || throw(DimensionMismatch("inds(A, 2) = $(inds(A, 2)) does not agree with inds(x, 1) = $(inds(x, 1))"))
    for j in inds(A, 2)
        xj = x[j]
        for (ib, iA) in sync(Follower(index(b)), index(stored(A, :, j)))
            b[ib] += A[iA]*xj
        end
    end
    b
end
function matvecmul_val!(b::AbstractVector, A::AbstractMatrix, x::AbstractVector)
    fill!(b, 0)
    inds(A, 2) == inds(x, 1) || throw(DimensionMismatch("inds(A, 2) = $(inds(A, 2)) does not agree with inds(x, 1) = $(inds(x, 1))"))
    for j in inds(A, 2)
        xj = x[j]
        for (ib, a) in sync(Follower(index(b)), stored(A, :, j))
            b[ib] += a*xj
        end
    end
    b
end

x = [1,-5,7,-13]
btrue = A*x
b = similar(btrue)
matvecmul_ind!(b, A, x)
@test b ≈ btrue
matvecmul_val!(b, A, x)
@test b ≈ btrue
@test_throws DimensionMismatch matvecmul_ind!(b, A, [1])
@test_throws DimensionMismatch matvecmul_ind!([0.1,0.2], A, x)
@test_throws DimensionMismatch matvecmul_val!(b, A, [1])
@test_throws DimensionMismatch matvecmul_val!([0.1,0.2], A, x)


### A Pluto.jl notebook ###
# v0.14.8

using Markdown
using InteractiveUtils

# ╔═╡ 4e019dd0-9a70-4034-915b-baf59e3416e5
using Fermi

# ╔═╡ e45bf90a-9cfe-46e8-bb0f-1c42723f1389
using TensorOperations

# ╔═╡ 819f62ec-da7c-11eb-300f-333f486858d2
html"""

  JuliaCon 2021

 Gustavo Aroeira

Center for Computational Quantum Chemistry

University of Georgia, Athens, GA

United States of America

"""

# ╔═╡ e1bd2ce0-a958-4460-b48f-620db0743e45
md"""
## 1. Quantum Chemistry
The main goal in quantum chemistry is to solve the electronic Schrödinger equation for many-electron systems

$\hat{H}|\Psi\rangle = E|\Psi\rangle$

A first approximation to the fermionic wave function, known as *mean-field*, is represented as a Slater determinant

$|\Phi_0\rangle = \mathcal{A}[\phi_1(\mathbf{x_1})\phi_2(\mathbf{x_2})\phi_3(\mathbf{x_3})...\phi_N(\mathbf{x_N})]$

where $\mathcal{A}$ is the antisymmetrizer operator and $\phi_i$ are spin-orbitals, that is, one-electron wave functions used to build the many-electron *ansatz*. Using the *basis set approximation*, we represent the orbitals as an expansion of $M$ pre-selected basis functions

$\phi_i = \sum_\mu^M C_{\mu i} \chi_\mu$

This is known as the **L**inear **C**ombination of **A**tomic **O**rbitals (LCAO) as the basis functions are designed to resemble atomic orbitals. Finally, more precise, i.e. correlated, wave functions are built as combinations of different Slater determinants (they differ by the orbitals within each one).

$|\Psi\rangle = \sum_i c_i|\Phi_i\rangle$

or, in cluster notation

$|\Psi\rangle =  e^{T}|\Phi_0\rangle$

**Fermi** is a quantum chemistry package in Julia specialized in wave function methods. Currently we have implemented the following methods

- Restricted Hartree-Fock
- Møller-Plesset Perturbation Theory
- Coupled Cluster Theory

Moreover, Fermi aims to serve as development platform for molecular quantum mechanics in Julia
"""

# ╔═╡ ff5a766b-716b-47dc-a94f-1f9c3f74ac8b
md"""
## 2. Molecular Integrals

Basic quantities required in all quantum chemistry computations

###### Overlap (S)

$S_{ij} = \int \phi_i^* \phi_j d\tau$

###### Kinetic (T)

$T_{ij} = \int \phi_i^* \left( -\frac{1}{2} \nabla^2 \right) \phi_j d\tau$

###### Nuclear (V)

$V_{ij} = \sum_k \int \phi_i^* \frac{Z_k}{r} \phi_j d\tau$

###### Eletron repulsion integral (ERI)

$G_{ijkl} = \int \phi_i^* \phi_j^* \frac{1}{r_{12}} \phi_k \phi_l d\tau$

Fermi utilizes the library [libcint](https://github.com/sunqm/libcint) to obtain molecular integrals.
"""

# ╔═╡ 67892086-7643-40a0-9e5e-eb56607412de
begin
@set {
	diis false
	oda false
}
nothing
end

# ╔═╡ e8585228-0be6-41dd-8434-045ed69d578c
md"""
The module `Fermi.Libcint` has minimal wraps around libcint functions
"""

# ╔═╡ 3a325265-7699-4f99-a850-207ec00f9c26
names(Fermi.Libcint)

# ╔═╡ 8dced24e-6cc6-46aa-bd50-658907190987
md"""
A more convinient way is to use the `IntegralHelper`.

First, define a molecule and basis set
"""

# ╔═╡ 01228770-9ad9-4634-9aba-afc798ae3e8c
@molecule {
	H 0.0 0.0 0.0
	H 0.76 0.0 0.0
}

# ╔═╡ 597aa086-dd0c-4ef9-b0ef-ab294b89b3a2
@set basis sto-3g

# ╔═╡ bd1dca4c-f126-4353-b09a-e2a641712a09
I = Fermi.Integrals.IntegralHelper()

# ╔═╡ c37cef53-ddfa-4bce-9328-21cde32bf8da
I["S"]

# ╔═╡ 73551bfb-c7c6-4f5b-a931-92c23b7feafd
md"""
## 3. Multiple dispatch in Fermi

`IntegralHelper` guides how amplitudes are updated in a CCSD computation
"""

# ╔═╡ b03d9e3f-8ef5-4e48-b2aa-f6b4295ac228
html"""

"""

# ╔═╡ fb42172a-33df-4b03-88d5-6db7224f132b
md"""
## 4. Tensor Contractions

In methods such as Coupled Cluster, tensor contractions appears in terms like

$t_{ij}^{ab} \leftarrow -2t_{ik}^{ac}\cdot t_{lj}^{bd}\cdot v_{kc}^{ld}$

The package [TensorOperations.jl](https://github.com/Jutho/TensorOperations.jl) is used to handle expressions of this type.

`FermiMDArray` is a thin wrap around a Julia array that allow us to use [TBLIS](https://github.com/devinamatthews/tblis) as a tensor contraction tool.
"""

# ╔═╡ d5942379-2e27-43ca-920e-889a64d32269
begin
	o_size = 10
	v_size = 100
	T2 = FermiMDrand(o_size, o_size, v_size, v_size)
	newT2 = FermiMDrand(o_size, o_size, v_size, v_size)
	Vovov = FermiMDrand(o_size, v_size, o_size, v_size)
	size(Vovov)
end

# ╔═╡ d0c5639e-49f7-45c0-a19e-d1da88b81ce0
t = @elapsed begin
	@tensoropt newT2[i,j,a,b] = -2*T2[i,k,a,c]*T2[l,j,b,d]*Vovov[k,c,l,d]
end

# ╔═╡ e89556da-e726-47bf-ae91-4312b24095b1
@set tblis true

# ╔═╡ Cell order:
# ╟─819f62ec-da7c-11eb-300f-333f486858d2
# ╟─e1bd2ce0-a958-4460-b48f-620db0743e45
# ╟─ff5a766b-716b-47dc-a94f-1f9c3f74ac8b
# ╟─67892086-7643-40a0-9e5e-eb56607412de
# ╠═4e019dd0-9a70-4034-915b-baf59e3416e5
# ╟─e8585228-0be6-41dd-8434-045ed69d578c
# ╠═3a325265-7699-4f99-a850-207ec00f9c26
# ╟─8dced24e-6cc6-46aa-bd50-658907190987
# ╠═01228770-9ad9-4634-9aba-afc798ae3e8c
# ╠═597aa086-dd0c-4ef9-b0ef-ab294b89b3a2
# ╠═bd1dca4c-f126-4353-b09a-e2a641712a09
# ╠═c37cef53-ddfa-4bce-9328-21cde32bf8da
# ╟─73551bfb-c7c6-4f5b-a931-92c23b7feafd
# ╟─b03d9e3f-8ef5-4e48-b2aa-f6b4295ac228
# ╟─fb42172a-33df-4b03-88d5-6db7224f132b
# ╠═e45bf90a-9cfe-46e8-bb0f-1c42723f1389
# ╠═d5942379-2e27-43ca-920e-889a64d32269
# ╠═d0c5639e-49f7-45c0-a19e-d1da88b81ce0
# ╠═e89556da-e726-47bf-ae91-4312b24095b1


using Test
using OrdinaryDiffEq
using DiffEqSensitivity
using DiffEqBase
using ForwardDiff
using QuadGK
using Zygote

function pendulum_eom(dx, x, p, t)
    dx[1] = p[1] * x[2]
    dx[2] = -sin(x[1]) + (-p[1]*sin(x[1]) + p[2]*x[2])  # Second term is a simple controller that stabilizes π
end

x0 = [0.1, 0.0]
tspan = (0.0, 10.0)
p = [1.0, -24.05, -19.137]
prob = ODEProblem(pendulum_eom, x0, tspan, p)
sol = solve(prob, Vern9(), abstol=1e-8, reltol=1e-8)

g(x, p, t) = 1.0*(x[1] - π)^2 + 1.0*x[2]^2 + 5.0*(-p[1]*sin(x[1]) + p[2]*x[2])^2
dgdu(out, y, p, t) = ForwardDiff.gradient!(out, y -> g(y, p, t), y)
dgdp(out, y, p, t) = ForwardDiff.gradient!(out, p -> g(y, p, t), p)

res_interp = adjoint_sensitivities(sol,Vern9(),g,nothing,(dgdu, dgdp),abstol=1e-8,
                                reltol=1e-8,iabstol=1e-8,ireltol=1e-8, sensealg=InterpolatingAdjoint())
res_quad = adjoint_sensitivities(sol,Vern9(),g,nothing,(dgdu, dgdp),abstol=1e-8,
                                reltol=1e-8,iabstol=1e-8,ireltol=1e-8, sensealg=QuadratureAdjoint())
#res_back = adjoint_sensitivities(sol,Vern9(),g,nothing,(dgdu, dgdp),abstol=1e-8,
#                                reltol=1e-8,iabstol=1e-8,ireltol=1e-8, sensealg=BacksolveAdjoint(checkpointing=true), sol=sol.t) # it's blowing up

function G(p)
    tmp_prob = remake(prob,p=p,u0=convert.(eltype(p), prob.u0))
    sol = solve(tmp_prob,Vern9(),abstol=1e-8,reltol=1e-8)
    res,err = quadgk((t)-> g(sol(t), p, t), 0.0,10.0,atol=1e-8,rtol=1e-8)
    res
end
res2 = ForwardDiff.gradient(G,p)

@test res_interp[2]' ≈ res2 atol=1e-5
@test res_quad[2]' ≈ res2 atol=1e-5

p = [2.0,3.0]
u0 = [2.0]
function f(du,u,p,t)
    du[1] = -u[1]*p[1]-p[2]
end

prob = ODEProblem(f,u0,(0.0,1.0),p)
sol = solve(prob,Tsit5(),abstol=1e-10,reltol=1e-10);

g(u,p,t) = -u[1]*p[1]-p[2]

dgdu(out, y, p, t) = ForwardDiff.gradient!(out, y -> g(y, p, t), y)
dgdp(out, y, p, t) = ForwardDiff.gradient!(out, p -> g(y, p, t), p)
du0,dp = adjoint_sensitivities(sol,Vern9(),g,nothing,(dgdu,dgdp);abstol=1e-10,reltol=1e-10)

function G(p)
    tmp_prob = remake(prob,p=p,u0=convert.(eltype(p), prob.u0))
    sol = solve(tmp_prob,Vern9(),abstol=1e-8,reltol=1e-8)
    res,err = quadgk((t)-> g(sol(t), p, t), 0.0,10.0,atol=1e-8,rtol=1e-8)
    res
end
res2 = ForwardDiff.gradient(G,p)

@test dp' ≈ res2 atol=1e-5

function model(p)
    N_oscillators = 30
    u0 = repeat([0.0; 1.0], 1, N_oscillators) # size(u0) = (2, 30)

    function du!(du, u, p, t)
        W, b  = p # Parameters
        dy  = @view du[1,:] # 30 elements
        dy′ = @view du[2,:]
        y = @view u[1,:]
        y′= @view u[2,:]
        @. dy′ = -y * W
        @. dy  = y′ * b
    end

    output = solve(
        ODEProblem(
            du!,
            u0,
            (0.0, 10.0),
            p,
            jac = true,
            abstol = 1e-12,
            reltol = 1e-12),
        Tsit5(),
        jac = true,
        saveat = collect(0:0.1:7),
        sensealg = QuadratureAdjoint(),
    )
    return Array(output[1, :, :]) # only return y, not y′
end

p=[1.5, 0.1]
y = model(p)
loss(p) = sum(model(p))
dp1 = Zygote.gradient(loss,p)[1]
dp2 = ForwardDiff.gradient(loss,p)
@test dp1 ≈ dp2


include("header.jl")

#=
# Fluid properties, flow quantities, and units
In this notebook, we will discuss *fluid properties* and *flow quantities*,
and the systems of units we use for them. First, some basic definitions:
- Fluid properties are properties of the material (the fluid)
- Flow quantities are characteristics of the flow of this fluid
We will also introduce some syntax we will use in this notebook and those that follow for dealing with units.
=#

# ### Set up the module
using MAE103

#=
We will generally focus on SI units, and these will be the default system
for displaying quantities. However, we also need to be familiar
with Imperial (sometimes called "English") units, which arise in many
situations in engineering. In the examples below, we will show that the tools
in these notebooks allow us to easily convert from one system to another.

For any quantity, we can see what the default units are by using `default_unit`.
For example, for pressure,
=#
default_unit(Pressure)
#-
#=
### Fluid properties
Let's start by discussing the basic properties of a fluid.

#### Density
The **density** provides a measure of the amount of fluid material per unit volume.
It is measured in units of mass/volume:
=#
default_unit(Density)

#=
For example, the density of water at a reference temperature of 15.6 degrees C is
=#
Density(Water)

#=
and for air at temperature of 15 degrees C and pressure of 1 atmosphere is
=#
Density(Air)

#=
Notice that water is around 800 times denser than air. This fact is very important
in fluid mechanics!
=#

#=
Remember that the density of gases, like air, depend on pressure and temperature (by
the ideal gas law), so the density of a gas may be sensitive to the local conditions.
This will happen in flows traveling relatively fast ([Q: compared to what?]). In such a case, density is a flow quantity,
not a fluid property, and we have to determine its value as part of the problem.

However, when a gas is traveling relatively slow [Q: compared to what?], we can
often treat a gas as having constant density, so it can be treated as a fluid property.

And liquid density is generally not sensitive. In fact, usually we can assume that the
density in a liquid is constant and uniform. We can treat
density in a liquid as a fluid property rather than a flow quantity.
=#
#-
#=
Also note that the density of seawater is larger (at average salt concentration)
than that of freshwater
=#
Density(Seawater)

#=
What if we want to set the density in other units? For example, in imperial
units, we would usually set it with lbm/ft^3. (In the notebooks, we use `lbm` for
pound (mass), and `lbf` is used for pound (force).) To set the units of a quantity,
rather than rely on the default units, we use the syntax to follow the number with
 `u"units"`. It will automatically convert it to the default units. For example, 2 lb/ft^3:
=#
Density(2u"lbm/ft^3")

#=
and to report the value of a quantity in other units, we use the `value`
function. For example, to get the density of water in lb/ft^3:
=#
value(Density(Water),u"lbm/ft^3")

#=
#### Viscosity
The internal friction in the fluid, called the *viscosity* (or, more specifically, the *dynamic viscosity*), given
by the Greek symbol $\mu$. [We can get this symbol by typing `\mu+TAB`.] The
viscosity controls the relationship between *shear stress*  $\tau$ (the frictional
force per unit area) and the *strain rate*, given by the gradient of velocity,
$\mathrm{d}u/\mathrm{d}y$. This latter quantity measures the difference in speeds in adjacent layers
of fluid, and larger differences suggest more shear stress. Viscosity is
the proportionality constant

$$\tau = \mu \dfrac{\mathrm{d}u}{\mathrm{d}y}$$

We will learn
much more about viscosity later, but for now, it is sufficient to
know that viscosity has units of kg/m/s:
=#
μw = Viscosity(Water)
#-
μa = Viscosity(Air)

#=
Note that water is much more viscous the air:
=#
μw/μa
#=
This is probably intuitive to you.
=#
#=
Note that both of these viscosities are much smaller than that of glycerin:
=#
Viscosity(Glycerin)
#-
#=
We will also occasionally make use of the ratio between viscosity and density.
This is called the *kinematic viscosity*. We use the symbol $\nu$ for this
[obtained by typing `\mu+TAB`.]
=#
νw = KinematicViscosity(Water)
#-
νa = KinematicViscosity(Air)

#=
#### Surface tension
The surface tension is a property associated with liquid interfaces, and particularly,
liquid interfaces with gases. The surface tension has units of force per unit length,
because if we imagine "cutting" a bit of the interface from the rest of it,
then this cut would form a perimeter of the snipped part of the interface, and
surface tension would act along this perimeter, representing how much the rest
of the liquid interface was pulling on it.

Surface tension also happens to have units of energy per unit area, so it
is sometimes referred to as "surface energy". Different liquids have different
surface energies, depending on the strength of their inter-molecular forces.
=#
SurfaceTension(Water)
#-
SurfaceTension(Glycerin)

#=
### Flow quantities
Now, let's discuss quantities that describe the fluid flow, or
at least, the *state* of the fluid. It is important to understand that
these quantities are, in general, *field quantities*: they vary from location
to location, and perhaps vary over time. So each of them should be
thought of as a *function* of the spatial coordinates, $(x,y,z)$ and time $t$.
Finding these functions is often our ultimate goal in solving a problem.
=#

#=
#### Pressure
Pressure represents the average force that the molecules exert per unit
area of surface. It is important to understand that pressure acts the same in every
direction. That is, it is an *isotropic* quantity. This means that, no matter
what the orientation of the surface, the pressure acts the same on it, and further,
it only acts *perpendicular* to the surface.

The default SI unit of pressure is the *Pascal* (Pa, equal to 1 N/m^2, or 1 kg/m/s^2).
=#
Pressure(20)

#=
But there are many other units for pressure in use. Some common ones. The
atmosphere (atm) represents the standard ambient pressure.
=#
p = Pressure(1u"atm")
#=
So 1 atm is about 101325 Pa, or 101.325 kPa. Let's see this written in other units
=#
value(p,u"psi")
#-
value(p,u"mmHg")
#=
These are all *absolute* pressures. For many of the flows we will
study, the absolute pressure will not be important. Consider,
for example, the flow of water through a pipe. Only the *difference* between
the pressure at the inlet and the pressure at the outlet matters for
driving this flow; the absolute pressure is irrelevant. So
we can also define a `PressureDifference`, measured in the
same units:
=#
PressureDifference(50)

#=
Pressure is a flow quantity. In general, we need to find its values
as part of the problem.
=#

#=
#### Velocity
The velocity of the fluid describes how fast it is moving in each direction,
and has units of length/time.
=#
u = Velocity(20)
#-
#=
In other units,
=#
value(u,u"cm/s")
#-
value(u,u"ft/s")

#=
We know that velocity is a vector, so it has, in general, three components $(u,v,w)$.
It is also a field quantity. This means that each of the three components
depends on $(x,y,z)$ and $t$. That's a lot of detail to determine, but once
we determine it, we know *everything* about a fluid flow.
=#


using GeometricIntegrators
using SafeTestsets

set_config(:verbosity, 2)

@safetestset "1st Poincaré Invariant Unit Tests                                               " begin include("poincare_invariant_1st_unittests.jl") end
@safetestset "1st Poincaré Invariant Function Tests                                           " begin include("poincare_invariant_1st_tests.jl") end
@safetestset "2nd Poincaré Invariant Unit Tests                                               " begin include("poincare_invariant_2nd_unittests.jl") end
@safetestset "2nd Poincaré Invariant Function Tests                                           " begin include("poincare_invariant_2nd_tests.jl") end


# ---
# title: 352. Data Stream as Disjoint Intervals
# id: problem352
# author: Tian Jun
# date: 2020-10-31
# difficulty: Hard
# categories: Binary Search, Ordered Map
# link: <https://leetcode.com/problems/data-stream-as-disjoint-intervals/description/>
# hidden: true
# ---
#
# Given a data stream input of non-negative integers a1, a2, ..., an, ...,
# summarize the numbers seen so far as a list of disjoint intervals.
#
# For example, suppose the integers from the data stream are 1, 3, 7, 2, 6, ...,
# then the summary will be:
#
#
#
#     [1, 1]
#     [1, 1], [3, 3]
#     [1, 1], [3, 3], [7, 7]
#     [1, 3], [7, 7]
#     [1, 3], [6, 7]
#
#
#
#
# **Follow up:**
#
# What if there are lots of merges and the number of disjoint intervals are
# small compared to the data stream's size?
#
#
## @lc code=start
using LeetCode

## add your code here:
## @lc code=end


### A Pluto.jl notebook ###
# v0.17.0

using Markdown
using InteractiveUtils

# ╔═╡ 2bb195db-311f-4191-95a1-04d6a4b75d74
md"""
#### Practice Problem: Some interesting nuggets about equilibrium and kinetics

a) Starting from the mass-action kinetic expression for the _net_ rate of reaction of $j$ denoted by $\hat{r}_{j}$:

$$\hat{r}_{j} = k^{+}_{j}\prod_{i\in\mathcal{M}^{+}}\left[X_{i}\right]^{-\sigma_{ij}} -
k^{-}_{j}\prod_{i\in\mathcal{M}^{-}}\left[X_{i}\right]^{-\sigma_{ij}}$$ show that equilibrium constant for reaction $j$, denoted by $K_{j}$, equals:

$$K_{j} = \frac{k_{j}^{+}}{k_{j}^{-}}$$

where $\left[X_{i}\right]$ denotes the concetration of component $i$, $k_{j}^{+}$ denotes the forward rate constant, $k_{j}^{+}$ denotes the backward rate constant, $i\in\mathcal{M}^{+}$ denotes the indicies of the reactants for the forward reaction, and $i\in\mathcal{M}^{-}$ denotes the indicies of the reactants in the reverse direction for reaction $j$. The quantity $\sigma_{ij}$ denotes the stoichiometric coefficient governing the behavior of component $i$ in reaction $j$.

b) If $k_{j}^{-}\rightarrow{0}$, what does that say about the $\Delta{G}$ of the reaction $j$?

c) Can chemical equilibrium occur in an _open_ system? Explain your answer in terms of an open system at steady-state with a single in/out stream, and the reaction $A\rightleftharpoons{B}$, versus a closed system with the same reaction.

d) If we gave you concentration measurements for $A$ as a function of time at different temperatures for a closed system in which the reaction $A\rightleftharpoons{B}$ was occurring, explain how you would use it to compute rate constants, activation energies, frequency factors, etc, assuming the kinetics were first-order?

"""

# ╔═╡ 40f90c69-76c9-41fa-9fe5-2c3410d6bc43
md"""
#### Solutions
"""

# ╔═╡ cf9767ca-133c-4322-959d-3eb680013f63
md"""
###### a) Starting from the mass-action kinetic expression for the _net_ rate of reaction of $j$ denoted by $\hat{r}_{j}$:

$$\hat{r}_{j} = k^{+}_{j}\prod_{i\in\mathcal{M}^{+}}\left[X_{i}\right]^{-\sigma_{ij}} -
k^{-}_{j}\prod_{i\in\mathcal{M}^{-}}\left[X_{i}\right]^{-\sigma_{ij}}$$ show that equilibrium constant for reaction $j$, denoted by $K_{j}$, equals:

$$K_{j} = \frac{k_{j}^{+}}{k_{j}^{-}}$$

At equilibrium, the _net rate_ of all reversible reactions is zero $\hat{r}_{j} = 0$. Thus,

$$k^{+}_{j}\prod_{i\in\mathcal{M}^{+}}\left[X_{i}\right]^{-\sigma_{ij}} = k^{-}_{j}\prod_{i\in\mathcal{M}^{-}}\left[X_{i}\right]^{-\sigma_{ij}}$$

of after some rearrangement:

$$\frac{k_{j}^{+}}{k_{j}^{-}} = \frac{\prod_{i\in\mathcal{M}^{-}}\left[X_{i}\right]^{-\sigma_{ij}}}{\prod_{i\in\mathcal{M}^{+}}\left[X_{i}\right]^{-\sigma_{ij}}}$$

The right-hand side of this expression is the product concentrations divided by the reactant concentrations , which is the equilibrium constant you've developed in your chemistry courses. Thus:

$$\frac{k_{j}^{+}}{k_{j}^{-}}=K_{j}$$

"""

# ╔═╡ bd2184c1-a656-4e54-9f0d-fcc9009e4b04

# ╔═╡ 64a749bc-423b-413a-afb3-c4e32d53da04
md"""
###### b) If $k_{j}^{-}\rightarrow{0}$, what does that say about the $\Delta{G}$ of the reaction $j$?

As the rate of the back reaction goes to zero, the reaction becomes irreversible. To link this to $\Delta{G}$ of the reaction, we know that:

$$\Delta{G} = -RT\ln{K}$$

or (after substitution) of the ratio of rate constants:

$$\Delta{G} = -RT\ln\left({\frac{k_{j}^{+}}{k_{j}^{-}}}\right)$$

As $k_{j}^{-}\rightarrow{0}$ the $\Delta{G}\rightarrow{-\infty}$, thus reaction becomes incresingly energetically favorable.

"""

# ╔═╡ 9abaaaaa-8bfb-45a2-8da2-c53e7310ba0a

# ╔═╡ 36bf5490-13b5-4da8-8ef9-00e519983c11
md"""
###### c) Can chemical equilibrium occur in an open system?

Open an system can reach steady-state, but equilibrium has the additional constraint that all net reactions are equal to zero. Thus, let's consider the case when we have a CSTR with a single input ($s=1$) and a single output ($s=2$) and a reversible single reaction $\hat{r}_{1}$. The species mole balance for component $i$ would be:

$$\dot{n}_{i1} - \dot{n}_{i2} + \sigma_{i1}\hat{r}_{1}V = 0$$

If an open system were at equilibrium (as opposed to steady-state), then $\hat{r}=0$ which implies:

$$\dot{n}_{i1} = \dot{n}_{i2}$$

which is clearly not true in general. Thus, an open system can not reach equilibrium, but rather operates at a steady-state.

"""

# ╔═╡ 85eaab6b-b7fe-421a-a2e7-f2634b7eca25
md"""
###### d) Estimation of kinetic parameters from data

If we have a closed system, we could wait until the system reached equilibrium and then get the ratio of the rate constants from the value of the equilibrium constant.

Next, as we approach eqilibrium we could measure the concentration of A (or B) as a function of time. Becuase we are in a closed system we know that:

$$\frac{d\left[A\right]}{dt} = -\hat{r}_{1}$$

meaning the slope of an $\left[A\right]$ versus time plot would give the reaction rate. Futher since rate constants are only functions of temperature, we could use _measurements_ of the equilibrium constant, and $\left[A\right]$ and $\left[B\right]$ concentrations at the same temperature $T$ to arrive at the expression (when we are near equilibrium):

$$\hat{r}_{1}\simeq k_{1}^{+}\left(\left[A\right] - \frac{\left[B\right]}{K}\right)$$

where $k_{1}^{-} = k_{1}^{+}K^{-1}$ to estimate the forward rate constant. Lastly, to get frequency factor and the activation energy we could use the linearized version of the [Arrhenius equation](https://en.wikipedia.org/wiki/Arrhenius_equation) in combination with the required measurements.

"""

# ╔═╡ 6a67ec88-3d92-11ec-0fca-f5f1f3e769b3
html"""
<style>
main {
    max-width: 1200px;
    width: 75%;
    margin: auto;
    font-family: "Roboto, monospace";
}

a {
    color: blue;
    text-decoration: none;
}

.H1 {
    padding: 0px 30px;
}
</style>"""

# ╔═╡ Cell order:
# ╟─2bb195db-311f-4191-95a1-04d6a4b75d74
# ╟─40f90c69-76c9-41fa-9fe5-2c3410d6bc43
# ╟─cf9767ca-133c-4322-959d-3eb680013f63
# ╟─bd2184c1-a656-4e54-9f0d-fcc9009e4b04
# ╟─64a749bc-423b-413a-afb3-c4e32d53da04
# ╟─9abaaaaa-8bfb-45a2-8da2-c53e7310ba0a
# ╟─36bf5490-13b5-4da8-8ef9-00e519983c11
# ╠═85eaab6b-b7fe-421a-a2e7-f2634b7eca25
# ╟─6a67ec88-3d92-11ec-0fca-f5f1f3e769b3


mutable struct expected_holding_cost{T} <: AbstractCost
    h::T
end

function (f::expected_holding_cost)(y, demand_dist::Union{CVNormal,Normal})
    μ = mean(demand_dist)
    σ = std(demand_dist)
    σ == 0 && return linear_holding_cost(f.h)(y-μ)
    h = f.h
    v = σ^2
    e = MathConstants.e
    π = MathConstants.pi
    fh = zero(y)
    if y >= 0
        fh = (((e^(-((y - μ)^2)/2v) -e^(-(μ^2)/2v))*sqrt(v))/(sqrt(2π))) + 1//2*(y-μ)*(erf((y-μ)/(sqrt(2)*sqrt(v))) + erf(μ/(sqrt(2)*sqrt(v))))
        fh *= h
    end
    return fh
end

mutable struct expected_stockout_cost{T} <: AbstractCost
    b::T
end

function (f::expected_stockout_cost)(y, demand_dist::Union{CVNormal,Normal})
    μ = mean(demand_dist)
    σ = std(demand_dist)
    σ == 0 && return linear_stockout_cost(f.h)(y-μ)
    b = f.b
    v = σ^2
    e = MathConstants.e
    π = MathConstants.pi
    fb = zero(y)
    if y >= 0
        fb = ((e^(-((y - μ)^2)/2v))*v + sqrt(π/2) * (μ - y) * sqrt(v) * erfc((y-μ)/(sqrt(2)*sqrt(v))))/(sqrt(2π)*sqrt(v))
        fb *= b
    else
        fb = (e^(-(μ^2)/2v)*v + sqrt(π/2) * (μ - y) * sqrt(v) * (1 + erf(μ/(sqrt(2)*sqrt(v)))))/(sqrt(2π)*sqrt(v))
        fb *= b
    end
    return fb
end


# Conradi, C., Shiu, A.,
# Dynamics of post-translational modification systems: recent progress and future directions
# Eq. 3.4
using Logging

using StructuralIdentifiability

logger = Logging.SimpleLogger(stdout, Logging.Info)
global_logger(logger)

ode = @ODEmodel(
    x1'(t) = -k1 * x1(t) * x2(t) + k2 * x4(t) + k4 * x6(t),
    x2'(t) = -k1 * x1(t) * x2(t) + k2 * x4(t) + k3 * x4(t),
    x3'(t) = k3 * x4(t) + k5 * x6(t) - k6 * x3(t) * x5(t),
    x4'(t) = k1 * x1(t) * x2(t) - k2 * x4(t) - k3 * x4(t),
    x5'(t) = k4 * x6(t) + k5 * x6(t) - k6 * x3(t) * x5(t),
    x6'(t) = -k4 * x6(t) - k5 * x6(t) + k6 * x3(t) * x5(t),
    y1(t) = x3(t),
    y2(t) = x2(t)
)

@time println(assess_identifiability(ode))


module TruncatedGaussQuadrature

#=
To import module...
push!(LOAD_PATH, "/Path/To/Src/")
... then load the module
using TruncatedGaussQuadrature
.. Test the module with:
truncatedGaussQuad(0.0, 0.1, -Inf, 10.0, 5)
=#

import Distributions
import Combinatorics
using LinearAlgebra # Hermitian, chol, permute, eigen

# include Distributions # loads other packages that subsequent code depends ons
include("normalDistMoment.jl")
include("momentGaussQuad.jl")
include("spd.jl")

export truncatedGaussQuad, normalDistMoment, momentGaussQuad, buildHankel, spdMatrix

function truncatedGaussQuad(μ::Real, σ::Real, lb::Real, ub::Real, n::Int;matrixFix::Bool=true)

  # 1) Calculate moments
  mm = normalDistMoment(μ, σ, 2n, lb, ub)
  # 2) Build Hankel Matrix
  M = buildHankel(mm,matrixFix)
  # 3) Calculate nodes:
  x, w = momentGaussQuad(M)
  return x, w
end

function truncatedGaussQuad(lb::Real, ub::Real, n::Int;matrixFix::Bool=true)
  # N(0,1)
  x, w = truncatedGaussQuad(0.0,1.0,lb,ub,n,matrixFix=matrixFix)
  return x, w
end

function buildHankel(mm::Array{T,1},matrixFix::Bool) where T<:Real
  n = div((length(mm)-1),2)
  M = zeros(eltype(mm),n+1,n+1)
  for i = 1:n+1
    M[:,i] = mm[i:i+n]
  end
  if  ~isposdef(M) && matrixFix
      M = spdMatrix(M)
  end
  return M
end

end


using Random

@testset "Testing random_mc.jl" begin
    @testset "Test random_markov_chain" begin
        n, k = 5, 3
        mc_dicts = (Dict("P" => random_markov_chain(n).p, "k" => n),
                    Dict("P" => random_markov_chain(n, k).p, "k" => k))
        for d in mc_dicts
            P = d["P"]
            @test size(P) == (n, n)
            @test all(x->(count(!iszero, x)==d["k"]),
                      [P[i, :] for i in 1:size(P)[1]]) == true
        end

        seed = 1234
        rngs = [MersenneTwister(seed) for i in 1:2]
        mcs = random_markov_chain.(rngs, n, k)
        @test mcs[2].p == mcs[1].p
    end

    @testset "Test random_stochastic_matrix" begin
        n, k = 5, 3
        Ps = (random_stochastic_matrix(n), random_stochastic_matrix(n, k))
        for P in Ps
            @test all(P .>= 0) == true
            @test all(x->isapprox(sum(x), 1),
                      [P[i, :] for i in 1:size(P)[1]]) == true
        end

        seed = 1234
        rngs = [MersenneTwister(seed) for i in 1:2]
        Ps = random_stochastic_matrix.(rngs, n, k)
        @test Ps[2] == Ps[1]
    end

    @testset "Test random_stochastic_matrix with k=1" begin
        n, k = 3, 1
        P = random_stochastic_matrix(n, k)
        @test all((P .== 0) .| (P .== 1))
        @test all(x->isequal(sum(x), 1),
                  [P[i, :] for i in 1:size(P)[1]]) == true
    end

    @testset "Test errors properly thrown" begin
        # n <= 0
        @test_throws ArgumentError random_markov_chain(0)

        # k <= 0
        @test_throws ArgumentError random_markov_chain(2, 0)

        # k > n
        @test_throws ArgumentError random_markov_chain(2, 3)
    end

    @testset "Test random_discrete_dp" begin
        num_states, num_actions = 5, 4
        num_sa = num_states * num_actions
        k = 3
        ddp = random_discrete_dp(num_states, num_actions; k=k)

        # Check shapes
        @test size(ddp.R) == (num_states, num_actions)
        @test size(ddp.Q) == (num_states, num_actions, num_states)

        # Check ddp.Q[:, a, :] is a stochastic matrix for all actions `a`
        @test all(ddp.Q .>= 0) == true
        for a in 1:num_actions
            P = reshape(ddp.Q[:, a, :], (num_states, num_states))
            @test all(x->isapprox(sum(x), 1),
                      [P[i, :] for i in 1:size(P)[1]]) == true
        end

        # Check number of nonzero entries for each state-action pair
        @test sum(ddp.Q .> 0, dims = 3) ==
            ones(Int, (num_states, num_actions, 1)) * k

        seed = 1234
        rngs = [MersenneTwister(seed) for i in 1:2]
        ddps = random_discrete_dp.(rngs, num_states, num_actions)
        @test ddps[2].R == ddps[1].R
        @test ddps[2].Q == ddps[1].Q
        @test ddps[2].beta == ddps[1].beta
    end

    @testset "Test random_probvec" begin
        k = 5
        x = random_probvec(k)
        @test size(x) == (k,)
    end
end  # @testset


using MultivariateStats
using LinearAlgebra
using Test
using StableRNGs
import Statistics: mean, var, cov
import SparseArrays

@testset "PCA" begin

    rng = StableRNG(34568)

    ## PCA with zero mean

    X = randn(rng, 5, 10)
    Y = randn(rng, 3, 10)

    P = qr(randn(rng, 5, 5)).Q[:, 1:3]
    pvars = [5., 4., 3.]
    l = [-0.236399  -0.706558  -1.21106
          1.27426   -0.047799  -0.212282
          0.990598   1.22718   -0.882179
         -1.51861    0.725771  -0.631665
          0.181386  -1.21069   -0.557706]
    M = PCA(Float64[], P, pvars, 15.0)

    @test size(M) == (5, 3)
    @test mean(M) == zeros(5)
    @test projection(M) == P
    @test principalvars(M) == pvars
    @test principalvar(M, 2) == pvars[2]
    @test var(M) == 15.0
    @test tprincipalvar(M) == 12.0
    @test tresidualvar(M) == 3.0
    @test principalratio(M) == 0.8
    @test isapprox(loadings(M), l, atol = 0.001)

    @test predict(M, X[:,1]) ≈ P'X[:,1]
    @test predict(M, X) ≈ P'X

    @test reconstruct(M, Y[:,1]) ≈ P * Y[:,1]
    @test reconstruct(M, Y) ≈ P * Y

    ## PCA with non-zero mean

    mval = rand(rng, 5)
    M = PCA(mval, P, pvars, 15.0)

    @test size(M) == (5,3)
    @test mean(M) == mval
    @test projection(M) == P
    @test principalvars(M) == pvars
    @test principalvar(M, 2) == pvars[2]
    @test var(M) == 15.0
    @test tprincipalvar(M) == 12.0
    @test tresidualvar(M) == 3.0
    @test principalratio(M) == 0.8

    @test predict(M, X[:,1]) ≈ P' * (X[:,1] .- mval)
    @test predict(M, X) ≈ P' * (X .- mval)

    @test reconstruct(M, Y[:,1]) ≈ P * Y[:,1] .+ mval
    @test reconstruct(M, Y) ≈ P * Y .+ mval

    ## prepare training data

    d = 5
    n = 1000

    R = collect(qr(randn(rng, d, d)).Q)
    @test R'R ≈ Matrix(I, 5, 5)
    rmul!(R, Diagonal(sqrt.([0.5, 0.3, 0.1, 0.05, 0.05])))

    X = R'randn(rng, 5, n) .+ randn(rng, 5)
    mval = vec(mean(X, dims=2))
    Z = X .- mval

    C = cov(X, dims=2)
    pvs0 = sort(eigvals(C); rev=true)
    tv = sum(diag(C))

    ## pcacov (default using cov when d < n)

    M = fit(PCA, X)
    P = projection(M)
    pvs = principalvars(M)

    @test size(M) == (5,5)
    @test mean(M) == mval
    @test P'P ≈ Matrix(I, 5, 5)
    @test C*P ≈ P*Diagonal(pvs)
    @test issorted(pvs; rev=true)
    @test pvs ≈ pvs0
    @test var(M) ≈ tv
    @test sum(pvs) ≈ var(M)
    @test reconstruct(M, predict(M, X)) ≈ X

    M = fit(PCA, X; mean=mval)
    @test projection(M) ≈ P

    M = fit(PCA, Z; mean=0)
    @test projection(M) ≈ P

    M = fit(PCA, X; maxoutdim=3)
    P = projection(M)

    @test size(M) == (5,3)
    @test P'P ≈ Matrix(I, 3, 3)
    @test issorted(pvs; rev=true)

    M = fit(PCA, X; pratio=0.85)

    @test size(M) == (5,3)
    @test P'P ≈ Matrix(I, 3, 3)
    @test issorted(pvs; rev=true)

    ## pcastd

    M = fit(PCA, X; method=:svd)
    P = projection(M)
    pvs = principalvars(M)

    @test size(M) == (5,5)
    @test mean(M) == mval
    @test P'P ≈ Matrix(I, 5, 5)
    @test isapprox(C*P, P*Diagonal(pvs), atol=1.0e-3)
    @test issorted(pvs; rev=true)
    @test isapprox(pvs, pvs0, atol=1.0e-3)
    @test isapprox(var(M), tv, atol=1.0e-2)
    @test sum(pvs) ≈ var(M)
    @test reconstruct(M, predict(M, X)) ≈ X

    M = fit(PCA, X; method=:svd, mean=mval)
    @test projection(M) ≈ P

    M = fit(PCA, Z; method=:svd, mean=0)
    @test projection(M) ≈ P

    M = fit(PCA, X; method=:svd, maxoutdim=3)
    P = projection(M)

    @test size(M) == (5,3)
    @test P'P ≈ Matrix(I, 3, 3)
    @test issorted(pvs; rev=true)

    M = fit(PCA, X; method=:svd, pratio=0.85)

    @test size(M) == (5,3)
    @test P'P ≈ Matrix(I, 3, 3)
    @test issorted(pvs; rev=true)

    # Different data types
    # --------------------

    XX = convert.(Float32, X)
    YY = convert.(Float32, Y)
    p = 0.085
    pp = convert(Float32, p)

    MM = fit(PCA, XX; maxoutdim=3)

    # mix types
    fit(PCA, X ; pratio=pp)
    fit(PCA, XX ; pratio=p)
    fit(PCA, XX ; pratio=pp)
    predict(M, XX)
    predict(MM, X)
    reconstruct(M, YY)
    reconstruct(MM, Y)

    # type consistency
    for func in (mean, projection, principalvars, tprincipalvar, tresidualvar, var, principalratio)
        @test eltype(func(M)) == Float64
        @test eltype(func(MM)) == Float32
    end

    # views
    M = fit(PCA, view(X, :, 1:500), pratio=0.85)

    # sparse
    @test_throws AssertionError fit(PCA, SparseArrays.sprandn(rng, 100d, n, 0.6))

end


include("convex_hull.jl")

using FactCheck

# Test 1
# convex_hull function accepts only array of tuples
# If other values was passed, then it raises an exception
facts("Validate convex_hull function arguments") do
  @fact main("STRING") --> "Please provide array of tuples where each element is Int64"
  @fact main(("string","string")) --> "Please provide array of tuples where each element is Int64"
  @fact main([(1.23, 5.55)]) --> "Please provide array of tuples where each element is Int64"
end

# Test 2
# If the number of points is smaller than 3, then return the input
facts("Return input when it contains less than 3 points") do
  @fact convex_hull!([(0,0)]) --> [(0,0)]
  @fact convex_hull!([(0,1),(0,2)]) --> [(0,1),(0,2)]
end

# Test 3
# Detect direction based on the three points
facts("Detect direction based on three points") do
  @fact counter_clockwise((0,1), (1,1),  (2,2))   > 0 --> true
  @fact counter_clockwise((0,1), (-1,-1),(-2,-2)) < 0 --> true
  @fact counter_clockwise((0,1), (0,2),  (0,3))   --> 0
end

# Test 4
facts("Generates array of points and returns it") do
  @fact size(generate_points(10, -30, 30),1) == 10 --> true
  @fact typeof(generate_points(10,-30,30)) == Array{Tuple{Int64,Int64},1} --> true
end

# Test 5
# Check if the convex hull part is build correctly

# general case:
input1 = [(2,3),(3,0),(6,6),(5,2),(5,4),(1,1),(4,5)]
# a case with collinear point:
input2 = [(2,3),(3,0),(6,6),(3,4),(5,2),(5,4),(1,1),(4,5)]

# sorted imputs
sorted_input1 = sort!(union(input1), alg = HeapSort)
sorted_input2 = sort!(union(input2), alg = HeapSort)

# result array
result1 = []
result2 = []

build_convex_hull_part!(result1, sorted_input1, 1:1:length(sorted_input1))
build_convex_hull_part!(result2, sorted_input2, length(sorted_input2):-1:1)

facts("Checking build_convex_hull_part method") do
    @fact result1 --> [(1,1),(2,3),(4,5),(6,6)]
    @fact result2 --> [(6,6),(5,2),(3,0),(1,1)]
end

# Test 6
# Check if convex hull is build correctly
input1 = [(0,3), (1,1), (2,2), (4,4), (0,0), (1,2), (3,1), (3,3)]
input2 = [(3,-3),(-5,-2),(0,-4),(4,3),(-2,0),(2,3),(1,6),(0,-4),(-5,5),(3,3),(-3,0),(6,0),(5,3),(2,2),(6,0),(4,-5),(-2,3),(4,5),(-3,-1),(-3,2)]
input3 = [(0,-4),(-5,-5),(-9,1),(-6,-1),(2,-10),(-4,1),(-1,-4),(2,7),(0,3),(-9,2),(3,-10),(9,5),(8,3),(10,-2),(-6,-7),(0,-9),(-10,-9),(-3,8),(9,4),(2,1),(-8,0),(-6,-9),(-10,3),(6,-7),(2,-2),(4,-5),(0,9),(-2,3),(8,-4),(5,-5)]

result1 = convex_hull!(input1)
result2 = convex_hull!(input2)
result3 = convex_hull!(input3)

facts("Check builing complete convex hull") do
  @fact result1 --> [(0,0),(0,3),(4,4),(3,1)]
  @fact result2 --> [(-5,-2),(-5,5),(1,6),(4,5),(5,3),(6,0),(4,-5),(0,-4)]
  @fact result3 --> [(-10,-9),(-10,3),(-3,8),(0,9),(9,5),(10,-2),(6,-7),(3,-10),(2,-10)]
end

println("All tests passed")


### A Pluto.jl notebook ###
# v0.14.2

using Markdown
using InteractiveUtils

# ╔═╡ 14d9371c-9a19-11eb-2c1a-8b680938c549
begin
	using Cairo
	using Compose
	using CSV
	using LightGraphs
	using GraphPlot
	using StatsPlots
	using DataFrames
	using Statistics
	using Base.Threads
end

# ╔═╡ 7c9d6b8c-68f0-4481-b042-b51e217a8abd
md"""
# Stochastic Streets Riddler Solution

In this post, we'll solve the [Riddler Problem](https://fivethirtyeight.com/features/can-you-navigate-the-one-way-streets/) from April 9th	 2021.

In this problem, we're asked to consider a square gridded city with randomly oriented one way streets, and asked to determine the probability of being able to navigate from the upper left-hand corner of the city to the lower right-hand corner on the road network.

In the original problem statement, we're asked to consider a city 2 blocks across by 2 blocks high... but I will consider cities of arbitrary size. In this notebook, I provide a simulation for arbitrary sized cities, and a lazy exact solution for small cities.

*Note*: Since it makes notation a bit simpler, I will use $n$ to denote the number of intersections across the city. That is $n= n_\text{blocks} + 1$ where $n_\text{blocks}$ is the width of the city in blocks.
"""

# ╔═╡ 891b9894-11b0-405f-aa85-232cc635db9a
md"""
## Estimate Probability via Simulation

First, we'll compute the probability that a path exists from the top left-hand corner of the city to the lower right-hand corner of the city via simulation.
"""

# ╔═╡ 03792d01-268a-451b-8d48-13afdf6f379f
function city_has_path(g, start=1, finish=nv(g))
	has_path_start_finish = has_path(g, start, finish)

	return has_path_start_finish
end

# ╔═╡ 713a8126-0a4d-4373-ad57-6872e6aed059
md"""
### Example Random Street Layouts
"""

# ╔═╡ caf46b55-87cf-4351-9b6a-f729cce80251
md"""
### Simulation Results
"""

# ╔═╡ 5327b953-268a-4b6d-9bf1-ce513f666b56
md"""
## Lazy Analytic Solution (Small Cities)

Now, we'll compute the probability that a path exists from the top left-hand corner of the city to the lower right-hand corner of the city via simulation. We do this in a brute force combinatoric way... by simply enumerating all possible city layouts, and counting how many have the desired path. Since the number of paths grows exponentially in city size, we'll only compute the result for small cities, i.e. $n = 2,3,4$.

There's probably a smarter way to do this... but I couldn't think of it 🙂
"""

# ╔═╡ fd5d8a0c-9a36-11eb-37af-b323787108b5
md"""
## Utility Functions and Package Imports
"""

# ╔═╡ 38abf0ba-c6ea-4f45-acf1-b92eebe462a8
function neighbors(i, j, n)
	ilo = max(1, i-1)
	ihi = min(n, i+1)

	jlo = max(1, j-1)
	jhi = min(n, j+1)

	return vcat(
		[(i_neighb, j) for i_neighb = ilo:ihi if (i_neighb, j) != (i, j)],
		[(i, j_neighb) for j_neighb = jlo:jhi if (i, j_neighb) != (i, j)]
	)
end

# ╔═╡ b5b7e758-e4de-4666-976b-583a732e8904
function tuple_to_idx(i, j, n)
	return n * (i - 1) + j
end

# ╔═╡ f05db774-9741-475b-a6fc-5df56154faa4
function random_street_graph(n)
	g = SimpleDiGraph(n^2)

	# create set of all adjacent nodes
	adjacencies = Set([Set([(i, j), neighb]) for i = 1:n for j = 1:n for neighb
						in neighbors(i, j, n)])

	for adjacent_pair in adjacencies
		# unpack nodes and get their indices
		(i1, j1), (i2, j2) = adjacent_pair

		k1 = tuple_to_idx(i1, j1, n)
		k2 = tuple_to_idx(i2, j2, n)

		# connect index of node k1 to k2 if 1, else k2 to k1
		direction = rand((0,1))

		# connect node k1 to k2 in the indicated direction
		if direction > 0
			add_edge!(g, k1, k2)
		else
			add_edge!(g, k2, k1)
		end
	end

	return g
end

# ╔═╡ fb2302ca-8b2c-4e23-8dc6-36cc353b16cc
function simulate_city_planning(n, n_simulations)
	rslts = zeros(Bool, n_simulations)

	Threads.@threads for i ∈ 1:n_simulations
		g = random_street_graph(n)

		has_path = city_has_path(g)
		rslts[i, :] = [has_path]
	end

	return DataFrame(has_path=rslts[:,1])
end

# ╔═╡ 285ee664-4c13-4a0b-8c0e-c9b9e5120078
begin
	n_sim = 10^6
	n_max = 20
	prob_sim_rslts = zeros(n_max - 1)

	for n = 2:n_max
		df_city_planning_sim = simulate_city_planning(n, n_sim)
		prob_sim_rslts[n-1] = sum(df_city_planning_sim[!, :has_path]) / n_sim
	end

	df_prob_sim = DataFrame(city_width = 2:n_max,
							prob_path = prob_sim_rslts)
end

# ╔═╡ cec180f1-68f3-42bc-a5d1-4a7eeca01ab4
CSV.write("./outputs/prob_path_sim_rslts.csv", df_prob_sim)

# ╔═╡ 6d86c2cc-e62d-46b2-a844-343dbe083eab
bar_plot_ref = @df df_prob_sim bar(
	:city_width,
	:prob_path,
	title="Probability of Navigable Path From Top Left to Bottom Right",
	xlabel="City Width (n)",
	legend=false
)

# ╔═╡ eb45d228-befd-46bc-939d-758fb289fd4a
png(bar_plot_ref, "./outputs/prob_path.png")

# ╔═╡ 377592fd-d9f5-4514-a488-2d9f35c4e822
g = random_street_graph(3)

# ╔═╡ caebbad7-460b-44d3-b6a6-0ff9d4684fe8
begin
	nn = Int(sqrt(nv(g)))

	street_plot = gplot(g,
		[j for i = 1:nn for j=1:nn],
		[i for i = 1:nn for j=1:nn],
		nodelabel=["$((i-1)*nn+j)" for i = 1:nn for j=1:nn],
	)
end

# ╔═╡ 019727cd-8aaa-4fc5-bd7e-f90b7e3269d5
draw(PNG("./outputs/street_layout_example_1.png", 16cm, 16cm), street_plot)

# ╔═╡ 3fda5c44-70f3-4202-861d-9c5c4a7415e3
city_has_path(g)

# ╔═╡ 943d0cb8-8c4d-4c4f-9119-7ad84014fc0d
g2 = random_street_graph(3)

# ╔═╡ ccec4661-f080-402d-90c1-0d0ac7238e23
begin
	nn2 = Int(sqrt(nv(g2)))

	street_plot2 = gplot(g2,
		[j for i = 1:nn2 for j=1:nn2],
		[i for i = 1:nn2 for j=1:nn2],
		nodelabel=["$((i-1)*nn2+j)" for i = 1:nn2 for j=1:nn2],
	)
end

# ╔═╡ 9a95a271-f4be-48e9-9b04-57ae7976dda9
draw(PNG("./outputs/street_layout_example_2.png", 16cm, 16cm), street_plot2)

# ╔═╡ dede6d7e-362a-49d7-9223-12b16b0bcdf9
city_has_path(g2)

# ╔═╡ 303cac4d-d222-414a-9b58-db2da01c3e9f
function count_connected(n = 3, start=(1,1), finish=(n,n))

	adjacencies = Set([Set([(i, j), neighb]) for i = 1:n for j = 1:n
					                         for neighb = neighbors(i, j, n)])
	n_adj = length(adjacencies)
	n_layouts = 2^n_adj
	n_connected = 0

	for i_layout = 0:(n_layouts-1)
		# use the binary representation of i_layout to determine connections
		#  1 represents forward connection
		#  0 represents reverse connection
		layout = [parse(Int, ss) for ss ∈ bitstring(i_layout)[end-n_adj+1:end]]

		# create a graph
		g = SimpleDiGraph(n^2)

		# connect the graph
		for (i_adj, adjacent_pair) in enumerate(adjacencies)

			# get the adjacent pair as indices
			(i1, j1), (i2, j2) = adjacent_pair
			k1 = tuple_to_idx(i1, j1, n)
			k2 = tuple_to_idx(i2, j2, n)

			# get the direction for the connection from k1 to k2 from layout
			direction = layout[i_adj]

			# connect node k1 to k2 in the indicated direction
			if direction == 1
				add_edge!(g, k1, k2)
			else
				add_edge!(g, k2, k1)
			end
		end

		# is there a path from start to finish?
		if has_path(g, 1, nv(g))
			n_connected += 1
		end
	end

	return n_connected / n_layouts
end

# ╔═╡ 27f9b76a-b998-4fc4-b89d-baee19104f25
begin
	n_max_analytic = 4
	analytic_rslts = zeros(n_max_analytic - 1)

	for n = 2:n_max_analytic
		analytic_rslts[n-1] = count_connected(n)
	end

	df_analytic = DataFrame(city_width = 2:n_max_analytic,
							prob_path = analytic_rslts)
end

# ╔═╡ 06bb3a39-49aa-42ad-988b-c33671fe1bdb
CSV.write("./outputs/prob_path_analytic_rslts.csv", df_analytic)

# ╔═╡ c2c4c421-8d79-498d-b246-4c1e94f5381b
function idx_to_tuple(k, n)
	i = div((k - 1), n) + 1

	j = mod(k, n)
	if j == 0
		j = n
	end

	return (i, j)
end

# ╔═╡ Cell order:
# ╟─7c9d6b8c-68f0-4481-b042-b51e217a8abd
# ╟─891b9894-11b0-405f-aa85-232cc635db9a
# ╠═f05db774-9741-475b-a6fc-5df56154faa4
# ╠═03792d01-268a-451b-8d48-13afdf6f379f
# ╠═fb2302ca-8b2c-4e23-8dc6-36cc353b16cc
# ╟─713a8126-0a4d-4373-ad57-6872e6aed059
# ╟─377592fd-d9f5-4514-a488-2d9f35c4e822
# ╟─caebbad7-460b-44d3-b6a6-0ff9d4684fe8
# ╟─019727cd-8aaa-4fc5-bd7e-f90b7e3269d5
# ╠═3fda5c44-70f3-4202-861d-9c5c4a7415e3
# ╟─943d0cb8-8c4d-4c4f-9119-7ad84014fc0d
# ╟─ccec4661-f080-402d-90c1-0d0ac7238e23
# ╠═dede6d7e-362a-49d7-9223-12b16b0bcdf9
# ╟─9a95a271-f4be-48e9-9b04-57ae7976dda9
# ╟─caf46b55-87cf-4351-9b6a-f729cce80251
# ╠═285ee664-4c13-4a0b-8c0e-c9b9e5120078
# ╟─cec180f1-68f3-42bc-a5d1-4a7eeca01ab4
# ╟─6d86c2cc-e62d-46b2-a844-343dbe083eab
# ╟─eb45d228-befd-46bc-939d-758fb289fd4a
# ╟─5327b953-268a-4b6d-9bf1-ce513f666b56
# ╠═303cac4d-d222-414a-9b58-db2da01c3e9f
# ╠═27f9b76a-b998-4fc4-b89d-baee19104f25
# ╟─06bb3a39-49aa-42ad-988b-c33671fe1bdb
# ╟─fd5d8a0c-9a36-11eb-37af-b323787108b5
# ╠═38abf0ba-c6ea-4f45-acf1-b92eebe462a8
# ╠═b5b7e758-e4de-4666-976b-583a732e8904
# ╠═c2c4c421-8d79-498d-b246-4c1e94f5381b
# ╠═14d9371c-9a19-11eb-2c1a-8b680938c549


# -*- coding: utf-8 -*-
# ---
# jupyter:
#   jupytext:
#     cell_metadata_json: true
#     formats: ipynb,jl:light
#     text_representation:
#       extension: .jl
#       format_name: light
#       format_version: '1.5'
#       jupytext_version: 1.11.3
#   kernelspec:
#     display_name: Julia 1.6.2
#     language: julia
#     name: julia-1.6
# ---

# + [markdown] {"slideshow": {"slide_type": "slide"}}
# # Transport Streamfunction And Divergent Component
#
# Here we apply such [Helmholtz Decomposition](https://en.wikipedia.org/wiki/Helmholtz_decomposition) to vertically integrated transports defined over a Global Ocean model (`C-grid`).
#
# 1. read vertically integrated transport from file
# 2. compute its convergence and apply land mask
# 3. decompose into rotational and divergent components
# 4. derive streamfunction from the rotational component

# + [markdown] {"slideshow": {"slide_type": "slide"}}
# ### Read Transports & Grid From File
#
# 1. pre-requisites
# 2. read variables

# + {"slideshow": {"slide_type": "subslide"}, "cell_style": "center"}
using MeshArrays, Plots, Statistics
pth=dirname(pathof(MeshArrays))
include(joinpath(pth,"../examples/Plots.jl"))
include("helper_functions.jl")

pth=MeshArrays.GRID_LLC90
γ=GridSpec("LatLonCap",pth)
Γ=GridLoad(γ;option="full")
(Tx,Ty,τx,τy,η)=trsp_read("LatLonCap",pth);

# +
μ =Γ.hFacC[:,1]
μ[findall(μ.>0.0)].=1.0
μ[findall(μ.==0.0)].=NaN

lon=[i for i=-179.:2.0:179., j=-89.:2.0:89.]
lat=[j for i=-179.:2.0:179., j=-89.:2.0:89.]

#(f,i,j,w)=InterpolationFactors(Γ,vec(lon),vec(lat))
#λ=(lon=lon,lat=lat,f=f,i=i,j=j,w=w);
#df = DataFrame(f=λ.f[:], i=λ.i[:], j=λ.j[:], w=Float32.(λ.w[:]));
#CSV.write("interp_coeffs.csv", df)

df=DataFrame(CSV.File("interp_coeffs.csv"))
λ=(f=reshape(df.f,length(lon[:]),4), i=reshape(df.i,length(lon[:]),4),
    j=reshape(df.j,length(lon[:]),4), w=reshape(df.w,length(lon[:]),4));

# + [markdown] {"slideshow": {"slide_type": "slide"}}
# ### Helmholtz Decomposition
#
# 1. convergence & land mask
# 2. compute scalar potential
# 3. subtract divergent component
# 4. compute vector potential / streamfunction
#
# _note: masking avoids singularities related to isolated canyons_

# + {"slideshow": {"slide_type": "subslide"}}
#convergence & land mask
TrspCon=μ.*convergence(Tx,Ty)

#scalar potential
TrspPot=ScalarPotential(TrspCon)

#Divergent transport component
(TxD,TyD)=gradient(TrspPot,Γ)
TxD=TxD.*Γ.DXC
TyD=TyD.*Γ.DYC

#Rotational transport component
TxR = Tx-TxD
TyR = Ty-TyD

#vector Potential
TrspPsi=VectorPotential(TxR,TyR,Γ);

# + [markdown] {"slideshow": {"slide_type": "skip"}}
# ### Verify The Results Consistency
#
# _Next plot below normally shows that TrspCon-tmpCon is negligible compared with TrspCon_

# + {"slideshow": {"slide_type": "skip"}, "cell_style": "center"}
tmpCon=convergence(TxD,TyD)
tmp1=TrspCon[3]
tmp2=tmp1[findall(isfinite.(tmp1))]
errCon=1/sqrt(mean(tmp2.^2)).*(tmpCon[3]-TrspCon[3])
errCon[findall(isnan.(errCon))].=0.0
extrema(errCon)

# + [markdown] {"slideshow": {"slide_type": "slide"}}
# ### Map Out Streamfunction And Scalar Potential
#
# _In plots below we interpolate MeshArrays to a regular global grid_

# + {"slideshow": {"slide_type": "subslide"}, "cell_style": "center"}
TrspPsiI=Interpolate(1e-6*μ*TrspPsi,λ.f,λ.i,λ.j,λ.w)
contourf(vec(lon[:,1]),vec(lat[1,:]),
    TrspPsiI,title="Streamfunction",clims=(-50,50))

# + {"slideshow": {"slide_type": "subslide"}, "cell_style": "center"}
TrspPotI=Interpolate(1e-6*μ*TrspPot,λ.f,λ.i,λ.j,λ.w)
contourf(vec(lon[:,1]),vec(lat[1,:]),
        TrspPotI,title="Scalar Potential",clims=(-0.4,0.4))
# -


"""
    nextoffset = flatten!(v, domain, val, offset=1)
Flatten `val`, taking values in `domain`, and transform to Cartesian coordinates. Insert the transformed value into `v` starting at `offset`. Returns the next offset to use.
"""
flatten!(v, domain, val) = flatten!(v, domain, val, 1)

# recursively apply `flatten`
function flatten!(v, ::Type{T}, val::T, offset) where {T}
  map(zip(fieldnames(T), domains(T))) do (fieldname, domain)
    offset = flatten!(v, domain, getfield(val, fieldname), offset)
  end
  return offset
end

function flatten!(v, ::Type{T}, val::T, offset) where {T<:Real}
  v[offset] = val
  return offset+1
end

"""
    v = flatten(domain, val)
Flatten `val`, taking values in `domain`, and transform to Cartesian coordinates. Returns a vector `v` of length `dimension(domain)`.
"""
function flatten(domain, val)
  v = zeros(Float64, dimension(domain))
  flatten!(v, domain, val)
  v
end

"""
    val, nextoffset = unflatten(v, domain, offset)
Transform from `v`, starting at `offset`, stored in Cartesian coordinates to value in `domain`. Returns the transformed value and the next offset.
"""
function unflatten end

"""
    val = unflatten(v, domain)
Transform from `v` stored in Cartesian coordinates to value in `domain`. Returns the transformed value.
"""
unflatten(v, domain) = first(unflatten(v, domain, 1))

function unflatten(v, ::Type{T}, offset) where {T}
  vals = map(domains(T)) do domain
    val, offset = unflatten(v, domain, offset)
    val
  end
  return T(vals...), offset
end

function unflatten(v, ::Type{T}, offset) where {T<:Real}
  v[offset], offset+1
end