import LinearAlgebra: mul!, ldiv!
import Base: getindex, iterate

using SparseArrays

struct DiagonalIndices{Tv, Ti <: Integer}
    matrix::SparseMatrixCSC{Tv,Ti}
    diag::Vector{Ti}

    function DiagonalIndices{Tv,Ti}(A::SparseMatrixCSC{Tv,Ti}) where {Tv,Ti}
        # Check square?
        diag = Vector{Ti}(undef, A.n)

        for col = 1 : A.n
            r1 = A.colptr[col]
            r2 = A.colptr[col + 1] - 1
            r1 = searchsortedfirst(A.rowval, col, r1, r2, Base.Order.Forward)
            if r1 > r2 || A.rowval[r1] != col || iszero(A.nzval[r1])
                throw(LinearAlgebra.SingularException(col))
            end
            diag[col] = r1
        end

        new(A, diag)
    end
end

DiagonalIndices(A::SparseMatrixCSC{Tv,Ti}) where {Tv,Ti} = DiagonalIndices{Tv,Ti}(A)

function ldiv!(y::AbstractVector{Tv}, D::DiagonalIndices{Tv,Ti}, x::AbstractVector{Tv}) where {Tv,Ti}
    @inbounds for row = 1 : D.matrix.n
        y[row] = x[row] / D.matrix.nzval[D.diag[row]]
    end
    y
end

@inline getindex(d::DiagonalIndices, i::Int) = d.diag[i]

struct FastLowerTriangular{Tv,Ti}
    matrix::SparseMatrixCSC{Tv,Ti}
    diag::DiagonalIndices{Tv,Ti}
end

struct FastUpperTriangular{Tv,Ti}
    matrix::SparseMatrixCSC{Tv,Ti}
    diag::DiagonalIndices{Tv,Ti}
end

struct StrictlyUpperTriangular{Tv,Ti}
    matrix::SparseMatrixCSC{Tv,Ti}
    diag::DiagonalIndices{Tv,Ti}
end

struct StrictlyLowerTriangular{Tv,Ti}
    matrix::SparseMatrixCSC{Tv,Ti}
    diag::DiagonalIndices{Tv,Ti}
end

struct OffDiagonal{Tv,Ti}
    matrix::SparseMatrixCSC{Tv,Ti}
    diag::DiagonalIndices{Tv,Ti}
end

"""
Forward substitution for the FastLowerTriangular type
"""
function forward_sub!(F::FastLowerTriangular, x::AbstractVector)
    A = F.matrix

    @inbounds for col = 1 : A.n

        # Solve for diagonal element
        idx = F.diag[col]
        x[col] /= A.nzval[idx]

        # Substitute next values involving x[col]
        for i = idx + 1 : (A.colptr[col + 1] - 1)
            x[A.rowval[i]] -= A.nzval[i] * x[col]
        end
    end

    x
end

"""
Forward substitution
"""
function forward_sub!(α, F::FastLowerTriangular, x::AbstractVector, β, y::AbstractVector)
    A = F.matrix

    @inbounds for col = 1 : A.n

        # Solve for diagonal element
        idx = F.diag[col]
        x[col] = α * x[col] / A.nzval[idx] + β * y[col]

        # Substitute next values involving x[col]
        for i = idx + 1 : (A.colptr[col + 1] - 1)
            x[A.rowval[i]] -= A.nzval[i] * x[col]
        end
    end

    x
end

"""
Backward substitution for the FastUpperTriangular type
"""
function backward_sub!(F::FastUpperTriangular, x::AbstractVector)
    A = F.matrix

    @inbounds for col = A.n : -1 : 1

        # Solve for diagonal element
        idx = F.diag[col]
        x[col] = x[col] / A.nzval[idx]

        # Substitute next values involving x[col]
        for i = A.colptr[col] : idx - 1
            x[A.rowval[i]] -= A.nzval[i] * x[col]
        end
    end

    x
end

function backward_sub!(α, F::FastUpperTriangular, x::AbstractVector, β, y::AbstractVector)
    A = F.matrix

    @inbounds for col = A.n : -1 : 1

        # Solve for diagonal element
        idx = F.diag[col]
        x[col] = α * x[col] / A.nzval[idx] + β * y[col]

        # Substitute next values involving x[col]
        for i = A.colptr[col] : idx - 1
            x[A.rowval[i]] -= A.nzval[i] * x[col]
        end
    end

    x
end

"""
Do mul! with the off-diagonal elements of a matrix.
"""
function mul!(α::T, O::OffDiagonal, x::AbstractVector, β::T, y::AbstractVector) where {T}
    # Specialize for β = 0 and β = 1
    A = O.matrix

    if β != one(T)
        if iszero(β)
            fill!(y, zero(T))
        else
            lmul!(β, y)
        end
    end

    @inbounds for col = 1 : A.n
        αx = α * x[col]
        diag_index = O.diag[col]
        for j = A.colptr[col] : diag_index - 1
            y[A.rowval[j]] += A.nzval[j] * αx
        end
        for j = diag_index + 1 : A.colptr[col + 1] - 1
            y[A.rowval[j]] += A.nzval[j] * αx
        end
    end

    y
end

"""
Computes z := α * U * x + β * y. Because U is StrictlyUpperTriangular
one can set z = x and update x in-place as x := α * U * x + β * y.
"""
function gauss_seidel_multiply!(α, U::StrictlyUpperTriangular, x::AbstractVector, β, y::AbstractVector, z::AbstractVector)
    A = U.matrix

    for col = 1 : A.n
        αx = α * x[col]
        diag_index = U.diag[col]
        @inbounds for j = A.colptr[col] : diag_index - 1
            z[A.rowval[j]] += A.nzval[j] * αx
        end
        z[col] = β * y[col]
    end
    z
end

"""
Computes z := α * L * x + β * y. Because A is StrictlyLowerTriangular
one can set z = x and update x in-place as x := α * L * x + β * y.
"""
function gauss_seidel_multiply!(α, L::StrictlyLowerTriangular, x::AbstractVector, β, y::AbstractVector, z::AbstractVector)
    A = L.matrix

    for col = A.n : -1 : 1
        αx = α * x[col]
        z[col] = β * y[col]
        @inbounds for j = L.diag[col] + 1 : (A.colptr[col + 1] - 1)
            z[A.rowval[j]] += A.nzval[j] * αx
        end
    end
    z
end

##
## Jacobi
##

mutable struct JacobiIterable{T <: Number, solT, vecT, rhsT}
    O::OffDiagonal

    x::solT
    next::vecT
    b::rhsT

    maxiter::Int
end

start(::JacobiIterable) = 1
done(j::JacobiIterable, iteration::Int) = iteration > j.maxiter
function iterate(j::JacobiIterable{T}, iteration::Int=start(j)) where {T}
    if done(j, iteration) return nothing end
    # tmp = D \ (b - (A - D) * x)
    copyto!(j.next, j.b)
    mul!(-one(T), j.O, j.x, one(T), j.next)
    ldiv!(j.x, j.O.diag, j.next)

    nothing, iteration + 1
end

jacobi_iterable(x::AbstractVector, A::SparseMatrixCSC, b::AbstractVector; maxiter::Int = 10) =
    JacobiIterable{eltype(x), typeof(x), typeof(x), typeof(b)}(
        OffDiagonal(A, DiagonalIndices(A)), x, similar(x), b, maxiter)

"""
    jacobi!(x, A::SparseMatrixCSC, b; maxiter=10) -> x

Performs exactly `maxiter` Jacobi iterations.

Allocates a temporary vector and precomputes the diagonal indices.

Throws `LinearAlgebra.SingularException` when the diagonal has a zero. This check
is performed once beforehand.
"""
function jacobi!(x::AbstractVector, A::SparseMatrixCSC, b::AbstractVector; maxiter::Int=10)
    iterable = jacobi_iterable(x, A, b, maxiter = maxiter)
    for item = iterable end
    iterable.x
end

##
## Gauss-Seidel
##

mutable struct GaussSeidelIterable{solT, rhsT}
    U::StrictlyUpperTriangular
    L::FastLowerTriangular

    x::solT
    b::rhsT

    maxiter::Int
end

function gauss_seidel_iterable(x::AbstractVector, A::SparseMatrixCSC, b::AbstractVector; maxiter::Int = 10)
    D = DiagonalIndices(A)
    GaussSeidelIterable(StrictlyUpperTriangular(A, D), FastLowerTriangular(A, D), x, b, maxiter)
end

start(::GaussSeidelIterable) = 1
done(g::GaussSeidelIterable, iteration::Int) = iteration > g.maxiter
function iterate(g::GaussSeidelIterable, iteration::Int=start(g))
    if done(g, iteration) return nothing end
    # x ← L \ (-U * x + b)
    T = eltype(g.x)
    gauss_seidel_multiply!(-one(T), g.U, g.x, one(T), g.b, g.x)
    forward_sub!(g.L, g.x)

    nothing, iteration + 1
end

"""
    gauss_seidel!(x, A::SparseMatrixCSC, b; maxiter=10) -> x

Performs exactly `maxiter` Gauss-Seidel iterations.

Works fully in-place, but precomputes the diagonal indices.

Throws `LinearAlgebra.SingularException` when the diagonal has a zero. This check
is performed once beforehand.
"""
function gauss_seidel!(x::AbstractVector, A::SparseMatrixCSC, b::AbstractVector; maxiter::Int = 10)
    iterable = gauss_seidel_iterable(x, A, b, maxiter = maxiter)
    for item = iterable end
    iterable.x
end

##
## SOR
##

mutable struct SORIterable{T, solT, vecT, rhsT, numT <: Real}
    U::StrictlyUpperTriangular
    L::FastLowerTriangular
    ω::numT

    x::solT
    next::vecT
    b::rhsT

    maxiter::Int
end

start(::SORIterable) = 1
done(s::SORIterable, iteration::Int) = iteration > s.maxiter
function iterate(s::SORIterable{T}, iteration::Int=start(s)) where {T}
    if done(s, iteration) return nothing end

    # next = b - U * x
    gauss_seidel_multiply!(-one(T), s.U, s.x, one(T), s.b, s.next)

    # next = ω * inv(L) * next + (1 - ω) * x
    forward_sub!(s.ω, s.L, s.next, one(T) - s.ω, s.x)

    # Switch current and next iterate
    s.x, s.next = s.next, s.x

    nothing, iteration + 1
end

function sor_iterable(x::AbstractVector, A::SparseMatrixCSC, b::AbstractVector, ω::Real; maxiter::Int = 10)
    D = DiagonalIndices(A)
    T = eltype(x)
    SORIterable{T,typeof(x),typeof(x),typeof(b),eltype(ω)}(
        StrictlyUpperTriangular(A, D), FastLowerTriangular(A, D), ω,
        x, similar(x), b, maxiter
    )
end

"""
    sor!(x, A::SparseMatrixCSC, b, ω::Real; maxiter=10)

Performs exactly `maxiter` SOR iterations with relaxation parameter `ω`.

Allocates a temporary vector and precomputes the diagonal indices.

Throws `LinearAlgebra.SingularException` when the diagonal has a zero. This check
is performed once beforehand.
"""
function sor!(x::AbstractVector, A::SparseMatrixCSC, b::AbstractVector, ω::Real; maxiter::Int = 10)
    iterable = sor_iterable(x, A, b, ω, maxiter = maxiter)
    for item = iterable end
    iterable.x
end

##
## SSOR
##

mutable struct SSORIterable{T, solT, vecT, rhsT, numT <: Real}
    sL::StrictlyLowerTriangular
    sU::StrictlyUpperTriangular
    L::FastLowerTriangular
    U::FastUpperTriangular
    ω::numT
    x::solT
    tmp::vecT
    b::rhsT
    maxiter::Int
end

function ssor_iterable(x::AbstractVector{T}, A::SparseMatrixCSC, b::AbstractVector, ω::Real; maxiter::Int = 10) where {T}
    D = DiagonalIndices(A)
    SSORIterable{T,typeof(x),typeof(x),typeof(b),typeof(ω)}(
        StrictlyLowerTriangular(A, D),
        StrictlyUpperTriangular(A, D),
        FastLowerTriangular(A, D),
        FastUpperTriangular(A, D),
        ω, x, similar(x), b, maxiter
    )
end

start(s::SSORIterable) = 1
done(s::SSORIterable, iteration::Int) = iteration > s.maxiter

function iterate(s::SSORIterable{T}, iteration::Int=start(s)) where {T}
    if done(s, iteration) return nothing end

    # tmp = b - U * x
    gauss_seidel_multiply!(-one(T), s.sU, s.x, one(T), s.b, s.tmp)

    # tmp = ω * inv(L) * tmp + (1 - ω) * x
    forward_sub!(s.ω, s.L, s.tmp, one(T) - s.ω, s.x)

    # x = b - L * tmp
    gauss_seidel_multiply!(-one(T), s.sL, s.tmp, one(T), s.b, s.x)

    # x = ω * inv(U) * x + (1 - ω) * tmp
    backward_sub!(s.ω, s.U, s.x, one(T) - s.ω, s.tmp)

    nothing, iteration + 1
end

"""
    ssor!(x, A::SparseMatrixCSC, b, ω::Real; maxiter=10)

Performs exactly `maxiter` SSOR iterations with relaxation parameter `ω`. Each iteration
is basically a forward *and* backward sweep of SOR.

Allocates a temporary vector and precomputes the diagonal indices.

Throws `LinearAlgebra.SingularException` when the diagonal has a zero. This check
is performed once beforehand.
"""
function ssor!(x::AbstractVector, A::SparseMatrixCSC, b::AbstractVector, ω::Real; maxiter::Int = 10)
    iterable = ssor_iterable(x, A, b, ω, maxiter = maxiter)
    for item = iterable end
    iterable.x
end


mutable struct ConstraintData
    nb_cstr::Int
    indices::Vector{Vector{Int}}
    coeffs::Vector{Vector{Float64}}
    lb::Vector{Float64}
    ub::Vector{Float64}
end

function ConstraintData(m)
    res = ConstraintData(
            length(m.linconstr),
            [map(v-> getfield(v, :col), CSTR.terms.vars) for CSTR in m.linconstr],
            [CSTR.terms.coeffs for CSTR in m.linconstr],
            [CSTR.lb for CSTR in m.linconstr],
            [CSTR.ub for CSTR in m.linconstr])

    to_discard = find(x-> x<=0 || x==Inf, res.ub)

    if !isempty(to_discard)
        res.nb_cstr = res.nb_cstr - length(to_discard)
        deleteat!(res.indices, to_discard)
        deleteat!(res.coeffs, to_discard)
        deleteat!(res.lb, to_discard)
        deleteat!(res.ub, to_discard)
    end

    res
end

isacover(C, a, b) = sum(a[C]) > b

function extend(C, a)
    m = maximum(a[C])
    res = copy(C)
    for i in eachindex(a)
        if a[i] >= m && !(i in res)
            push!(res, i)
        end
    end
    sort!(res)
end

function isminimal(C, a, b)
    sum(a[C]) - minimum(a[C]) <= b
end

function isastrongcover(C, a, b)
    !isminimal(C, a, b) && return false
    EC = extend(C, a)
    length(EC) == length(a) && return true
    aij1 = maximum(a[C])
    aij2 = maximum(a[setdiff(1:length(a), EC)])
    return sum(a[C]) - aij1 + aij2 <= b
end

function isviolated(C, x, card)
    sum(x[C]) > card + 1e-3
end

function separateHeur(x, a, b)
    order = sortperm((1 .- x)./ a)
    C = Int[]
    astar = b
    while !isempty(order)
        p = popfirst!(order)
        a[p] <= astar && push!(C, p)
        isempty(C) && return (Int[], 0)
        if isacover(C, a, b)
            ECI = extend(C, a)
            if isviolated(ECI, x, length(C)-1)
                return (C, length(C)-1)
            else
                kstar = argmax(a[C])
                astar = a[kstar]
                deleteat!(C, kstar)
            end
        end
    end
    Int[], 0
end

function find_cover_cuts(n::Node, cstrData, lift_covers)
    success = false
    for i = 1:cstrData.nb_cstr
        inds = cstrData.indices[i] #indices of variables involved in the constraint
        x = view(n.m.colVal, inds) #value of the variables in the current solution
        a = cstrData.coeffs[i] #coefficients of the variables in the current constraint
        b = cstrData.ub[i] #RHS of the constraint

        if dot(x, a) > b - 0.5 #if the constraint is activated
            C, card = separateHeur(x, a, b) #try to find a cover
            if !isempty(C) #if success, add it to the list of cover cuts
                success = true
                if lift_covers
                    if isastrongcover(C, a, b)
                        res = lifting_procedure(C, a, b, card)
                        @constraint(n.m, dot([JuMP.Variable(n.m, j) for j in inds], res) <= card)
                    else
                        if lifting_procedure_M2!(C, a, b)
                            res = lifting_procedure(C, a, b, card)
                            @constraint(n.m, dot([JuMP.Variable(n.m, j) for j in inds], res) <= card)
                        else
                            @constraint(n.m, sum(JuMP.Variable(n.m, j) for j in inds[extend(C, a)]) <= card)
                        end
                    end
                else
                    @constraint(n.m, sum(JuMP.Variable(n.m, j) for j in inds[extend(C, a)]) <= card)
                end
            end
        end
    end
   success
end

function find_cover_cuts(n::NodeParragh, cstrData, lift_covers)
    cuts = Set{Tuple{Vector{Int}, Int}}()
    lifted_cuts = Set{Tuple{Vector{Int}, Vector{Int}, Int}}()
    success = false
    for i = 1:cstrData.nb_cstr
        inds = cstrData.indices[i] #indices of variables involved in the constraint
        a = cstrData.coeffs[i] #coefficients of the variables in the current constraint
        b = cstrData.ub[i] #RHS of the constraint

        for ind_x in 1:length(n.x) #for each solution of the convex linear relaxation
            x =  view(n.x[ind_x], inds) #value of the variables in the solution

            if dot(x, a) > b - 0.5 #if the constraint is activated
                C, card = separateHeur(x, a, b) #try to find a cover
                if !isempty(C) #if success, add it to the list of cover cuts
                    success = true
                    if lift_covers
                        if isastrongcover(C, a, b)
                            res = lifting_procedure(C,  a, b, card)
                            push!(lifted_cuts, (inds, res, card))
                        else
                            if lifting_procedure_M2!(C, a, b)
                                res = lifting_procedure(C, a, b, card)
                            else
                                push!(cuts, (inds[extend(C, a)], card))
                            end
                        end
                        # @constraint(n.m, dot([JuMP.Variable(n.m, j) for j in inds], res) <= card)
                    else
                        push!(cuts, (inds[extend(C, a)], card))
                        # @constraint(n.m, sum(JuMP.Variable(n.m, j) for j in inds[extend(C, a)]) <= card)
                    end
                end
            end
        end
    end
    if success
        for (cut, card) in cuts
            @constraint(n.m, sum(JuMP.Variable(n.m, k) for k in cut) <= card)
        end
        for (inds, cut, card) in lifted_cuts
            @constraint(n.m, dot([JuMP.Variable(n.m, k) for k in inds], cut) <= card)
        end
    end
   success
end

function lifting_procedure(C, a, b, card)
    sorted_Ci = sort(C, by = x->a[x], rev=true)
    sorted_ECi = sort(extend(C, a), by = x->a[x], rev=true)
    q = card
    n0 = setdiff(1:length(a), sorted_ECi)
    aj = ones(Int, length(a))
    aj[n0] .= 0

    lb = a[sorted_Ci[1]]
    ub = lb + a[sorted_Ci[2]]
    for h = 2:q
        lb += a[sorted_Ci[h]]
        ub += a[sorted_Ci[h+1]]
        maximum(a) < lb && break
        nh = find(x-> lb <= x < ub, a)
        aj[nh] .= h
    end

    return aj

end

function lifting_procedure_M2!(C, a, b)
    EC = extend(C, a)
    N_minus_EC = setdiff(1:length(a), EC)
    isempty(N_minus_EC) && return false
    aij1, j1 = findmax(a[C])
    aij2 = maximum(a[N_minus_EC])
    temp = copy(a)
    Δ = sum(a[C]) - aij1 + aij2 - b
    while Δ > 0
        deleteat!(C, j1)
        copyto!(temp, a)
        temp[EC] .= typemax(Int)
        temp[temp .<= (aij2 - Δ)] .= typemax(Int)
        j = indmin(temp)
        push!(C, j) ; sort!(C)
        EC = extend(C, a)
        if length(EC) == length(a)
            isastrongcover(C, a, b) && return true
            return false
        end
        N_minus_EC = setdiff(1:length(a), EC)
        aij1, j1 = findmax(a[C])
        aij2 = maximum(a[N_minus_EC])
        Δ = sum(a[C]) - aij1 + aij2 - b
    end
    true
end


"""
```
transform_transition_matrices(m::AbstractDSGEModel, T::Matrix{T},
                              R::Matrix{T}, C::Matrix{T}) where {T<:AbstractFloat}
```

### Arguments

-`m`: the m object
-`T`, `R`, and `C`: matrices of the state transition equation

### Return values

- `TTT`, `RRR`, and `CCC`: transform transition equation matrices to accomodate states simulated
    between observations

### Description
"""
function transform_transition_matrices(m::AbstractDSGEModel, TT::Matrix{T},
                                       R::Matrix{T}, C::Vector{T};
                                       track_lag::Bool = true) where {T<:AbstractFloat}

    freq = get_setting(m, :state_simulation_freq)
    if track_lag
        TTT = zeros(size(TT, 1) * (freq+1), size(TT, 2) * (freq+1)) # pre-allocate
        RRR = zeros(size(R, 1) * (freq+1), size(R, 2) * freq) # pre-allocate
    else
        TTT = zeros(size(TT, 1) * freq, size(TT, 2) * freq) # pre-allocate
        RRR = zeros(size(R, 1) * freq, size(R, 2) * freq) # pre-allocate
    end

    # Construct powers of TT iteratively
    TT_powers = Dict{Int64, Matrix{T}}()
    for i = 1:freq
        if i == 1
            TT_powers[i] = TT
        else
            TT_powers[i] = TT_powers[i - 1] * TT
        end
    end

    # Insert directly into TTT
    if track_lag
        for i = 1:freq+1
            if i == 1
                TTT[1:size(TT,1), 1 + size(TT,2) * freq:end] = eye(TT) # track last lag
            else
                TTT[1 + size(TT,1)*(i-1):size(TT,1)*i, 1 + size(TT,2)*freq:end] = TT_powers[i-1]
            end
        end
    else
        for i = 1:freq
            TTT[1 + size(TT,1)*(i-1):size(TT,1)*i, 1 + size(TT,2)*(freq-1):end] = TT_powers[i]
        end
    end

    # Create RRR
    TR_powers = Dict{Int64, Matrix{T}}()
    TR_powers[0] = R

    # Construct T^i * R iteratively
    for i = 1:(freq - 1)
        TR_powers[i] = TT * TR_powers[i - 1]
    end

    # Fill in RRR and CCC
    if track_lag
        RRR[1+size(R,1):size(R, 1)*2, 1:size(R,2)] = TR_powers[0]
        for i in 3:freq+1
            for j = (i-2):-1:0
                RRR[1+size(R,1)*(i-1):size(R,1)*i, 1+size(R,2)*j:size(R,2)*(j+1)] = TR_powers[i-j-2]
            end
        end
        CCC = vcat(zeros(C), repeat(C, freq))
    else
        RRR[1:size(R, 1), 1:size(R,2)] = TR_powers[0]
        for i in 2:freq
            for j = (i-1):-1:0
                # inserts values from right nonzero blocks to left
                RRR[1+size(R,1)*(i-1):size(R,1)*i, 1+size(R,2)*j:size(R,2)*(j+1)] = TR_powers[i-j-1]
            end
        end
        CCC = repeat(C, freq) #vec(repeat(C, 1, freq))
    end

    return TTT, RRR, CCC
end


using OrdinaryDiffEq, Test, Random, LinearAlgebra, SparseArrays

# Parameters
Nc = 22
η = 1.0
κ = 1.0
T = (0.0,100.0)

# Matrix definitions
A = sparse(diagm(1 => sqrt.(Complex[1:Nc;])))
H = η*(A + A') - 1.0im*κ*A'*A

u0 = zeros(ComplexF64, Nc+1)
u0[1] = 1.0
function f_psos(du, u, t, p)
    du .= -1.0im*H*u
end

# Callback
rng = MersenneTwister(rand(UInt))
jumpnorm = Ref(rand(rng))
djumpnorm(x::Vector{ComplexF64}, t, integrator) = norm(x)^2 - (1-jumpnorm[])
function dojump(integrator)
    x = integrator.u
    t = integrator.t

    x .= normalize(A*x)
    jumpnorm[] = rand(rng)
end

cb = ContinuousCallback(djumpnorm,dojump)

prob = ODEProblem{true}(f_psos,u0,T)

sol_tot = []
Ntraj = 100
for i=1:Ntraj
    rng = MersenneTwister(rand(UInt))
    # Tweaking tolerances and dtmax also is not reliable
    sol = solve(prob,DP5(),save_everystep=true,callback=cb,
            abstol=1e-8,reltol=1e-6,dtmax=10)
    push!(sol_tot, sol)
end

# This number has to be η^2/κ^2 in steady-state; all trajectories should converge there
n = [[norm(A * normalize(s.u[j]))^2 for j=1:length(s.t)] for s=sol_tot]

@test all(η^2/κ^2 .≈ [k[end] for k in n])

#=
using Plots
gr()

p1 = plot(sol_tot[1].t, n[1], lw = 2)
for i=2:Ntraj
    plot!(p1,sol_tot[i].t, n[i])
end

p2 = plot(sol_tot[1].t, norm.(sol_tot[1].u).^2)
for i=2:Ntraj
    plot!(p2,sol_tot[i].t, norm.(sol_tot[i].u).^2)
end
=#

using OrdinaryDiffEq, DiffEqCallbacks, Test

# Initial state
u0=[0, -0.25, 0.42081, 0]

function hheom!(du,u,p,t)
    du[1] = u[3]
    du[2] = u[4]
    du[3] = -u[1] - 2u[1]*u[2]
    du[4] = -u[2] - (u[1]^2 - u[2]^2)
    return nothing
end

@inline Vhh(q1, q2) = 1//2 * (q1^2 + q2^2 + 2q1^2 * q2 - 2//3 * q2^3)
@inline Thh(p1, p2) = 1//2 * (p1^2 + p2^2)
@inline Hhh(q1, q2, p1, p2) = Thh(p1, p2) + Vhh(q1, q2)
@inline Hhh(u::AbstractVector) = Hhh(u...)

# Energy
const E = Hhh(u0)

function ghh(resid, u, p)
    resid[1] = Hhh(u[1],u[2],u[3],u[4]) - E
    resid[2:4] .= 0
end

# energy conserving callback:
# important to use save = false, I dont want rescaling points
cb = ManifoldProjection(ghh, nlopts=Dict(:ftol=>1e-13), save = false)

# Callback for Poincare surface of section
function psos_callback(j, direction = +1, offset::Real = 0,
    callback_kwargs = Dict{Symbol, Any}(:abstol=>1e-9))

    # Prepare callback:
    s = sign(direction)
    cond = (u,t,integrator) -> s*(u - offset)
    affect! = (integrator) -> nothing

    cb = DiffEqBase.ContinuousCallback(cond, nothing, affect!; callback_kwargs...,
    save_positions = (true,false), idxs = j)
end

# with this callback, the saved values of variable 1 should be zero
poincarecb = psos_callback(1)

totalcb = CallbackSet(poincarecb, cb)

prob = ODEProblem(hheom!, u0, (0.0, 100.0), callback = totalcb)

extra_kw = Dict(:save_start=>false, :save_end=>false)
DEFAULT_DIFFEQ_KWARGS = Dict{Symbol, Any}(:abstol => 1e-9, :reltol => 1e-9)

sol = solve(prob, Vern9(); extra_kw..., DEFAULT_DIFFEQ_KWARGS..., save_everystep = false)

Es = [Hhh(sol[:, i]) for i in 1:length(sol)]
Eerror = maximum(@. abs(E - Es))

a = sol[1, :]

@test Eerror < 1e-10
for el in a
    @test abs(el) < 1e-10
end


# # Эксперимент №1

# ## Набор параметров
α       = 0.004;        # Параметр регуляризации
w       = 0.0005;       # Эмпирический параметр регуляризации
S       = 35000;        # Количество итераций
β       = 0.003;        # Шаг минимизации
#
x_tp    = 0.25;         # Стартовое местоположение фронта
T_end   = 0.2;          # Регулируем конечное местоположение фронта
ε       = 10^(-1.5);    # Крутизна фронта
Nx      = 500;          # Число интервалов по ``X``
Mt      = 1000;         # Число интервалов по ``T``

# -----------------------------------------------------------------------------
using NonLinearReactionAdvectionDiffusionWithFrontData
using NonLinearReactionAdvectionDiffusionWithFrontData: heterogeneity_map;
using Serialization;
using Plots; pyplot();
# -----------------------------------------------------------------------------

# ## Решение на точных данных

# ### Решение прямой задачи для генерирования экспериментальной информации
# -----------------------------------------------------------------------------
a, b, t₀, T, N, M, ε, Xₙ, Tₘ, qₙ, ulₘ, urₘ, u₀ = dparams(x_tp = x_tp,
                                                         ε = ε,
                                                         Nx = Nx,
                                                         Mt = Mt,
                                                         T_end = T_end);
u, XX, TP = solve(u₀, Xₙ, N, Tₘ, M, ε, ulₘ, urₘ, qₙ);
ϕl, ϕr, ϕ, f1_data, f2_data = generate_obs_data(u, Xₙ, N, Tₘ, M, qₙ, ulₘ, urₘ);
directP = draft(u, Xₙ, N, Tₘ, M, title = "Эскиз прямого решения")
savefig(directP, "direct1.png");
# -----------------------------------------------------------------------------

# ### Найдем начальное приближение
# -----------------------------------------------------------------------------
q_guess = initial_guess(f1_data, Xₙ, N, Tₘ, M, ulₘ, urₘ, α);
plot(Xₙ, qₙ, label="Истинное");
plot!(Xₙ, q_guess, label="Найденное")
# -----------------------------------------------------------------------------

# -----------------------------------------------------------------------------
# В следующей строке, мы записываем все параметры в latex формате,
# не забываем экранировать все спецсимволы.
# Дальше, мы отобразим все эти параметры на графике.
using Printf;
params = "\$w = $(w), \\varepsilon = $(ε), "*
"\\beta = $(β), f_1 \\in $(@sprintf("[%.2f, %.2f]", extrema(f1_data)...)), "*
"\\alpha = $(α)\$"
nothing #hide
# -----------------------------------------------------------------------------

# ### Старт с найденного приближения на точных данных
# -----------------------------------------------------------------------------
q₀ = q_guess;
@time qs, Js, Qs = minimize(q₀, u₀, ulₘ, urₘ, Xₙ, N, Tₘ, M, ε, f1_data, f2_data,
                            S = S, β = β, w = w, showProgress = true)
serialize("qs.jld", qs);
serialize("Js.jld", Js);
serialize("Qs.jld", Qs);

# ### Эскиз процесса минимизации
# -----------------------------------------------------------------------------
a, b, c = minimization_draft(qₙ, Qs, Xₙ, N, Js, zoom = true,
                              annotate_string = params, zoom_chunk = 17//18)
plot(a, size = (800, 800))
plot(b)
plot(c)
withguessP = plot(a, b, c, layout = (3,1), size = (800, 2400));
savefig(withguessP, "withguess1.png")
nothing; #hide
# -----------------------------------------------------------------------------

# ## Слабо зашумленные данные
# -----------------------------------------------------------------------------
using Random
using NonLinearReactionAdvectionDiffusionWithFrontData: front_velocity;
rng = MersenneTwister(1234)

# Создадим мелкий шум.
f1_data_noised = f1_data + randn(rng, length(f1_data)) .* 0.00008
f2_data_noised = f2_data + randn(rng, length(f2_data)) .* 0.00008
plot(Tₘ, f1_data_noised, label="Зашумленные \$f_1\$")
plot!(Tₘ, f2_data_noised, label="Зашумленные \$f_2\$")

# Убедимся, что численное дифференцирование зашумленных функций --- это плохо.
v_f1 = front_velocity(f1_data_noised, Tₘ, M);
plot(Tₘ, v_f1[5:end-5], label="Скорость зашумленного \$f_1\$")

# Регуляризация во время поиска начального приближения справляется
q_guess = initial_guess(f1_data_noised, Xₙ, N, Tₘ, M, ulₘ, urₘ, α*3);
plot(Xₙ, qₙ, label="Истинное");
plot!(Xₙ, q_guess, label="Найденное")

# ### Старт с найденного приближения на зашумленных данных
# -----------------------------------------------------------------------------
q₀ = q_guess;
@time qs, Js, Qs = minimize(q₀, u₀, ulₘ, urₘ, Xₙ, N, Tₘ, M, ε, f1_data, f2_data,
                            S = S, β = β, w = w, showProgress = true)
serialize("qs_noised.jld", qs);
serialize("Js_noised.jld", Js);
serialize("Qs_noised.jld", Qs);

# ### Эскиз процесса минимизации
a, b, c = minimization_draft(qₙ, Qs, Xₙ, N, Js, zoom = true,
                              annotate_string = params, zoom_chunk = 17//18)
plot(a)
plot(b)
plot(c)
noisedP = plot(a, b, c);
savefig(noisedP, "noised.png");
nothing; #hide
# -----------------------------------------------------------------------------

# ### Старт с найденного приближения на точных данных
# -----------------------------------------------------------------------------
q₀ = q_guess;
@time qs, Js, Qs = minimize(q₀, u₀, ulₘ, urₘ, Xₙ, N, Tₘ, M, ε, f1_data, f2_data,
                            S = S, β = β, w = w, showProgress = true)
serialize("qs_noised.jld", qs);
serialize("Js_noised.jld", Js);
serialize("Qs_noised.jld", Qs);

# ### Эскиз процесса минимизации
a, b, c = minimization_draft(qₙ, Qs, Xₙ, N, Js, zoom = true,
                              annotate_string = params, zoom_chunk = 17//18)
plot(a)
plot(b)
plot(c)
noisedP = plot(a, b, c);
savefig(noisedP, "noised.png");
nothing; #hide
# -----------------------------------------------------------------------------


# ======== Extension of diff rules to complex functions ==========#

# preempt other function diffrules. The two entries correspond to d/dz and d/dz*
DiffRules.@define_diffrule Base.abs2(z) = :(conj($z)), :($z)
DiffRules.@define_diffrule Base.sqrt(z) = :(inv(2 * sqrt($z))), :(0)
DiffRules.@define_diffrule Base.log(z) = :(inv($z)), :(0)
DiffRules.@define_diffrule Base.:^(z,p) = :($p * ($z^($p - 1))), :(0)

DiffRules.@define_diffrule Base.abs(z) = :(0.5*conj($z)*inv(abs($z))), :(0.5*$z*inv(abs($z)))

# ======== Train functions to deal with complex duals ==========#

# Extend functions of a single complex argument to accept dual argument
macro extend_unary_dual_to_complex(fcn)
    M = :Base
    f = :($fcn)
    dfz, dfzstar = DiffRules.diffrule(M,f,:v)
    fdef = quote
        function $M.$f(z::ComplexDual{T}) where {T}
            v = value(z)
            dvz, dvzstar = dz_partials(z)

            # chain rule
            dfdz =     $dfz*dvz     + $dfzstar*conj(dvzstar)
            dfdzstar = $dfz*dvzstar + $dfzstar*conj(dvz)

            return ComplexDual{T}($M.$f(v),dfdz,dfdzstar)
        end
    end
    return esc(fdef)
end

# Extend functions of a single complex argument and another (non-dual) argument to accept
# dual argument in place of first argument
macro extend_binary_dual_to_complex(fcn)
    M = :Base
    f = :($fcn)
    dfz, dfzstar = DiffRules.diffrule(M,f,:v,:p)
    defs = quote end
    for R in AMBIGUOUS_TYPES
        expr = quote
            function $M.$f(z::ComplexDual{T},p::$R) where {T}
                v = value(z)
                dvz, dvzstar = dz_partials(z)

                # chain rule
                dfdz =     $dfz*dvz     + $dfzstar*conj(dvzstar)
                dfdzstar = $dfz*dvzstar + $dfzstar*conj(dvz)

                return ComplexDual{T}($M.$f(v,p),dfdz,dfdzstar)

            end
        end
        append!(defs.args, expr.args)
    end
    return esc(defs)
end

@extend_unary_dual_to_complex abs2
@extend_unary_dual_to_complex sqrt
@extend_unary_dual_to_complex log
@extend_unary_dual_to_complex abs

@extend_binary_dual_to_complex ^


# f(x)=A'*X
function FAtransposeX(A,X)
    if size(A)==(1,1)
        return (A[1].*X,nothing)
    elseif size(X)==(1,1)
        return (A'.*X[1],nothing)
    else
        return (A'*X,nothing)
    end
end

function FAtransposeX_inplace(value,auxvalue,A,X)
    if size(A)==(1,1)
        copy!(value,A[1]*X)
    elseif size(X)==(1,1)
        copy!(value,A'*X[1])
    else
        BLAS.gemm!('T','N',1.0,A,X,0.0,value)
    end
end

function DAtransposeX(derivativeIDX,f_c,faux_c,grad_c,grad_n,A,X)
    if derivativeIDX==1
        if size(A)==(1,1)
            axpy!(1.0,[sum(X.*grad_c)],grad_n)
        elseif size(X)==(1,1)
            axpy!(X[1],grad_c',grad_n)
        else
            BLAS.gemm!('N','T',1.0,X,grad_c,1.0,grad_n) # X*grad_c'
        end
    elseif derivativeIDX==2
        if size(A)==(1,1)
            axpy!(A[1],grad_c,grad_n)
        elseif size(X)==(1,1)
            axpy!(1.0,[sum(A'.*grad_c)],grad_n)
        else
            BLAS.gemm!('N','N',1.0,A,grad_c,1.0,grad_n) # A*grad_c
        end
    end
end

if PROC=="GPU"
#    function FAtransposeX(A::CudaArray,X::CudaArray)
#        if size(A)==(1,1)
#            value=CudaArray(Float64,size(X));
#            copy!(value,X); scale!(A,value)
#            return (value,[])
#        elseif size(X)==(1,1)
#            value=CudaArray(Float64,size(A));
#            copy!(value,A); scale!(X,value) # nb argument converse of Base.scale!
#            return (value,[])
#        else
#            return (CUBLAS.gemm('T','N',A,X),[])
#        end
#    end

    function FAtransposeX_inplace(value::CudaArray,auxvalue,A::CudaArray,X::CudaArray)
        if size(A)==(1,1)
            copy!(value,X); scale!(A,value) # nb argument converse of Base.scale!
        elseif size(X)==(1,1)
            CUBLAS.geam!('T','N',1.0,A,0.0,value,value) # nb argument converse of Base.scale!
            scale!(X,value) # nb argument converse of Base.scale!
        else
            CUBLAS.gemm!('T','N',1.0,A,X,0.0,value)
        end
    end

    function DAtransposeX(derivativeIDX,f_c,faux_c,grad_c,grad_n,A::CudaArray,X::CudaArray)
        if derivativeIDX==1
            if size(A)==(1,1)
                tmp=CudaArray(Float64,size(X))
                vmult!(1.0,X,grad_c,tmp)
                sum_update!(1.0,tmp,1.0,grad_n); free(tmp)
            elseif size(X)==(1,1)
                alphaaxpy!(1.0,X,grad_c,grad_n)
            else
                CUBLAS.gemm!('N','T',1.0,X,grad_c,1.0,grad_n) # X*grad_c'
            end
        elseif derivativeIDX==2
            if size(A)==(1,1)
                alphaaxpy!(1.0,A,grad_c,grad_n)
            elseif size(X)==(1,1)
                tmp=CudaArray(Float64,size(A))
                vmult!(1.0,A,grad_c,tmp)
                sum_update!(1.0,tmp,1.0,grad_n);free(tmp)
            else
                CUBLAS.gemm!('N','N',1.0,A,grad_c,1.0,grad_n) # A*grad_c
            end
        end
    end
end

Derivative[FAtransposeX]=DAtransposeX
export FAtransposeX

Inplace[FAtransposeX]=FAtransposeX_inplace

#import Base.*
#*(A::ADnode,B::ADnode)=ADnode(FAX,[A B])

#*(A::Real,B::ADnode)=ADnode(FAX,[ADconst(A) B])
#*(A::ADnode,B::Real)=ADnode(FAX,[A ADconst(B)])

#@gpu *(A::CudaArray,B::CudaArray)=CUBLAS.gemm('N','N',A,B)
#export *


abstract type AbstractBernstein{T,X} <: AbstractSpecialPolynomial{T,X} end

"""

    Bernstein{N, T, X}

A [Bernstein  polynomial](https://en.wikipedia.org/wiki/Bernstein_polynomial) is a polynomial expressed in terms of
Bernstein basic polynomials. For each degree, `𝐍`, this is a set of `𝐍+1` degree `𝐍` polynomials of the form:
`β_{𝐍,ν} =  (ν choose 𝐍) x^ν  (1-x)^{𝐍-ν}`, `0 ≤ x ≤ 1.`

The `Bernstein{𝐍,T}` type represents a polynomial of degree `𝐍` or
less with a linear combination of the basis vectors using coefficients
of type `T`.

```jldoctest Bernstein
julia> using Polynomials, SpecialPolynomials

julia> p = basis(Bernstein{3},  2)
Bernstein(1⋅β₂,₂(x))

julia> convert(Polynomial, p)
Polynomials.Polynomial(1.0*x^2)
```

!!! note
    [StaticUnivariatePolynomials](https://github.com/tkoolen/StaticUnivariatePolynomials.jl) Offers a  more  performant version.

"""
struct Bernstein{𝐍,T,X} <: AbstractBernstein{T,X}
    coeffs::Vector{T}
    function Bernstein{𝐍,T,X}(coeffs::AbstractVector{T}) where {𝐍,T,X}
        N = findlast(!iszero, coeffs)
        N == nothing && return new{𝐍,T,X}(zeros(T, 0))
        (N > 𝐍 + 1) && throw(ArgumentError("Wrong length for coefficents"))
        return new{𝐍,T,X}(coeffs[1:N])
    end
end

export Bernstein
Polynomials.@registerN Bernstein 𝐍
Polynomials.:⟒(::Type{<:Bernstein{𝐍}}) where {𝐍} = Bernstein{𝐍}

function Polynomials.showterm(
    io::IO,
    ::Type{Bernstein{N,T,X}},
    pj::T,
    var,
    j,
    first::Bool,
    mimetype,
) where {N,T,X}
    iszero(pj) && return false
    !first && print(io, " ")
    print(io, Polynomials.hasneg(T) && Polynomials.isneg(pj) ? "- " : (!first ? "+ " : ""))
    print(io, "$(abs.(pj))⋅β")
    Polynomials.unicode_subscript(io, N)
    print(io, ",")
    Polynomials.unicode_subscript(io, j)
    print(io, "($var)")
    return true
end

# add default case with no N specified
function Bernstein(coeffs::AbstractVector{T}, var::Symbol=:x) where {T}
    l = findlast(!iszero, coeffs)
    𝐍 = l == nothing ? -1 : l - 1
    Bernstein{𝐍,T,Symbol(var)}(coeffs[1:(𝐍 + 1)])
end

# promotion between family
function Base.promote(
    p::P,
    q::Q,
) where {𝐍,T,X,P<:Bernstein{𝐍,T,X},𝐌,S,Y,Q<:Bernstein{𝐌,S,Y}}
    X == Y || throw(ArgumentError("different indeterminates"))
    MN = max(𝐌, 𝐍)
    R = promote_type(T, S)
    PQ = Bernstein{MN,R}
    convert(PQ, p), convert(PQ, q)
end

# promote up
function Base.convert(P::Type{<:Bernstein{𝐍}}, p::Bernstein{𝐌,S}) where {𝐍,𝐌,S}
    @assert 𝐍 >= 𝐌
    𝐑 = 𝐍 - 𝐌
    R = eltype(one(promote_type(eltype(P), S)) / 1)
    cs = zeros(R, 𝐍 + 1)
    for j in 0:𝐌
        pⱼ = p[j]
        iszero(pⱼ) && continue
        for i in j:(j + 𝐑)
            cs[1 + i] += pⱼ * binomial(𝐌, j) * binomial(𝐑, i - j) / binomial(𝐍, i)
        end
    end
    X = Polynomials.indeterminate(P, p)
    Bernstein{𝐍,R,X}(cs)
end

function Base.convert(
    P::Type{<:Polynomials.StandardBasisPolynomial},
    p::Bernstein{N,T},
) where {N,T}
    x = variable(P)
    p(x)
end

function Base.convert(C::Type{<:Bernstein{𝐍,T}}, p::Polynomial) where {𝐍,T}
    @assert degree(p) <= 𝐍

    R = eltype(one(T) / one(Int))
    cs = zeros(R, 𝐍 + 1)

    for (i, a_nu) in enumerate(coeffs(p))
        k = i - 1
        nk = binomial(𝐍, k)
        for j in k:𝐍
            cs[j + 1] += a_nu / nk * binomial(j, k)
        end
    end

    X = Polynomials.indeterminate(C, p)
    Bernstein{𝐍,R,X}(cs)
end

function Base.convert(P::Type{<:AbstractBernstein}, p::Polynomial{T}) where {T}
    𝐍 = degree(p)
    X = Polynomials.indeterminate(P, p)
    convert(Bernstein{𝐍,T,X}, p)
end

Polynomials.domain(::Type{<:AbstractBernstein}) = Polynomials.Interval(0, 1)
Polynomials.degree(p::Bernstein{N,T,X}) where {N,T,X} = degree(convert(Polynomial{T,X}, p))

Polynomials.constantterm(p::Bernstein) = p[0] # β.,ᵢ = 0 + ... when i > 0

## need to do one  and variable differently,as coefficients have length N+1
Base.one(P::Type{Bernstein{N,T,X}}) where {N,T,X} = Bernstein{N,T,X}(ones(T, N + 1))
Base.one(P::Type{Bernstein{N,T}}, var::Polynomials.SymbolLike) where {N,T} =
    Bernstein{N,T,Symbol(var)}(ones(T, N + 1))
function Base.one(P::Type{<:Bernstein{N}}, var::Polynomials.SymbolLike=:x) where {N}
    N′ = max(N, 0)
    one(Bernstein{N′,eltype(P),Symbol(var)})
end

function Polynomials.variable(P::Type{Bernstein{𝐍,T,X}}) where {𝐍,T,X}
    𝐍 <= 0 && throw(ArgumentError("Need  𝐍 ≥ 1"))
    R = typeof(one(T) / 1)
    Bernstein{𝐍,R,X}([1 - (i * one(T)) / 𝐍 for i in 𝐍:-1:0])
end

function Polynomials.variable(
    P::Type{<:Bernstein{𝐍}},
    var::Polynomials.SymbolLike=:x,
) where {𝐍}
    S = eltype(P)
    variable(Bernstein{𝐍,S,Symbol(var)})
end
Polynomials.variable(P::Type{Bernstein}, var::Polynomials.SymbolLike=:x) =
    variable(Bernstein{1,eltype(P),Symbol(var)})

function basis(::Type{P}, k::Int) where {𝐍,T,X,P<:Bernstein{𝐍,T,X}}
    zs = zeros(Int, k + 1)
    zs[end] = 1
    P(zs)
end

function basis(P::Type{<:Bernstein}, k::Int, _var::Polynomials.SymbolLike=:x; var=_var)
    zs = zeros(Int, k + 1)
    zs[end] = 1
    Bernstein(zs, var)
end

# poly evaluation
Polynomials.evalpoly(x, p::Bernstein) = simple_deCasteljau_polyval(p, x)

# we could  use the compensated one
# from https://doi.org/10.1016/j.camwa.2010.05.021
# Accurate evaluation of a polynomial and its derivative in Bernstein form
# Hao Jiang, Shengguo Li, Lizhi Cheng, and Fang Su
#
# for increased accuracy, but is it really necessary here?
function simple_deCasteljau_polyval(p::Bernstein{𝐍,T,X}, t::S) where {𝐍,T,X,S}
    p₀ = p[0]
    R = eltype(p₀ * t)
    zR = p₀ * zero(t)

    R = promote_type(T, S)
    𝐍 == -1 && return zR
    𝐍 == 0 && return p[0] * one(t)
    n = length(coeffs(p))
    iszero(n) && return zR

    bs = [zR for _ in 1:(𝐍 + 1)]
    for i in eachindex(coeffs(p))
        bs[i] = p[i - 1]
    end

    r = 1 - t
    for j in 1:𝐍
        for i in 0:(𝐍 - j)
            ii = i + 1
            bs[ii] = bs[ii] * r + bs[ii + 1] * t
        end
    end
    bs[1]
end

# function deCasteljau_polyval(p::Bernstein{N,T}, t) where {N, T}
#     b̂ = coeffs(p) .* one(t)
#     errb̂ = 0 .* b̂ #zeros(eltype(b̂), N+1)
#     r̂,ρ = twoSum(1, -t)
#     for j in 1:N
#         for ii in 0:(N-j)
#             i = ii+1
#             s, π̂ = twoProd(b̂[i], r̂)
#             v, σ̂ = twoProd(b̂[i+1], t)
#             b̂[i],β̂ = twoSum(s, v)
#             ŵ = π̂ .+ σ̂ .+ β̂ .+ (b̂[i] .* ρ)
#             errb̂[i] = errb̂[i] * r̂ + (errb̂[i+1] * t + ŵ)
#         end
#     end
#     b̂[1]+ errb̂[1]

# end

# # compensated form to  find p'(t)
# function deCasteljau_derivative_val(p::Bernstein{N,T}, t) where {N,T}
#     R = eltype(one(T)*one(eltype(t)))
#     r̂,ρ  = twoSum(one(R), -t)
#     Δb = zeros(R, N).*one(t)
#     errΔb = similar(Δb) #zeros(R, N).*one(t)
#     for ii=0:N-1
#         i = ii+1
#         Δb[i], errΔb[i] = twoSum(p[ii+1], -p[ii])
#     end
#     for j in 1:N-1
#         for ii in 0:N-1-j
#             i = ii+1
#             s, π̂ = twoProd(Δb[i], r̂)
#             v, σ̂ = twoProd(Δb[i+1], t)
#             Δb[i], β̂ = twoSum(s, v)
#             ŵ = π̂ + σ̂ + β̂ + (Δb[i] * ρ)
#             errΔb[i] = errΔb[i] * r̂ + (errΔb[i+1] * t + ŵ)
#         end
#     end
#     (Δb[1] + errΔb[1]) * N
# end

# # Moller-Knuth from Exact computations with approximate arithmetic
# # http://perso.ens-lyon.fr/jean-michel.muller/
# function twoSum(a::Number, b::Number)
#     s = a + b
#     â = s - b
#     b̂ = s - â
#     δ_a = a - â
#     δ_b = b - b̂
#     s, δ_a + δ_b
# end
# twoSum(x, y) = (x + y, zero(x)+zero(y))

# function twoProd(a::Number, b::Number)
#     ab = a * b
#     ab, fma(a,b,-ab)
# end
# twoProd(x, y)  = x*y, 0

# not needed, but speeds things along
function Base.:+(p1::P, p2::Q) where {𝐍,T,X,P<:Bernstein{𝐍,T,X},𝐌,S,Q<:Bernstein{𝐌,S,X}}
    p, q = promote(p1, p2)
    𝐎, R = length(p.coeffs), eltype(p)
    P′ = Bernstein{𝐎,R,X}
    return Bernstein([p[i] + q[i] for i in 0:𝐎], X)
end

# no promote(p1,p2)  called here
function Base.:*(p::P, q::Q) where {𝐍,T,X,P<:Bernstein{𝐍,T,X},𝐌,S,Y,Q<:Bernstein{𝐌,S,Y}}
    ## use b(n,k) * b(m,j) = choose(n,k)choose(m,j)/choose(n+m,k+j) b(n+m, k+j)

    isconstant(p) && return q * constantterm(p)
    isconstant(q) && return p * constantterm(q)
    Polynomials.assert_same_variable(p, q) ||
        throw(ArgumentError("p1 and p2 must have the same symbol"))

    R = typeof(one(promote_type(T, S)) / 1)
    x = variable(Polynomial{R})
    return convert(Bernstein, p(x) * q(x))

    ## multiplication  through conversion  is as  efficient
    # n,  m =  degree(p), degree(q)
    # 𝐍 = n + m
    # cs = zeros(R, 𝐍 + 1)

    # for i in 0:𝐍
    #     for j in max(0,i-m):min(n, i)
    #         cs[1+i] += p[j] * q[i-j] * (one(R) * binomial(n,j)  *  binomial(m,i-j))
    #     end
    #     cs[1+i] /= binomial(m+n, i)
    # end

    # Bernstein{𝐍,R}(cs, p.var)

end

#  use p'(t) = n ∑ (b_{i+1} - b_i) B_i,n-1
function Polynomials.derivative(p::Bernstein{𝐍,T,X}, order::Integer=1) where {𝐍,T,X}
    cs = coeffs(p)
    𝐍 < -1 && return p
    iszero(𝐍) && return zero(p)
    order == 0 && return p

    cs = zeros(T, 𝐍)
    dp = Bernstein{𝐍 - 1,T,X}(𝐍 * diff(coeffs(p)))

    order > 1 ? derivative(dp, order - 1) : dp
end

# use ∫b_n,nu  = 1/(n+1)  * ∑_{nu+1}^{n+1} b_n+1, j
function Polynomials.integrate(p::Bernstein{𝐍,T,X}) where {𝐍,T,X}
    R = eltype(one(T) / 1)
    𝐍𝐍 = 𝐍 + 1

    cs = zeros(R, 𝐍𝐍 + 1)

    @inbounds for ν in eachindex(p)
        for j in (ν + 1):𝐍𝐍
            cs[1 + j] += p[ν] / 𝐍𝐍
        end
    end

    ∫p = Bernstein{𝐍𝐍,R,X}(cs)
    return ∫p
end

#  could do p 33 of http://mae.engr.ucdavis.edu/~farouki/bernstein.pdf
function Base.divrem(num::Bernstein{𝐍,T}, den::Bernstein{𝐌,S}) where {𝐍,T,𝐌,S}
    R = eltype((one(T) + one(S)) / 1)
    p1 = convert(Polynomial{R}, num)
    p2 = convert(Polynomial{R}, den)
    q, r = divrem(p1, p2)
    convert.(Bernstein, (q, r))
end

# Replace me. Doesn't seem necessary and poor stability of conversion.
# cf https://arxiv.org/pdf/1910.01998.pdf for one possible way, though for approximate GCD
function Base.gcd(num::Bernstein{𝐍,T}, den::Bernstein{𝐌,S}) where {𝐍,T,𝐌,S}
    ps = convert.(Polynomial, (num, den))
    qs = gcd(ps...)
    convert.(Bernstein, qs)
end

function Polynomials.vander(
    P::Type{<:AbstractBernstein},
    xs::AbstractVector{T},
    n::Integer,
) where {T<:Number}
    N = length(xs) - 1 # xs = [x0, x1, ...]
    R = typeof(one(T) / one(T))
    V = zeros(R, N + 1, n + 1)
    for j in 0:n
        bnj = binomial(n, j)
        for i in 0:N
            x = xs[i + 1]
            V[i + 1, j + 1] = bnj * x^j * (1 - x)^(n - j)     #beta(n,j)(xj)
        end
    end
    V
end

# direct compuation of roots using numerical stable method of y Guðbjörn and Jónsson as
# noted by Corless, Sevyeri  in §2.1 of
# https://arxiv.org/pdf/1910.01998.pdf
function Polynomials.roots(p::Bernstein{𝐍,T}) where {𝐍,T}
    R = eltype(one(T) / 1)

    as = zeros(R, 𝐍 + 1)
    for (i, pᵢ) in pairs(coeffs(p))
        as[i] = pᵢ
    end

    Ap = diagm(-1 => ones(R, 𝐍 - 1))
    for j in 1:𝐍
        Ap[1, j] = -as[end - j]
    end

    Bp = diagm(-1 => ones(R, 𝐍 - 1))
    for j in 1:𝐍
        Bp[1, j] = -as[end - j]
    end

    for j in 2:𝐍
        Bp[j, j] = j / (𝐍 - j + 1)
    end
    Bp[1, 1] += as[end] / 𝐍

    eigvals(Ap, Bp)
end

# http://www.cecm.sfu.ca/personal/pborwein/MITACS/papers/LinMatPol.pdf
# Sec 4.1

function comrade_matrix(p::P) where {𝐍,T,P<:Bernstein{𝐍,T}}
    C₀, C₁ = comrade_pencil(p)
    C₀ * inv(C₁)
end

function comrade_pencil(p::P) where {𝐍,T,P<:Bernstein{𝐍,T}}
    n = 𝐍
    a, b = 0, 1
    C₀ = diagm(
        0 => a / (b - a) .* ((n:-1:1) ./ (1:n)) .* ones(T, n),
        -1 => b / (b - a) .* ones(T, n - 1),
    )
    C₁ = diagm(
        0 => 1 / (b - a) .* ((n:-1:1) ./ (1:n)) .* ones(T, n),
        -1 => 1 / (b - a) .* ones(T, n - 1),
    )

    C₀[end, end] *= p[end]
    C₁[end, end] *= p[end]

    for i in 0:(n - 1)
        C₀[1 + i, end] -= b / (b - a) * p[i]
        C₁[1 + i, end] -= 1 / (b - a) * p[i]
    end
    #    C₁[1,end] = -(n-1)/(b-a) * p[0]

    C₀, C₁
end

# block version
function comrade_pencil(p::P) where {𝐍,T,M<:AbstractMatrix{T},P<:Bernstein{𝐍,M}}
    m, m′ = size(p[0])
    @assert m == m′

    Iₘ = I(m)
    n = 𝐍

    a, b = 0, 1 # could generalize with Bernsteing
    R = eltype(one(T) / 1)
    C₀ = zeros(R, n * m, n * m)
    C₁ = zeros(R, n * m, n * m)
    Δ = 1:m

    for i in 1:(n - 1)
        C₀[(i - 1) * m .+ Δ, (i - 1) * m .+ Δ] .= (n + 1 - i) / i * a / (b - a) .* Iₘ
        C₀[i * m .+ Δ, (i - 1) * m .+ Δ] .= b / (b - a) .* Iₘ

        C₁[(i - 1) * m .+ Δ, (i - 1) * m .+ Δ] .= (n + 1 - i) / i * 1 / (b - a) .* Iₘ
        C₁[i * m .+ Δ, (i - 1) * m .+ Δ] .= 1 / (b - a) .* Iₘ
    end

    C₀[(n - 1) * m .+ Δ, (n - 1) * m .+ Δ] .= 1 / n * a / (b - a) .* p[n]
    C₁[(n - 1) * m .+ Δ, (n - 1) * m .+ Δ] .= 1 / n * 1 / (b - a) .* p[n]

    for i in 0:(n - 1)
        C₀[i * m .+ Δ, (n - 1) * m .+ Δ] .-= b / (b - a) .* p[i]
        C₁[i * m .+ Δ, (n - 1) * m .+ Δ] .-= 1 / (b - a) .* p[i]
    end

    #    C₁[Δ, (n-1)*Δ] .= -(n-1) * 1 / (b-a) .* p[0]

    C₀, C₁
end


using Bessels
using Test
import SpecialFunctions

@time @testset "besseli" begin include("besseli_test.jl") end
@time @testset "besselk" begin include("besselk_test.jl") end
@time @testset "besselj" begin include("besselj_test.jl") end
@time @testset "bessely" begin include("bessely_test.jl") end


function diffeq(du,u,p,t)
    Rb    = u[V.uRb]  - u[V.tE2f]  + u[V.E2f]
    pRb   = p[C.tRb]  - u[V.uRb]
    RbE2f = u[V.tE2f] - u[V.E2f]

    uERK = p[C.tERK] - u[V.pERK] - u[V.pERK_DUSP]

    Ce  = u[V.tCe]  - u[V.CeP21]

    P21 = p[C.tP21]-u[V.CeP21]
    ###
    rSyE2f = p[C.kSyE2f] + p[C.kSyE2fE2f]*u[V.E2f]/(p[C.jSyE2f] + u[V.E2f])
    rPhRb  = p[C.kPhRbCd]*u[V.Cd] + p[C.kPhRbCe]*Ce
    rDeP21 = p[C.kDeP21] + p[C.kDeP21Cy]*p[C.Skp2]*Ce + p[C.pAkt]

    ###
    du[V.pERK] = (p[C.kf]*p[C.mitogen]*uERK - p[C.kr]*u[V.pERK]) - (p[C.kon]*p[C.DUSP]*u[V.pERK] - p[C.koff]*u[V.pERK_DUSP])
    du[V.pERK_DUSP] = p[C.kon]*p[C.DUSP]*u[V.pERK] - p[C.koff]*u[V.pERK_DUSP]
    du[V.pRSK] = p[C.kfRSK]*u[V.pERK]/(p[C.KmRSK]+u[V.pERK]) - p[C.krRSK]*u[V.pRSK]
    du[V.cFos] = p[C.kfFOS]*u[V.pERK]*u[V.pRSK]/(p[C.KmFOS]+u[V.pRSK]) - p[C.krFOS]*u[V.cFos]
    du[V.Cd] = p[C.kSyCd]*u[V.cFos] - p[C.kDeCd]*u[V.Cd]
    du[V.uRb]  = - rPhRb*u[V.uRb] + p[C.kDpRb]*pRb
    du[V.tE2f] = rSyE2f - p[C.kDeE2f]*u[V.tE2f]
    du[V.E2f]  = rSyE2f - p[C.kDeE2f]*u[V.E2f] - p[C.kAsRbE2f]*Rb*u[V.E2f] + (p[C.kDsRbE2f]+rPhRb)*RbE2f
    du[V.tCe]  = p[C.kSyCe] *u[V.E2f] - p[C.kDeCe]*u[V.tCe]
    du[V.CeP21] = p[C.kAsCyP21]*P21*Ce - (p[C.kDsCyP21]+p[C.kDeCe]+rDeP21)*u[V.CeP21]

end

function get_steady_state(p::Vector{Float64})
    tspan_ss::Tuple{Float64,Float64} = (0.0,Inf)
    u0::Vector{Float64} = zeros(V.NUM)

    prob = ODEProblem(diffeq,u0,tspan_ss,p)
    prob = SteadyStateProblem(prob)
    sol = solve(prob,DynamicSS(CVODE_BDF()),dt=1.0)

    return sol.u
end


function vor( u :: Array{Float64,2}, v :: Array{Float64,2}, m :: MITgcmDatas)
  dudy = dimDiff(u, 2, OnePointR()) /m.dyspacing
  dvdx = dimDiff(v, 1, OnePointR()) /m.dxspacing
  return (dvdx - dudy)
end

function vor( u :: Array{Float64,3}, v :: Array{Float64,3}, m :: MITgcmDatas)
  vv = zeros(size(u))
  for iz = 1: size(u,3)
    vv[:,:,iz] = vor(u[:,:,iz],v[:,:,iz],m)
  end
  return vv
end

#= include("intercept.jl") =#
function MITgcmPV( u :: Array{Float64,3}, v :: Array{Float64,3},
             T :: Array{Float64,3}, m :: MITgcmDatas;
             dzspacing :: Float64 = 2.5, f :: Float64 = 7.29e-5, BbyT :: Float64 = 2.0e-4 * 9.8)
  parIndent = OnePointR()
  dudy = begin
    fdy =  dimDiff(u, 2, OnePointLR()) /(m.dyspacing)
    dimAverage(fdy,1, parIndent)
    #= dimMoves(fdy,2,OnePointR()) =#
  end
  dvdx = begin
    fdx = dimDiff(v,1,OnePointLR())/(m.dxspacing)
    dimAverage(fdx, 2,parIndent)
    #= dimMoves(fdx,1,OnePointR()) =#
  end
  dudz = begin
    fdz = -dimDiff(u,3,OnePointLR())/dzspacing
    dimAverage(fdz,1, parIndent)
    #= dimMoves(fdz,2,OnePointR()) =#
  end
  dvdz = begin
    fdz = -dimDiff(v,3,OnePointLR())/dzspacing
    dimAverage(fdz,2, parIndent)
    #= dimMoves(fdz,1,OnePointR()) =#
  end
  avor = ( -dvdz, dudz, f+(dvdx - dudy))

  dB = begin
    TT = map( x-> BbyT * dimDiff(T,x, OnePointLR()), (1,2,3))
    (TT[1]/m.dxspacing, TT[2]/m.dyspacing, -TT[3]/dzspacing)
  end

  return mapreduce(+, zip(avor, dB))do x
    x[1].*x[2]
  end
end

export vor
export MITgcmPV, MITgcmVor

include("intercept.jl")
function MFPV( u :: Array{Float64,3}, v :: Array{Float64,3},
             T :: Array{Float64,3}, m :: MITgcmDatas;
             dzspacing :: Float64 = 2.5, f :: Float64 = 7.29e-5, BbyT :: Float64 = 2.0e-4 * 9.8,
             whether_smooth_T :: Bool = true
             )
  dudy = begin
    fdy =  dimDiff(u, 2, OnePointLR()) /(m.dyspacing)
    shift = GridMoving(uGrid, GridPosition(mGrid.x, mGrid.y - 0.5),2.0,2.0)
    intercept(fdy, shift)/(m.dyspacing)
  end
  dvdx = begin
    fdx = dimDiff(v,1,OnePointLR())/(m.dxspacing)
    shift = GridMoving(vGrid, GridPosition(mGrid.x - 0.5, mGrid.y),2.0,2.0)
    intercept(fdx, shift)
  end
  dudz = begin
    fdz = -dimDiff(u,3,OnePointLR())/dzspacing
    shift = GridMoving(uGrid, mGrid,2.,2.)
    intercept(fdz,shift)
  end
  dvdz = begin
    fdz = -dimDiff(v,3,OnePointLR())/dzspacing
    shift = GridMoving(vGrid, mGrid,2.,2.)
    intercept(fdz,shift)/dzspacing
  end
  avor = ( -dvdz, dudz, f+(dvdx - dudy))

  dB = begin
    T1 = intercept(T,GridMoving(0.0, 0.0, 2.0,2.0);  keep_when_notmoving=!whether_smooth_T)
    TT = map( x-> BbyT * dimDiff(T1,x, OnePointLR()), (1,2,3))
    (TT[1]/m.dxspacing, TT[2]/m.dyspacing, -TT[3]/dzspacing)
  end

  return mapreduce(+, zip(avor, dB))do x
    x[1].*x[2]
  end
end
export MFPV


using Plots

using simulator.materials.spesifics
using simulator.fresnelltools
using simulator.matrixcore

λ = 1200e-9

n1 = 1.5
nm = Pd_Werner(λ)
n2 = 1.2

θs = 0.7:1e-4:1

function point(θ, n2, which)
    i_1_2_te, i_1_2_tm, θm = FresnellBoundrary(n1, nm, θ)
    bulk = FresnellSlab(nm, 2*π/λ, 50e-9, θm)
    i_2_3_te, i_2_3_tm, _ = FresnellBoundrary(nm, n2, θm)

    system_te = CascadeScattering([i_1_2_te, bulk, i_2_3_te])
    system_tm = CascadeScattering([i_1_2_tm, bulk, i_2_3_tm])

    te_r = abs(system_te[2, 1])^2
    tm_r = abs(system_tm[2, 1])^2

    [te_r, tm_r][which]
end

anim = @animate for n ∈ 1:0.001:1.2
    plot(θs, point.(θs, n, 1), ylim = [0.5, 1], title="Plasmon reflectance in varying index environment", legend=:bottomright, label="TE")
    plot!(θs, point.(θs, n, 2), label="TM")
end

mp4(anim, "~/Downloads/temp.mp4", fps=30)


using FdeSolver
using Test
using SpecialFunctions
using Statistics

@testset "FdeSolver.jl" begin

    tSpan = [0, 1.5]
    y0 = 0
    β = 0.9

    function F(t, n, β, y)

        return (40320 ./ gamma(9 - β) .* t[n] .^ (8 - β) .- 3 .* gamma(5 + β / 2) ./ gamma(5 - β / 2) .* t[n] .^ (4 - β / 2) .+ 9/4 * gamma(β + 1) .+ (3/2 .* t[n] .^ (β / 2) .- t[n] .^ 4) .^ 3 .- y[n] .^ (3/2))

    end

    JacobF(t, n, β, y) = -(3 / 2) .* y[n] .^ (1 / 2)

    t, Yapp = FDEsolver(F, tSpan, y0, β, nothing, StopIt = "Convergence", tol = 10e-8, itmax = 15)
    t1, Yapp1 = FDEsolver(F, tSpan, y0, β, JacobF, StopIt = "Convergence", tol = 10e-8, itmax = 15)

    @test @isdefined(t)
    @test @isdefined(Yapp)

    @test @isdefined(t1)
    @test @isdefined(Yapp1)

    @test abs(mean(Yapp) - 0.99971) < 0.05
    @test abs(mean(Yapp1) - 0.99971) < 0.05

end


"""

A fast language can be beat
by a slow one `if` the code in the slow language solves the problem
more efficiently
 (if slower is uses a  smarter, better Algorithm )

takeaway: Implementation:
1. Language
2. Solution Strategy

You can also get the opposite

In a slow language , the converse can also be true

A poor algorithm might run faster than a good algorithm
(due to language limitations): #unlike Julia (doesn't have overhead, has optimization)
Tim Holy: the actual `impractice performance` you get from a language that doesn't do the Optimziations(julia does ) depends upon
what's Built-in for you
basically relies on people having written  `the things` that your particular Implementation requires
in a way that allows it to run quickly
invariably, you'll find certain concepts are efficient
other construcs arent' efficient
[some examples would change your impression on how the best way to solve a problem is]

Certain constructs are implemented efficiently
(Certain other constructs are not )

"""

"""Big-O Notation @22:34

you may notice there's a little work outside the Innermost loop
Point of Big-O
freedom to ignore terms that are less imprtant
me: I have to digress, it's a measure of complexity , as implied by demonstrating an inner-loop i.e. raising the already complex complexity
-Tim: "you might object", say  you're ignoring certain things  !
"""

for i = i:m
    s=0              # happens `m` times  - little work before time
    for j = 1:n
        s += M[i,j] * v[j] # happens `m*n` times

    end
    out[i] = s          # happens `m` times  # little work after time
end
#=more accurately, the running time is `a*m + b*m*n`, where a &b are constants # the actual running time
(since inner-loop runs m*n times )

Takeaway: Pay attention to the things that grow most quickly
so O(mn)
this algorithm is big-O -            a lisence to ignore small ones, yet focus on big ones that matter
=#

"""Key Rules about memory @41:38

1. Computers have multiple layers of caches
2. the fastest caches have only ~1/10^6 of storage of main memory
3. Main memory gets transferred to cache in small chuncks of adjacent memory (cache line)
4. Performace is Better If you Exploit what you have in cache before flushing it & swapping in new Memory
5. CPUs try to Predict what Chunck will be Needed next & start fetching it before

Tip: `LoopVectorization` will Reorder Your Loops & Make Other changes that help not Slow Code become "fast"
"""

"""Know the Costs & Benefits of Threading

A `Thread.@threads` has an overhead of a few microseconds.
equivalnet to 1000s of computations
(Reverse threding for sizable jobs)

- In cases where threads can work independently. it's sometimes possible to get an `~n` fold speedup from `n` threads

Communicating between threads can be a source of bugs &/or slowdowns

sum values of vector v
threaded loop (without it would look the same as normal loops) - enable us to Distribute the work
"""

function threadedsum(v)
    s = 0
    Threads.@threads for i in eachindex(v)
        s = s + v[i] #readout old value, then increments with next item on the list, store result back in s again
    end
    return s
end

#Define v
@time v = rand(1:5, 10^4) #0.000105 seconds (2allocations:78.203KiB)

@time threadedsum(v) # 0.092380 seconds (93.25 k allocations: 5.186 MiB, 99.55% compilation time)
# 30144
@time sum(v) # 0.017116 seconds (262 allocations: 13.641 KiB, 96.04% compilation time)

"""Race-condition

one of common things to happen (me:is it because not locking the common working memory? ) - possible
multiple working threads

3 working fast
1 of them gets interrupted
it gets that core, to do sth else
then it comes back,  picks off execution, where it left off

read out a v. old value of s
we don't know its' refreshed
then overwrite the value you had before
this causes it to drop many terms of v

sum (subset (value of v) )

Bugs arise from one thread being overwritted by another thread

1 way :
prevent 1 thread from accessing s

caled blocks new texes, atomic operation (happening in time )
in this case

problem blocking (& its generalizations)
while one
one thread holding lock, other (simply) has to wait - until the first one is done
(that delays the thread )
Unfavorable circumstances that take way way longer than a simple implentation
Solution:
by giving each threads to Maintin their counter

"""
function threadedsum2(v)
    spartial = zeros(Int, Threads.nthreads()) # s = 0 #allocates a n(counter) = to n(threads)
    Threads.@threads for i in eachindex(v)
        #s = s + v[i] #readout old value, then increments with next item on the list, store result back in s again
        spartial[Threads.threadid()] += v[i]  # each counter implements just its own counter
    end
    return sum(spartial) #return sum of 4 values Accumulated by each thread
end

@time threadedsum2(v) # 0.000087 seconds (9allocations:736bytes)

sum(v)

"""Other opportunities in parallelism

1. GPU Computing: julia can generate native code for GPUs (see JuliaGPU)
2. Distributed computing clusters
- each server is working on the portion of the same problem

cluseter computing is supported by julia support library (& other packages in julia ecosystem)

-Caution: in slow languages, sometimes people reach almost unthinkingly for parallelism
So if you're coming from a slow language, you may essentially reach parallelism almost inflexible
You get most out of it optimizing single mode
You'll iron out any performance gotchas (wasts lots of time ) -regardless if you parallize it
-Julia's advantage , you don't need parallelism to get the speed you need

Then, if you have done all that , really taken a serious look at your code
Good algorithm,
Implemented it well, using Idiomatic Julia ( & if it's still not passing up)
Then  that's the time to reach for parallelism
It's a well justified choice (in that circumstance )
But think it's worth emphasizing how important it is to

1.First, try to optimize simpler implementation -that's all you need- for many practical problems
2.When you Parallelize it, you'll get a Better Result - for having done so

good notes:

it's important to use julia well, because you can handicap your code (tell me about it)
and there's consequences of your performance quite severe if you do so

the goodside: there aren't many things you have to know in order to write good idiomatic julia code
you can then use julia pragmatic experience -testing differnet thing out-to quickly develop intuition about
what makes a good algorithm
it's a more reliable ally to your inituition
it really helps you grow as a programmer

Next steps :

1. Algorithms & data structures: Hash tables, Heaps, Priority Queues, et al.
2. extending knowledge of idiomatic Julia: read package code & Julia's `Base` Module
3. learning more about threading & High-performance computing

This course is a taste of lots of things that Prof. Tim Holy found useful
There's a lot out there - many fish in the sea
Tastes offered is just a basis
Exploration is a v. worthwhile experience
Active maintainer: continously spending time  as an
Odds are: v. good! you can get accustumed tuturing from more experienced julia development


module DAslabTests

using Test
using ForwardDiff
using LightPropagation: fluence_DA_semiinf_CW, fluence_DA_semiinf_TD
using LightPropagation: fluence_DA_slab_CW, fluence_DA_slab_TD
using LightPropagation: flux_DA_slab_CW, flux_DA_slab_TD, flux_DA_semiinf_TD, flux_DA_semiinf_CW

# Test slabs against semi-inf for sufficiently thick slabs

# CW
@test fluence_DA_slab_CW(1.0, 0.1, 10.0, n_ext = 1.0, n_med = 1.0, s = 20.0, z = 0.0) == 0.024035998288332954
@test fluence_DA_slab_CW(2.4, 0.6, 18.0, n_ext = 1.3, n_med = 1.2, s = 15.0, z = 1.0) == 1.77784095787581e-7
@test fluence_DA_semiinf_CW(1.0, 0.1, 10.0, n_ext = 1.0, n_med = 1.0, z = 0.0) ≈ fluence_DA_slab_CW(1.0, 0.1, 10.0, n_ext = 1.0, n_med = 1.0, z = 0.0, s = 20.0)
@test fluence_DA_semiinf_CW(4.0, 0.01, 50.0, n_ext = 1.0, n_med = 1.0, z = 0.0) ≈ fluence_DA_slab_CW(4.0, 0.01, 50.0,n_ext = 1.0, n_med = 1.0, z = 0.0, s = 20.0)

# TD
t = 0.1:0.5:8
@test fluence_DA_semiinf_TD(t, 1.0, 0.1, 10.0, n_ext = 1.0, n_med = 1.0, z = 0.0) ≈ fluence_DA_slab_TD(t, 1.0, 0.1, 10.0, n_ext = 1.0, n_med = 1.0, z = 0.0, s = 40.0)
@test fluence_DA_semiinf_TD(t, 1.0, 0.01, 70.0, n_ext = 1.0, n_med = 1.0, z = 0.0) ≈ fluence_DA_slab_TD(t, 1.0, 0.01, 70.0, n_ext = 1.0, n_med = 1.0, z = 0.0, s = 40.0)

@test fluence_DA_semiinf_TD(t, 3.0, 0.1, 10.0, n_ext = 1.0, n_med = 1.0, z = 0.0) ≈ fluence_DA_slab_TD(t, 3.0, 0.1, 10.0, n_ext = 1.0, n_med = 1.0, z = 0.0, s = 40.0)
@test fluence_DA_semiinf_TD(t, 3.0, 0.01, 70.0, n_ext = 1.0, n_med = 1.0, z = 0.0) ≈ fluence_DA_slab_TD(t, 3.0, 0.01, 70.0, n_ext = 1.0, n_med = 1.0, z = 0.0, s = 40.0)
@test fluence_DA_semiinf_TD(t, 3.0, 0.01, 70.0, n_ext = 1.4, n_med = 1.2, z = 0.1) ≈ fluence_DA_slab_TD(t, 3.0, 0.01, 70.0, n_ext = 1.4, n_med = 1.2, z = 0.1, s = 40.0)

## check flux

# reflectance

D = 1 / (3 * 10.0)
@test (D * ForwardDiff.derivative(dz -> fluence_DA_slab_TD(1.0, 1.0, 0.1, 10.0, n_ext = 1.0, n_med = 1.0, s = 1.0, z = dz), 0.0)) ≈ flux_DA_slab_TD(1.0, 1.0, 0.1, 10.0,  n_ext = 1.0, n_med = 1.0, s = 1.0, z = 0.0)

D = 1 / (3 * 20.0)
@test (D * ForwardDiff.derivative(dz -> fluence_DA_slab_TD(1.0, 1.3, 0.3, 20.0, n_ext = 1.0, n_med = 1.0, s = 1.0, z = dz), 0.0)) ≈ flux_DA_slab_TD(1.0, 1.3, 0.3, 20.0,  n_ext = 1.0, n_med = 1.0, s = 1.0, z = 0.0)

# transmittance
D = 1 / (3 * 20.0)
@test (-D * ForwardDiff.derivative(dz -> fluence_DA_slab_TD(1.0, 1.3, 0.3, 20.0, n_ext = 1.0, n_med = 1.0, s = 1.0, z = dz), 1.0)) ≈ flux_DA_slab_TD(1.0, 1.3, 0.3, 20.0,  n_ext = 1.0, n_med = 1.0, s = 1.0, z = 1.0)

## CW
@test flux_DA_slab_CW(1.0, 0.1, 10.0; n_ext = 1.0, n_med = 1.0, s = 40.0, z = 0.0, xs = 15) ≈ flux_DA_semiinf_CW(1.0, 0.1, 10.0; n_ext = 1.0, n_med = 1.0)

end


using FieldTracer, Meshes
using Test

@testset "FieldTracer.jl" begin
   @testset "Seeding" begin
      x = 0.0:0.1:1.0
      y = 1.0:0.1:2.0
      seeds = select_seeds(x, y; nsegment=(2,2))
      @test seeds[:,4] == [0.75, 1.75]
      z = 2.0:0.1:3.0
      seeds = select_seeds(x, y, z; nsegment=(2,2,2))
      @test seeds[:,2] == [0.75, 1.25, 2.25]
   end

   @testset "2D structured mesh" begin
      # bilinear interpolation function in a normalized rectangle
      x, y = 0.1, 0.2
      Q00, Q10, Q01, Q11 = 3.0, 40.0, 5.0, 60.0
      out = FieldTracer.bilin_reg(x, y, Q00, Q10, Q01, Q11)
      @test out ≈ 7.460 atol=1e-5

      # Euler2 and RK4 functions
      # https://se.mathworks.com/help/matlab/ref/streamline.html
      ds, maxstep = 0.1, 150
      x = range(0, 1, step=0.1)
      y = range(0, 1, step=0.1)
      # ndgrid
      xgrid = [i for i in x, _ in y]
      ygrid = [j for _ in x, j in y]
      u = copy(xgrid)
      v = -copy(ygrid)
      xstart = 0.1
      ystart = 0.9

      xt, yt = FieldTracer.euler(maxstep, ds, xstart, ystart, x, y, u, v)

      @test length(xt) == 141

      xt, yt = FieldTracer.rk4(maxstep, ds, xstart, ystart, x, y, u, v)

      @test length(xt) == 140

      grid = CartesianGrid((length(x)-1,length(y)-1),(0.,0.),(0.1,0.1))

      # default direction is both
      xt, yt = trace(u, v, xstart, ystart, grid; alg=Euler(), maxstep, ds)

      @test length(xt) == 148

      include("test_structured.jl")
      # asymptotic field
      @test test_trace_asymptote()     # single precision
      @test test_trace_asymptote(true) # double precision

      # dipole field
      @test test_trace_dipole()        # dipole tracing in 2D
   end

   @testset "3D structured mesh" begin
      # trilinear interpolation function in a normalized box
      x, y, z = 0.1, 0.2, 0.3
      Q = [3.0, 40.0, 5.0, 60.0, 3.0, 40.0, 5.0, 60.0]
      out = FieldTracer.trilin_reg(x, y, z, Q)
      @test out ≈ 7.460 atol=1e-5

      x = range(0, 10, length=15)
      y = range(0, 10, length=20)
      z = range(0, 10, length=25)
      bx = fill(1.0, length(x), length(y), length(z))
      by = fill(1.0, length(x), length(y), length(z))
      bz = fill(1.0, length(x), length(y), length(z))
      xs, ys, zs = 1.0, 1.0, 1.0

      # Euler 2nd order
      x1, y1, z1 = trace(bx, bz, bz, xs, ys, zs, x, y, z;
         alg=Euler(), ds=0.2, maxstep=200)

      @test length(x1) == 170 && x1[end] ≈ y1[end] ≈ z1[end]

      # RK4 by default
      x1, y1, z1 = trace(bx, bz, bz, xs, ys, zs, x, y, z;
         ds=0.2, maxstep=200, direction="forward")

      @test length(x1) == 152 && x1[end] ≈ y1[end] ≈ z1[end]

      Δx = x[2] - x[1]
      Δy = y[2] - y[1]
      Δz = z[2] - z[1]

      grid = CartesianGrid((length(x)-1, length(y)-1, length(z)-1),
         (0., 0., 0.),
         (Δx, Δy, Δz))

      # default direction is both
      x1, y1, z1 = trace(bx, bz, bz, xs, ys, zs, grid;
         alg=Euler(), ds=0.2, maxstep=200)

      @test length(x1) == 170 && x1[end] ≈ y1[end] ≈ z1[end]
   end

   @testset "2D unstructured mesh" begin
      include("test_unstructured.jl")
      @test test_trace_unstructured2D()
   end

end


"""
    struct OLS
        X::Array{Float64,2}
        y::Array{Float64,1}
        betas::Array{Float64,1}
    end

    Immutable data structure that holds design matrix, response vector, and estimated regression parameters.

# Arguments
- `X::Array{Float64, 2}`: Design matrix.
- `y::Array{Float64, 1}`: Response vector.
- `betas::Array{Float64, 1}`: Regression coefficients.

See also: [`ols`](@ref), [`wls`](@ref)
"""
struct OLS
    X::Array{Float64,2}
    y::Array{Float64,1}
    betas::Array{Float64,1}
end

"""
    ols(X, y)

    Create OLS object with estimated regression coefficients.

# Arguments
- `X::Array{Float64, 2}`: Design matrix.
- `y::Array{Float64, 1}`: Response vector.

# Examples
```julia-repl
julia> X = hcat(ones(24), phones[:,"year"]);
julia> y = phones[:,"calls"];
julia> reg = ols(X, y)
julia> reg.betas
2-element Array{Float64,1}:
 -260.0592463768119
    5.04147826086957
```

See also: [`OLS`](@ref), [`wls`](@ref)

"""
ols(X::Array{Float64,2}, y::Array{Float64,1})::OLS = OLS(X, y, qr(X, Val(true)) \ y)

"""
    wls(X, y, wts)

    Estimate weighted least squares regression and create OLS object with estimated parameters.

# Arguments
- `X::Array{Float64, 2}`: Design matrix.
- `y::Array{Float64, 1}`: Response vector.
- `wts::Array{Float64, 1}`: Weights vector.

# Examples
```julia-repl
julia> X = hcat(ones(24), phones[:,"year"]);
julia> y = phones[:,"calls"];
julia> w = ones(24)
julia> w[15:20] .= 0.0
julia> reg = wls(X, y, w)
julia> reg.betas
2-element Array{Float64,1}:
 -63.481644325290425
   1.3040571939231453
```

See also: [`ols`](@ref), [`OLS`](@ref)

"""
function wls(X::Array{Float64,2}, y::Array{Float64,1}, wts::Array{Float64,1})
    W = Diagonal(sqrt.(wts))
    #  I commented this because passing weighted values of X and y to OLS
    #  causes wrong calculations of residuals.
    #  return ols(W * X, W * y)
    return OLS(X, y, qr(W * X, Val(true)) \ (W * y))
end

"""
    residuals(ols)

Estimate weighted least squares regression and create OLS object with estimated parameters.

# Arguments
- `ols::OLS`: OLS object, possible created using `ols` or `wls`.

See also: [`ols`](@ref), [`wls`](@ref)
"""
residuals(ols::OLS) = ols.y .- ols.X * ols.betas

"""
    coef(ols)

Extract regression coefficients from an `OLS` object.

# Arguments
- `ols::OLS`: OLS object, possible created using `ols` or `wls`.

See also: [`ols`](@ref), [`residuals`](@ref), [`predict`](@ref)
"""
coef(ols::OLS) = ols.betas

"""
    predict(ols)

Calculate estimated response using an `OLS` object.

# Arguments
- `ols::OLS`: OLS object, possible created using `ols` or `wls`.

See also: [`ols`](@ref), [`residuals`](@ref)
"""
predict(ols::OLS) = ols.X * ols.betas

predict(ols::OLS, X::Array{Float64,2}) = X * ols.betas


"""

This example learns to classify images of fashion products(trousers, shirts, bags...)
from the [Fashion-MNIST](https://github.com/zalandoresearch/fashion-mnist) dataset.
There are 60000 training and 10000 test examples. Each input x
consists of 784 pixels representing a 28x28 image. The pixel values are
normalized to [0,1]. Each output y is converted to a ten-dimensional
one-hot vector (a vector that has a single non-zero component) indicating
the correct class (0-9) for a given image. 10 is used instead of 0.
Labels and descriptions are shown below.

    Label   Description
    1       T-shirt/top
    2       Trouser
    3       Pullover
    4       Dress
    5       Coat
    6       Sandal
    7       Shirt
    8       Sneaker
    9       Bag
    10      Ankle boot

You can run the demo using `julia fashion-mnist.jl` on the command line or
by first including `julia> include("fashion-mnist.jl")` and typing `julia> FashionMNIST.main()`
at the Julia prompt.  Options can be used like `julia fashion-mnist.jl --epochs 3`
or `julia> FashionMNIST.main("--epochs 3")`. Use `julia fashion-mnist.jl --help`
for a list of options.  The dataset will be automatically downloaded.
By default a softmax model will be trained for 10 epochs. You can also
train a multi-layer perceptron by specifying one or more --hidden sizes.
The accuracy for the training and test sets will be printed at every epoch
and optimized parameters will be returned.

"""
module FashionMNISTMLP
using Knet,CUDA,MLDatasets,ArgParse

function predict(w,x; pdrop=0)
    x = mat(x)
    for i=1:2:length(w)
        x = w[i]*dropout(x, pdrop) .+ w[i+1]
        if i<length(w)-1
            x = relu.(x) # max(0,x)
        end
    end
    return x
end

loss(w,x,ygold;pdrop=0) = nll(predict(w,x;pdrop=pdrop), ygold)

lossgradient = grad(loss)

function train(w, dtrn; lr=.5, epochs=10, pdrop=0)
    for epoch=1:epochs
        for (x,y) in dtrn
            g = lossgradient(w, x, y; pdrop=pdrop)
            update!(w, g, lr=lr)
        end
    end
    return w
end

function weights(h...; atype=Array{Float32}, winit=0.1)
    w = Any[]
    x = 28*28
    for y in [h..., 10]
        push!(w, convert(atype, winit*randn(y,x)))
        push!(w, convert(atype, zeros(y, 1)))
        x = y
    end
    return w
end

function main(args="")
    s = ArgParseSettings()
    s.description="fashion-mnist.jl (c) 2017 Adapted by Emre Unal based on Deniz Yuret’s MNIST example https://github.com/denizyuret/Knet.jl/tree/master/examples/mnist-mlp/mlp.jl.\nMulti-layer perceptron model on the Fashion-MNIST dataset from https://github.com/zalandoresearch/fashion-mnist.\n"
    s.exc_handler=ArgParse.debug_handler
    @add_arg_table s begin
        ("--seed"; arg_type=Int; default=-1; help="random number seed: use a nonnegative int for repeatable results")
        ("--batchsize"; arg_type=Int; default=100; help="minibatch size")
        ("--epochs"; arg_type=Int; default=10; help="number of epochs for training")
        ("--hidden"; nargs='*'; arg_type=Int; help="sizes of hidden layers, e.g. --hidden 128 64 for a net with two hidden layers")
        ("--lr"; arg_type=Float64; default=0.15; help="learning rate")
        ("--winit"; arg_type=Float64; default=0.1; help="w initialized with winit*randn()")
        ("--fast"; action=:store_true; help="skip loss printing for faster run")
        ("--atype"; default="$(Knet.array_type[])"; help="array type: Array for cpu, KnetArray for gpu")
        ("--gcheck"; arg_type=Int; default=0; help="check N random gradients per parameter")
        ("--dropout"; arg_type=Float64; default=0.5; help="Dropout probability.")
        # These are to experiment with sparse arrays
        # ("--xtype"; help="input array type: defaults to atype")
        # ("--ytype"; help="output array type: defaults to atype")
    end
    isa(args, AbstractString) && (args=split(args))
    if in("--help", args) || in("-h", args)
        ArgParse.show_help(s; exit_when_done=false)
        return
    end
    o = parse_args(args, s; as_symbols=true)
    if !o[:fast]
        println(s.description)
        println("opts=",[(k,v) for (k,v) in o]...)
    end
    o[:seed] > 0 && Knet.seed!(o[:seed])
    atype = eval(Meta.parse(o[:atype]))
    w = weights(o[:hidden]...; atype=atype, winit=o[:winit])
    xtrn,ytrn = FashionMNIST.traindata()
    xtst,ytst = FashionMNIST.testdata()
    global dtrn = minibatch(xtrn, ytrn, o[:batchsize]; xtype=atype, xsize=(size(xtrn,1),size(xtrn,2),1,o[:batchsize]))
    global dtst = minibatch(xtst, ytst, o[:batchsize]; xtype=atype, xsize=(size(xtrn,1),size(xtrn,2),1,o[:batchsize]))
    report(epoch)=println((:epoch,epoch,:trn,accuracy(w,dtrn,predict),:tst,accuracy(w,dtst,predict)))
    if o[:fast]
        @time (train(w, dtrn; lr=o[:lr], epochs=o[:epochs], pdrop=o[:dropout]); CUDA.functional() && CUDA.synchronize())
    else
        report(0)
        @time for epoch=1:o[:epochs]
            train(w, dtrn; lr=o[:lr], epochs=1)
            report(epoch)
            if o[:gcheck] > 0
                gradcheck(loss, w, first(dtrn)...; gcheck=o[:gcheck], verbose=true)
            end
        end
    end
    return w
end

# This allows both non-interactive (shell command) and interactive calls like:
# $ julia mnist.jl --epochs 10
# julia> FashionMNIST.main("--epochs 10")
PROGRAM_FILE == "fashion-mnist.jl" && main(ARGS)

end # module


using Neet, Test

@testset "BooleanSpace" begin
    @testset "Construction" begin
        @test_throws DomainError BooleanSpace(-1)
        @test_throws DomainError BooleanSpace(0)

        let space = BooleanSpace(1)
            @test space.ndim == 1
            @test space.volume == 2
            @test space.isuniform == true
        end

        let space = BooleanSpace(2)
            @test space.ndim == 2
            @test space.volume == 4
            @test space.isuniform == true
        end
    end

    @testset "size" begin
        @test size(BooleanSpace(1)) == (2,)
        @test size(BooleanSpace(2)) == (2,2)
        @test size(BooleanSpace(3)) == (2,2,2)
    end

    @testset "volume" begin
        @test volume(BooleanSpace(1)) == 2
        @test volume(BooleanSpace(2)) == 4
        @test volume(BooleanSpace(3)) == 8
    end

    @testset "encode" begin
        let space = BooleanSpace(1)
            @test encode(space, [0]) == 0
            @test encode(space, [1]) == 1
        end
        let space = BooleanSpace(2)
            @test encode(space, [0,0]) == 0
            @test encode(space, [0,1]) == 1
            @test encode(space, [1,0]) == 2
            @test encode(space, [1,1]) == 3
        end
        let space = BooleanSpace(3)
            @test encode(space, [0,0,0]) == 0
            @test encode(space, [0,0,1]) == 1
            @test encode(space, [0,1,0]) == 2
            @test encode(space, [0,1,1]) == 3
            @test encode(space, [1,0,0]) == 4
            @test encode(space, [1,0,1]) == 5
            @test encode(space, [1,1,0]) == 6
            @test encode(space, [1,1,1]) == 7
        end
    end

    @testset "decode" for ndim in 1:3
        let space = BooleanSpace(ndim)
            @testset "$(ndim)D Space" for d in 0:volume(space)-1
                @test encode(space, decode(space, d)) == d
            end
        end
    end

    @testset "iterate" for ndim in 1:3
        let space = BooleanSpace(ndim)
            @testset "$(ndim)D Space" for (i, state) in enumerate(space)
                @test encode(space, state) == i-1
            end
        end
    end
end


## Legendre Polynomials = Gegenbauer{1/2}

# cf FastGaussquadrature.jl
#include("FastLegendre.jl")

@register0 Legendre AbstractCCOP0
export Legendre

"""
    Legendre{T}

Implements the [Legendre](https://en.wikipedia.org/wiki/Legendre_polynomials) polynomials. These have weight function `w(x) = 1` over the domain `[-1,1]`.

```jldoctest
julia> using Polynomials, SpecialPolynomials

julia> p = Legendre([1,2,3])
Legendre(1⋅P₀(x) + 2⋅P₁(x) + 3⋅P₂(x))

julia> convert(Polynomial, p)
Polynomials.Polynomial(-0.5 + 2.0*x + 4.5*x^2)

julia> p2m, p2m1 = basis.(Legendre, (8,9)) # evaluation P_{2m+k}(-1) =  (-1)^k
(Legendre(1.0⋅P₈(x)), Legendre(1.0⋅P₉(x)))

julia> p2m(-1) == 1
false

julia> p2m1(-1) == -1
false

julia> n = 5  # verify  Rodrigues' formula
5

julia> x = Polynomial(:x)
Polynomials.Polynomial(1.0*x)

julia> derivative((x^2-1)^n, n) - 2^n *  factorial(n) * basis(Legendre, n)
Polynomials.Polynomial(0.0)

julia> p4, p5  =  basis.(Legendre, (4,5)) # verify  orthogonality  of  P₄,P₅
(Legendre(1.0⋅P₄(x)), Legendre(1.0⋅P₅(x)))

julia> SpecialPolynomials.innerproduct(Legendre, p4,  p5)
-1.5309832087675112e-16
```
"""
Legendre

abcde(::Type{<:Legendre}) = NamedTuple{(:a, :b, :c, :d, :e)}((-1, 0, 1, -2, 0))

basis_symbol(::Type{<:Legendre}) = "P"
Polynomials.domain(::Type{<:Legendre}) = Polynomials.Interval(-1, 1)

k0(P::Type{<:Legendre}) = one(eltype(P))
k1k0(P::Type{<:Legendre}, n::Int) = (one(eltype(P)) * (2n + 1)) / (n + 1) #k1k0(Gegenbauer{1/2, eltype(P)}, n)

norm2(::Type{<:Legendre}, n) = 2 / (2n + 1)
ω₁₀(P::Type{<:Legendre}, n) = sqrt((one(eltype(P)) * (2n + 1)) / (2n + 3))
weight_function(::Type{<:Legendre}) = x -> one(x)
generating_function(::Type{<:Legendre}) = (t, x) -> 1 / sqrt(1 - 2x * t + t^2)
function classical_hypergeometric(::Type{<:Legendre}, n, x)
    as = (-n, n + 1)
    bs = (1,)
    pFq(as, bs, (1 - x) / 2)
end
#eval_basis(::Type{P}, n, x::Float64) where {P <: Legendre} = FastLegendre.fastlegendre(n,x)
# cf. fastgauss.jl
# eval_basis(::Type{P}, n, x::Union{Int, Int32, Float16, Float32}) where {P <: Legendre} = FastLegendre.fastlegendre(n,float(x))
#gauss_nodes_weights(p::Type{P}, n) where {P <: Legendre} =
#    FastGaussQuadrature.gausslegendre(n)
# overrides

B̃n(P::Type{<:Legendre}, n::Int) = zero(eltype(P))
C̃n(P::Type{<:Legendre}, n::Int) = (one(eltype(P)) * n^2) / (4n^2 - 1)

B̃n(P::Type{<:Legendre}, ::Val{0}) = B̃n(Gegenbauer{1 / 2,eltype(P)}, Val(0))
C̃n(P::Type{<:Legendre}, ::Val{0}) = zero(eltype(P))
b̂̃n(P::Type{<:Legendre}, ::Val{0}) = b̂̃n(Gegenbauer{1 / 2,eltype(P)}, Val(0))
ĉ̃n(P::Type{<:Legendre}, ::Val{0}) = ĉ̃n(Gegenbauer{1 / 2,eltype(P)}, Val(0))

function Polynomials.derivative(p::Legendre{T,X}, order::Integer=1) where {T,X}
    order < 0 && throw(ArgumentError("Order of derivative must be non-negative"))
    order == 0 && return convert(Legendre{T}, p)
    hasnan(p) && return Legendre{T,X}(T[NaN])
    order > length(p) && return zero(Legendre{T})

    d = degree(p)
    qs = zeros(T, d)
    for i in 0:(d - 1)
        gamma = 2 * i + 1
        qs[i + 1] = gamma * sum(p[j] for j in (i + 1):2:d)
    end

    q = Legendre{T,X}(qs)

    if order > 1
        derivative(q, order - 1)
    else
        q
    end
end

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

# Monic
@register0 MonicLegendre AbstractCCOP0
export MonicLegendre
ϟ(::Type{<:MonicLegendre}) = Legendre
ϟ(::Type{<:MonicLegendre{T}}) where {T} = Legendre{T}
@register_monic(MonicLegendre)

## fast  gauss nodes
pqr_symmetry(::Type{<:MonicLegendre}) = true
pqr_weight(P::Type{<:MonicLegendre}, n, x, dπx) = (one(eltype(P)) * 2) / (1 - x^2) / dπx^2

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

@register0 ShiftedLegendre AbstractCCOP0
"""
    ShiftedLegendre

Type for the shifted Legendre polynomials: `Pˢᵢ(x) =  Pᵢ(2x-1)` for `x ∈ [0,1]`.
"""
ShiftedLegendre

ϟ(::Type{<:ShiftedLegendre}) = Legendre
ϟ(::Type{<:ShiftedLegendre{T}}) where {T} = Legendre{T}
export ShiftedLegendre
@register_shifted(ShiftedLegendre, 2, -1)

B̃n(P::Type{<:ShiftedLegendre}, ::Val{0}) = -one(eltype(P)) / 2

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

@register0 MonicShiftedLegendre AbstractCCOP0
export MonicShiftedLegendre
ϟ(::Type{<:MonicShiftedLegendre}) = ShiftedLegendre
ϟ(::Type{<:MonicShiftedLegendre{T}}) where {T} = ShiftedLegendre{T}
@register_monic(MonicShiftedLegendre)

@register0 OrthonormalLegendre AbstractCCOP0
export OrthonormalLegendre
ϟ(::Type{<:OrthonormalLegendre}) = Legendre
ϟ(::Type{<:OrthonormalLegendre{T}}) where {T} = Legendre{T}
@register_orthonormal(OrthonormalLegendre)


struct BabyPOMG
    babyPOMDP::CryingBaby
end

function MultiCaregiverCryingBaby()
    BabyPOMDP = CryingBaby()
    return BabyPOMG(BabyPOMDP)
end

n_agents(pomg::BabyPOMG) = 2

ordered_states(pomg::BabyPOMG) = [SATED, HUNGRY]
ordered_actions(pomg::BabyPOMG, i::Int) = [FEED, IGNORE, SING]
ordered_joint_actions(pomg::BabyPOMG) = vec(collect(Iterators.product([ordered_actions(pomg, i) for i in 1:n_agents(pomg)]...)))

n_actions(pomg::BabyPOMG, i::Int) = length(ordered_actions(pomg, i))
n_joint_actions(pomg::BabyPOMG) = length(ordered_joint_actions(pomg))

ordered_observations(pomg::BabyPOMG, i::Int) = [CRYING, QUIET]
ordered_joint_observations(pomg::BabyPOMG) = vec(collect(Iterators.product([ordered_observations(pomg, i) for i in 1:n_agents(pomg)]...)))

n_observations(pomg::BabyPOMG, i::Int) = length(ordered_observations(pomg, i))
n_joint_observations(pomg::BabyPOMG) = length(ordered_joint_observations(pomg))

function transition(pomg::BabyPOMG, s, a, s′)
    # Regardless, feeding makes the baby sated.
    if a[1] == FEED || a[2] == FEED
        if s′ == SATED
            return 1.0
        else
            return 0.0
        end
    else
        # If neither caretaker fed, then one of two things happens.
        # First, a baby that is hungry remains hungry.
        if s == HUNGRY
            if s′ == HUNGRY
                return 1.0
            else
                return 0.0
            end
        # Otherwise, it becomes hungry with a fixed probability.
        else
            probBecomeHungry = 0.5 #pomg.babyPOMDP.p_become_hungry
            if s′ == SATED
                return 1.0 - probBecomeHungry
            else
                return probBecomeHungry
            end
        end
    end
end

function joint_observation(pomg::BabyPOMG, a, s′, o)
    # If at least one caregiver sings, then both observe the result.
    if a[1] == SING || a[2] == SING
        # If the baby is hungry, then the caregivers both observe crying/silent together.
        if s′ == HUNGRY
            if o[1] == CRYING && o[2] == CRYING
                return pomg.babyPOMDP.p_cry_when_hungry_in_sing
            elseif o[1] == QUIET && o[2] == QUIET
                return 1.0 - pomg.babyPOMDP.p_cry_when_hungry_in_sing
            else
                return 0.0
            end
        # Otherwise the baby is sated, and the baby is silent.
        else
            if o[1] == QUIET && o[2] == QUIET
                return 1.0
            else
                return 0.0
            end
        end
    # Otherwise, the caregivers fed and/or ignored the baby.
    else
        # If the baby is hungry, then there's a probability it cries.
        if s′ == HUNGRY
            if o[1] == CRYING && o[2] == CRYING
                return pomg.babyPOMDP.p_cry_when_hungry
            elseif o[1] == QUIET && o[2] == QUIET
                return 1.0 - pomg.babyPOMDP.p_cry_when_hungry
            else
                return 0.0
            end
        # Similarly when it is sated.
        else
            if o[1] == CRYING && o[2] == CRYING
                return pomg.babyPOMDP.p_cry_when_not_hungry
            elseif o[1] == QUIET && o[2] == QUIET
                return 1.0 - pomg.babyPOMDP.p_cry_when_not_hungry
            else
                return 0.0
            end
        end
    end
end

function joint_reward(pomg::BabyPOMG, s, a)
    r = [0.0, 0.0]

    # Both caregivers do not want the child to be hungry.
    if s == HUNGRY
        r += [pomg.babyPOMDP.r_hungry, pomg.babyPOMDP.r_hungry]
    end

    # One caregiver prefers to feed.
    if a[1] == FEED
        r[1] += pomg.babyPOMDP.r_feed / 2.0
    elseif a[1] == SING
        r[1] += pomg.babyPOMDP.r_sing
    end

    # One caregiver prefers to sing.
    if a[2] == FEED
        r[2] += pomg.babyPOMDP.r_feed
    elseif a[2] == SING
        r[2] += pomg.babyPOMDP.r_sing / 2.0
    end

    # Note that caregivers only experience a cost if they do something.

    return r
end

joint_reward(pomg::BabyPOMG, b::Vector{Float64}, a) = sum(joint_reward(pomg, s, a) * b[s] for s in ordered_states(pomg))

function POMG(pomg::BabyPOMG)
    return POMG(
        pomg.babyPOMDP.γ,
        vec(collect(1:n_agents(pomg))),
        ordered_states(pomg),
        [ordered_actions(pomg, i) for i in 1:n_agents(pomg)],
        [ordered_observations(pomg, i) for i in 1:n_agents(pomg)],
        (s, a, s′) -> transition(pomg, s, a, s′),
        (a, s′, o) -> joint_observation(pomg, a, s′, o),
        (s, a) -> joint_reward(pomg, s, a)
    )
end


# This file is part of IntegerSequences.
# Copyright Peter Luschny. License is MIT.

(@__DIR__) ∉ LOAD_PATH && push!(LOAD_PATH, (@__DIR__))

module SeqUtils

using Nemo

export ModuleSeqUtils
export SeqShow, SeqPrint, SeqName, SeqNum, Println, ZArray

"""

Nemo is a library designed, developed and maintained by William Hart with the help of others. Many functions in our project are based on Nemo.
"""
const ModuleSeqUtils = ""

"""

The basic integer type of our project is 'fmpz'. Think of a 'fmpz' as a 'BigInt'. 'fmpz' stands for 'fast multiple precision zahl' (zahl=integer). The mathematical symbol for the ring of integers is the blackboard (double-struck) Z, also written fmpz. In reference to this fmpz(n) defines the integer n as a fmpz.
"""
function Nemofmpz(n::Int)
    return fmpz(n)
end

"""

Return an array of type `fmpz` of length ``n`` preset with ``0``.
"""
ZArray(len::Int) = fill(fmpz(0), len)

"""

Return an array of type fmpz and of size ``n`` filled by the values of the function ``f`` in the range `offset:n`.
"""
function ZArray(n::Int, f::Function, offset = 1)
    n ≤ 0 && return fmpz[]
    fmpz[f(k) for k ∈ offset:offset+n-1] # !
end

"""

Return the size of the array ``A``.
"""
SeqSize(A) = Base.length(A)

"""

Return the range of a SeqArray.
"""
SeqRange(A) = 1:SeqSize(A)

"""

Print the array ``A`` in the format ``n ↦ A[n]`` for n in the given range.
"""
function SeqShow(A::Array, R::AbstractRange, offset = 0)
    for k ∈ R
        if isassigned(A, k)
            println(k + offset, " ↦ ", A[k])
        else
            println(k + offset, " ↦ ", "undef")
        end
    end
end

"""

Print the array ``A`` in the format ``n ↦ A[n]`` for n in the range first:last.
"""
SeqShow(A, first::Int, last::Int) = SeqShow(A, first:last)

"""

Print the array ``A`` in the format ``n ↦ A[n]``.
"""
SeqShow(A::Array, offset = 0) = SeqShow(A, 1:length(A), offset - 1)

function print_without_type(io, v::AbstractVector)
    print(io, "[")
    for (i, el) ∈ enumerate(v)
        i > 1 && print(io, ", ")
        print(io, el)
    end
    println(io, "]")
end

"""

Print the array without typeinfo.
"""
Println(v::AbstractVector) = print_without_type(IOContext(stdout), v)

"""

Print the array without typeinfo.
"""
function Println(V::Array{Array{fmpz,1},1})
    for v ∈ V  v |> Println end
end

"""

Print the array with or without typeinfo.
"""
function SeqPrint(v::AbstractVector, typeinfo = false)
    typeinfo ? println(v) : Println(v)
end

"""

Valid prefixes to the numerical part of the OEIS A-numbers.

* C  => Coroutine (channel)
* F  => Filter (all below n)
* G  => Generating function
* I  => Iteration
* L  => List (array based)
* M  => Matrix (2-dim square)
* P  => Polynomial (over fmpz or QQ)
* R  => RealFunction (Float64)
* S  => Staircase (list iteration)
* T  => Triangle (list iteration)
* V  => Value (single term)
* is => is a (predicate), boolean
"""
const ValidPrefixes = ['C', 'F', 'G', 'I', 'L', 'M', 'R', 'S', 'T', 'V']

"""

Return the name of a OEIS sequence given a similar named function as a string.
"""
function SeqName(fun)
    aname = string(fun)
    for X ∈ ValidPrefixes
        aname = replace(aname, X => 'A')
    end

    if !occursin(r"^A[0-9]{6}$", aname)
        fullname = split(aname, ".")
        aname = String(fullname[end])
        if !occursin(r"^A[0-9]{6}$", aname)
            @error(aname, " is not a valid A-name!")
            return
        end
    end
    aname
end

"""

Return the A-number of a OEIS sequence given a similar named function as an integer.
"""
function SeqNum(seq)
    aname = SeqName(seq)
    parse(Int, aname[2:end])
end

#START-TEST-########################################################

using Test

function test() end

function demo()
    A = fmpz[1, 2, 3, 4, 5]

    println()
    SeqShow(A)
    println()

    SeqShow(A, -4)
    println()

    A |> Println

    println()
    V = Array{fmpz,1}[[1], [0, 1], [0, 1, 70], [0, 1, 990, 34650]]
    for v ∈ V  v |> Println end
end

function perf() end

function main()
    test()
    demo()
    perf()
end

main()

end # module

# a(f) = applicable(f, 1) ? f(1) : f(1,2)
# m(f) = first(methods(h)).nargs == 1 ? f(1) : f(1,2)


function randomAssignments(dataset::Dataset; seed::Union{Nothing, Int}=nothing)::Matrix{Int}
    """ Generates random assignments of n nodes to q groups

    Parameters
    ----------
    dataset : Dataset
        Dataset representing an observed graph.

    Returns
    -------
    x : Array{Int64,2}
        Matrix of assignments of nodes to groups
    """
    n = dataset.n
    q = dataset.n_communities
    if seed != nothing
        Random.seed!(seed)
    end

    assign = rand(1:q, n)

    x = zeros(Int, n, q)
    for i=1:n
        x[i,assign[i]] = 1
    end
    @assert sum(x) == n
    return x
end

function findImprovingRelocation(dataset::Dataset, w::Matrix{Float64}, x::Matrix{Int},
                                 reeval_w::Bool; best_imp::Bool=true, time_limit::Float64=400.0)
    """
    Parameters
    ----------
    dataset : Dataset
        Dataset representing an observed graph.
    x : Matrix{Int}
        Matrix of assignments of nodes to groups.
    w : Matrix{Float64}
        Matrix of probabilities of the SBM.
    reeval_w : Bool
        Whether to re-evaluate w to evaluate every possible move.
    best_imp : Bool
        Whether to take the best possible improving move (true) or to take the first improving move (false).
    """
    if best_imp == false
        throw("not implemented")
    end

    start = time(); # Start counting time

    n = dataset.n
    q = dataset.n_communities

    # Calculate objective value
    best_obj = calculateObjective(dataset, w, x)
    best_solution = (copy(w), copy(x))
    improved = false

    # loop through all nodes
    for i=1:n
        # Get current assignment of node i
        current_assign = argmax(x[i,:])
        for g=1:(q-1) # Find the best relocate move for node i

            # Break if time limit is exceeded
            availableTime = time_limit - (time() - start)
            if availableTime <= 0
                (w, x) = best_solution
                return improved, best_obj, w, x
            end

            # Set new assignment for node i
            new_assign = (current_assign + g)
            if new_assign > q
                new_assign = new_assign % q
            end

            # Reassign move
            x[i,current_assign] = 0
            x[i,new_assign] = 1

            # Evaluate move
            if reeval_w ##### Evaluate move on corresponding optimal w
                w = optimalProbMatrix(dataset, x)
            end
            #####
            obj_value = calculateObjective(dataset, w, x)
            if obj_value < best_obj
                best_obj = obj_value
                best_solution = (copy(w), copy(x))
                improved = true
            end

            # Undo reassign move
            x[i,current_assign] = 1
            x[i,new_assign] = 0
        end
    end

    # Reassign using the best move
    ##### Update matrix of probabilities w
    (w, x) = best_solution
    return improved, best_obj, w, x
end

function localSearchAssignments(estimator::SBMEstimator, dataset::Dataset, w::Union{Nothing,Matrix{Float64}}, x::Matrix{Int}; time_limit::Float64=400.0)
    """ Local Search in the space of the assignments x,
    by looking for the best relocate move.
    The value of w (probability matrix) is fixed and
    does not get updated inside this function

    Parameters
    ----------
    estimator : SBMEstimator
        Optimization method specifications.
    dataset : Dataset
        Dataset representing an observed graph.
    w : Union{Nothing,Matrix{Float64}}
        Matrix of probabilities. If nothing, then the optimal w will be calculated for the given x
    x : Matrix{Int}
        Matrix of assignments of nodes to groups.
    time_limit : Float64
        Time limit in seconds.

    Returns
    -------
    x : Matrix{Int}
        Matrix of assignments of nodes to groups (local optimum).
    """

    # Check if a matrix w was given
    if w == nothing
        w = optimalProbMatrix(dataset, x)
    end

    start = time(); # Start counting time

    improved = true
    while improved
        improved = false
        # Break if time limit is exceeded
        availableTime = time_limit - (time() - start)
        if availableTime <= 0
            break
        end
        improved, best_obj, w, x = findImprovingRelocation(dataset, w, x, false, best_imp=true, time_limit=availableTime)
    end
    return x
end


#***********************************************************************
#  Copyright 2006 John A. Trangenstein
#
#  This software is made available for research and instructional use
#  only.
#  You may copy and use this software without charge for these
#  non-commercial purposes, provided that the copyright notice and
#  associated text is reproduced on all copies.
#  For all other uses (including distribution of modified versions),
#  please contact the author at
#    John A. Trangenstein
#    Department of Mathematics
#    Duke University
#    Durham, NC 27708-0320
#    USA
#  or
#    johnt@math.duke.edu
#
#  This software is made available "as is" without any assurance that it
#  is completely correct, or that it will work for your purposes.
#  Use the software at your own risk.
#***********************************************************************

function consdiff(fic::Int,lac::Int,fif::Int,laf::Int,fix::Int,lax::Int,
                  ifirst::Int,ilast::Int,
                  x::OffsetArray{Float64,1},
                  flux::OffsetArray{Float64,1},
                  conservd::OffsetArray{Float64,1})
#   ******************************************************************
#   update conservd to new time
#   ******************************************************************
    @inbounds for ic=ifirst:ilast
        conservd[ic] -= (flux[ic+1]-flux[ic]) / (x[ic+1]-x[ic])
    end
end

function stabledt(fc::Int,lc::Int,fm::Int,lm::Int,ifirst::Int,ilast::Int,
                  conserved::OffsetArray{Float64,1},
                  lambda_cell::OffsetArray{Float64,1},
                  x::OffsetArray{Float64,1})
#   ******************************************************************
#   compute stable timestep
#   ******************************************************************
    sdt=huge
    @inbounds for ic=ifirst:ilast
        abs_lambda=abs(lambda_cell[ic])
        dx = x[ic+1]-x[ic]
        if abs_lambda > roundoff*dx
            sdt=min(sdt,dx/abs_lambda)
        end
    end
    @assert sdt > 0.0
    sdt
end