"""
     fdhinfminus(sys,freq) -> (β, ind, fr)

Compute for a stable descriptor system `sys = (A-λE,B,C,D)` the `H∞-` index `β` of its
transfer function matrix `G(λ)`. If `freq = missing` (default), then `β` is the
minimum `H∞-norm` of the columns of `G`, `ind` is the index of the minimum-norm column and `fr` is
the frequency where the minimum `H∞-norm` of the columns is achieved. If `freq` is a real value or
a real vector of frequency values, then `β` is the minimum of the 2-norms of the columns of the
frequency responses of `G` evaluated for all values contained in `freq`, `ind` is the index of column
for which the minimum is achieved and `fr` is the corresponding frequency.
"""
function fdhinfminus(sys::DescriptorStateSpace{T}, freq::Union{AbstractVector{<:Real},Real,Missing} = missing) where T
   p, m = size(sys)

   m == 0 && (return T[], Int[], T[])

   isstable(sys) || error("the system is unstable")
   β = Inf; ind = 0; f = -1;
   if ismissing(freq)
      # evaluate β as the minimum of H-infinity norms of the columns of G
      for j = 1:m
         temp, fr = ghinfnorm(sys[:,j])
         β <= temp || (β = temp; ind = j; f = fr)
      end
   else
      # evaluate β as the minimum of the norms of columns of the frequency
      # responses of G evaluated over all frequencies contained in FREQ
      if !isa(freq, Vector)
         # use evalfr if only one frequency is present
         H = evalfr(sys; fval = freq)
         for j = 1:m
             temp = norm(view(H,:,j))
             β <= temp || (β = temp; ind = j)
         end
         f = freq
      else
         T1 = T <: BlasFloat ? T : promote_type(Float64,T)
         ONE = one(T1)
         a, e, b, c, d = dssdata(T1,sys)
         Ts = abs(sys.Ts)
         disc = !iszero(Ts)
         sw = disc ? im*Ts : im
         desc = !(e == I)
         # Determine the complex Hessenberg-form system to be used for efficient
         # frequency response computation.
         ac, ec, bc, cc, dc = chess(a, e, b, c, d)
         H = similar(dc, eltype(dc), p, m)
         bct = similar(bc)
         for i = 1:length(freq)
             if isinf(freq[i])
                # exceptional call to evalfr
                H = evalfr(sys, fval=Inf)
             else
                copyto!(H, dc)
                copyto!(bct, bc)
                w = disc ? -exp(sw*freq[i]) : -sw*freq[i]
                desc ? ldiv!(UpperHessenberg(ac+w*ec),bct) : ldiv!(ac,bct,shift = w)
                mul!(H, cc, bct, -ONE, ONE)
             end
             for j = 1:m
                 temp = norm(view(H,:,j))
                 β <= temp || (β = temp; ind = j; f = freq[i])
             end
         end
      end
   end
   return β, ind, f
end
"""
     fdhinfmax(sys,freq) -> (γ, ind, fr)

Compute for a descriptor system `sys = (A-λE,B,C,D)`, `γ` - the maximum norm of the columns of its
transfer function matrix `G(λ)`. If `freq = missing` (default), then `γ` is the
maximum `H∞-norm` of the columns of `G`, `ind` is the index of the maximum-norm column and `fr` is
the frequency where the maximum `H∞-norm` of the columns is achieved. If `freq` is a real value or
a real vector of frequency values, then `γ` is the maximum of the 2-norms of the columns of the
frequency responses of `G` evaluated for all values contained in `freq`, `ind` is the index of column
for which the maximum is achieved and `fr` is the corresponding frequency.
"""
function fdhinfmax(sys::DescriptorStateSpace{T}, freq::Union{AbstractVector{<:Real},Real,Missing} = missing) where T
   p, m = size(sys)

   m == 0 && (return T[], Int[], T[])

   γ = 0; ind = 0; f = -1;
   if ismissing(freq)
      # evaluate γ as the maximum of H-infinity norms of the columns of G
      for j = 1:m
         temp, fr = ghinfnorm(sys[:,j])
         γ >= temp || (γ = temp; ind = j; f = fr)
      end
   else
      # evaluate γ as the maximum of the norms of columns of the frequency
      # responses of G evaluated over all frequencies contained in FREQ
      if !isa(freq, Vector)
         # use evalfr if only one frequency is present
         H = evalfr(sys; fval = freq)
         for j = 1:m
             temp = norm(view(H,:,j))
             γ >= temp || (γ = temp; ind = j)
         end
         f = freq
      else
         T1 = T <: BlasFloat ? T : promote_type(Float64,T)
         ONE = one(T1)
         a, e, b, c, d = dssdata(T1,sys)
         Ts = abs(sys.Ts)
         disc = !iszero(Ts)
         sw = disc ? im*Ts : im
         desc = !(e == I)
         # Determine the complex Hessenberg-form system to be used for efficient
         # frequency response computation.
         ac, ec, bc, cc, dc = chess(a, e, b, c, d)
         H = similar(dc, eltype(dc), p, m)
         bct = similar(bc)
         for i = 1:length(freq)
             if isinf(freq[i])
                # exceptional call to evalfr
                H = evalfr(sys, fval=Inf)
             else
                copyto!(H, dc)
                copyto!(bct, bc)
                w = disc ? -exp(sw*freq[i]) : -sw*freq[i]
                desc ? ldiv!(UpperHessenberg(ac+w*ec),bct) : ldiv!(ac,bct,shift = w)
                mul!(H, cc, bct, -ONE, ONE)
             end
             for j = 1:m
                 temp = norm(view(H,:,j))
                 γ >= temp || (γ = temp; ind = j; f = freq[i])
             end
         end
      end
   end
   return γ, ind, f
end
function binmat2dec(S::Union{BitArray,VecOrMat{<:Number}})
# Convert a binary matrix with zero/non-zero elements to an equivalent decimal number integer vector
   len = size(S,2)
   n = Int.(S[:,1] .!= 0)
   for i = 2:len
       n .= 2*n .+ Int.(S[:,i] .!= 0)
   end
   return n
end
function dec2binmat(n::Union{Int,Vector{Int}}, numbits::Int = 1)
# Convert a decimal non-negative integer or integer vector with decimal non-negative values
# to their binary matrix representation with `numbits` bits.
   minimum(n) < 0 && error("only non-negative values allowed")
   maxn = maximum(n)
   len = maxn == 0 ? max(1,numbits) : max(Int(floor(log2(maxn)+1)),numbits)
   ns = size(n,1)
   v = isa(n,Int) ? string.([n], base = 2, pad = len) : string.(n, base = 2, pad = len)
   S = falses(ns,len)
   for i = 1:ns
       for j = 1:len
           S[i,j] = parse(Int,v[i][j]) == 1
       end
   end
   return S
end
"""
    S = fditspec_(sysrf::DescriptorStateSpace; FDfreq = missing, block = false, poleshift = false,
                 FDtol, FDStol, atol = 0, atol1 = atol, atol2 = atol, rtol, fast = true)

Compute the weak or strong binary structure matrix `S` of the transfer function matrix of a
linear time-invariant system `sysrf`
(typically representing the transfer channel from the fault inputs to residuals).
`sysrf` has a descriptor system realization of the form `sysrf = (Af-lambda*Ef,Bf,Cf,Df)`
with a  `q x mf` transfer function matrix `Rf(λ)`.
For the description of keyword parameters see the documentation of [`fditspec`](@ref).
"""
function fditspec_(sysrf::DescriptorStateSpace{T}; FDfreq::Union{AbstractVector{<:Real},Real,Missing} = missing,
                  block::Bool = false, poleshift::Bool = false, FDtol::Real = 0., FDStol::Real = 0.,
                  atol::Real = zero(float(real(T))), atol1::Real = atol, atol2::Real = atol,
                  rtol::Real =  (size(sysrf.A,1)+1)*eps(float(one(real(T))))*iszero(max(atol1,atol2)), fast::Bool = true) where T
   p, mf = size(sysrf)
   n = order(sysrf)
   q = block ? 1 : p  # number of rows of S
   Sstrong = !ismissing(FDfreq)
   if Sstrong
      isa(FDfreq,Vector) || (FDfreq = [FDfreq])
      lfreq = length(FDfreq);
      w = im*FDfreq;                   # w = j*freq
      Ts = abs(sysrf.Ts);
      Ts > 0 && ( w = exp(Ts*w))     # w = exp(j*Ts*freq)
      S = falses(q,mf,lfreq)
   else
      S = falses(q,mf)
   end
   a, e, b, c, d = dssdata(sysrf)
   standard = isequal(e,I)
   if n == 0
      S1 =  abs.(d) .> (FDtol <= 0 ? 0.0001*max(1., norm(d,1)) : FDtol)
      S = block ? maximum(S1, dims=1) : S1
      return Sstrong ? repeat(S,1,1,lfreq) : S
   end
   if Sstrong
      #FDStol <= 0. && (FDStol = 0.0001*max(1., norm(a,1), norm(b,1), norm(c,Inf), norm(d,1), standard ? 0 : norm(e,1)))
      FDStol <= 0. && (FDStol = eps(max(1., norm(a,1), norm(b,1), norm(c,Inf), norm(d,1), standard ? 0 : norm(e,1))))
      S = trues(q, mf, lfreq)
      p1, n1 = size(c,1), size(c,2)
      # employ structural analysis to compute weak/strong structure matrix
      ispole = false
      for k = 1:lfreq
         # check if freq(k) is a pole
         ispole = isinf(FDfreq[k]) ? (e == I ? false : rank(e) < n1) : rank(a-w[k]*e) < n1
         ispole || break
      end
      ispole && !poleshift && error("sysf has poles in Ω")
      if block
         ispole && (f = rand(T,n1,p1); mul!(a,f,c,1,1); mul!(b,f,d,1,1) )
         for j = 1:mf
             # elliminate uncontrollable and unobservable eigenvalues for the j-th column of B
             a1, e1, b1, c1, d1 = lsminreal(a, e, view(b,:,j), c, view(d,:,j); noseig = false, atol1, atol2, rtol, fast)
             for k = 1:lfreq
                # check if freq(k) is a zero of the j-th column
                s = isinf(FDfreq[k]) ? svdvals!([e1 b1; c1 d1]) : svdvals!([a1-w[k]*e1 b1; c1 d1])
                S[1, j, k] = (s[end] > FDStol)
            end
         end
      else
         for i = 1:p
             # elliminate unobservable eigenvalues for the i-th row of C
             a1, e1, b1, c1, d1 = lsminreal(a, e, b, view(c,i:i,:), view(d,i:i,:); contr = false, noseig = false, atol1, atol2, rtol, fast)
             n1 = size(a1,1)
             ispole && (f1 = rand(T,n1,1); mul!(a1,f1,c1,1,1); mul!(b1,f1,d1,1,1) )
             for j = 1:mf
                 # elliminate uncontrollable and non-dynamic eigenvalues for the j-th column of B
                 a2, e2, b2, c2, d2 = lsminreal(a1, e1, view(b1,:,j:j), c1, view(d1,:,j:j); obs = false, noseig = true, atol1, atol2, rtol, fast)
                 for k = 1:lfreq
                    # check if freq(k) is a zero of the (i,j)-th element
                    s = isinf(FDfreq[k]) ? svdvals!([e2 b2; c2 d2]) : svdvals!([a2-w[k]*e2 b2; c2 d2])
                    S[i,j,k] = s[end] > FDStol
                 end
             end
         end
      end
   else
      FDtol <= 0. && (FDtol = 0.0001*max(1., norm(b,1), norm(d,1)))
      standard && all(abs.(d) .> FDtol) && (return trues(q,mf))
      S = falses(q, mf)
      if block
         for j = 1:mf
             # compute minimal realization of the j-th column of Rf
             _, _, b1, _, d1 = lsminreal(a, e, view(b,:,j:j), c, view(d,:,j:j); atol1, atol2, rtol, fast)
             S[1,j] |=  (any(abs.(view(d1,:,1)) .> FDtol) || any(abs.(view(b1,:,1)) .> FDtol))
         end
      else
         for j = 1:mf
             # elliminate uncontrollable eigenvalues for the j-th column of B
             a1, e1, b1, c1, d1 = lsminreal(a, e, view(b,:,j), c, view(d,:,j); obs = false, noseig = false, atol1, atol2, rtol, fast)
             for i = 1:p
                # elliminate unobservable and non-dynamic eigenvalues for the (i,j)-th element
                _, _, b2, _, d2 = lsminreal(a1, e1, b1, view(c1,i:i,:),view(d1,i:i,:); contr = false, noseig = true, atol1, atol2, rtol, fast)
                S[i,j] |=  (abs(d2[1,1]) > FDtol || any(abs.(view(b2,:,1)) .> FDtol))
             end
         end
      end
   end
   return S
end
"""
     fdisspec_(sysrf::DescriptorStateSpace, freq; block = false, stabilize = false, FDGainTol = 0.01,
                     atol, atol1, atol2, atol3, rtol, fast = true) -> (S, gains)

Compute the strong binary structure matrix `S` of the transfer function matrix of a
linear time-invariant system `sysrf`
(typically representing the transfer channel from the fault inputs to residuals).
`sysrf` has a descriptor system realization of the form `sysrf = (Af-lambda*Ef,Bf,Cf,Df)`
with a  `q x mf` transfer function matrix `Rf(λ)`.
For the description of keyword parameters see the documentation of [`fdisspec`](@ref).
"""
function fdisspec_(sysrf::DescriptorStateSpace{T}, freq::Union{AbstractVector{<:Real},Real} = 0; stabilize::Bool = false, FDGainTol::Real = 0.01, block::Bool = false,
                   atol::Real = zero(float(real(T))), atol1::Real = atol, atol2::Real = atol, atol3::Real = atol,
                   rtol::Real =  ((max(size(sysrf.A)...))+1)*eps(float(one(real(T))))*iszero(max(atol1,atol2)), fast::Bool = true) where T
   p, mf = size(sysrf)
   isa(freq,Vector) || (freq = [freq])
   lfreq = length(freq);
   w = im*freq;                   # w = j*freq
   Ts = abs(sysrf.Ts);
   Ts > 0 && ( w = exp(Ts*w))     # w = exp(j*Ts*freq)

   ispole = false
   n = size(sysrf.A,1)
   for k = 1:lfreq
      # check if freq(k) is a pole
      ispole = isinf(freq[k]) ? (sysrf.E == I ? false : rank(sysrf.E) < n) : rank(sysrf.A-w[k]*sysrf.E) < n
      ispole || break
   end
   ispole && !stabilize && error("sysf has poles in Ω")
   if block
      stabilize && (sysrf = glcf(sysrf; atol1, atol2, atol3, rtol)[1])
      # gs = evalfr(sysrf, w[1]; atol1, atol2, rtol, fast)
      # any(isinf.(gs)) && error("fdisspec:pole - the frequency $(w[1]) is a system pole")
      smat = falses(1, mf, lfreq)
      gains = zeros(T, 1, mf, lfreq)
      for i = 1:lfreq
          gs = evalfr(sysrf, w[i]; atol1, atol2, rtol, fast)
         any(isinf.(gs)) && error("fdisspec_:pole - the frequency $(w[i]) is a system pole")
          for j = 1:mf
              gsj = norm(view(gs,:,j))
              gains[1,j,i] = gsj
              smat[1,j,i] = (gsj .> FDGainTol)
          end
      end
   else
      smat = falses(p, mf, lfreq)
      gains = zeros(T, p, mf, lfreq)
      for k = 1:p
         t = gir(sysrf[k,:]; contr = false, atol1, atol2, rtol)
         stabilize && (t = glcf(t; atol1, atol2, atol3, rtol)[1])
         for i = 1:lfreq
            #  gs = stabilize ? abs.(evalfr(glcf(t; atol1, atol2, atol3, rtol)[1], w[i]; atol1, atol2, rtol, fast)) :
            #                   abs.(evalfr(t, w[i]; atol1, atol2, rtol, fast))
             gs = abs.(evalfr(t, w[i]; atol1, atol2, rtol, fast))
             any(isinf.(gs)) && error("fdisspec_:pole - the frequency $(w[i]) is a system pole")
             gains[k,:,i] = gs
             smat[k,:,i] = (gs .> FDGainTol)
         end
      end
   end
   return smat, gains
end
"""
     fdiscond_(sysrf::DescriptorStateSpace, freq) -> (scond, β, γ)

Compute for a stable descriptor system `sysrf = (A-λE,B,C,D)` with the transfer function matrix `Rf(λ)`,
`β` - the H∞- index of `Rf(λ)`, `γ` - the maximum of the columns norms of `Rf(λ)` and
`scond` - the column-gains sensitivity condition evaluated as `scond := β/γ`.
If `freq` is a vector of real frequency values, then `β` and `γ`
are evaluated over the frequencies contained in `freq`.
"""
function fdiscond_(sysrf::DescriptorStateSpace{T}, freq::Union{AbstractVector{<:Real},Real,Missing} = missing) where T
   p, m = size(sysrf)

   m == 0 && (return T[], T[], T[])

   isstable(sysrf) || error("the system is unstable")

   β = Inf
   γ = 0
   if ismissing(freq)
      # evaluate β and γ as the minimum and maximum of H-infinity norms of the columns of G
      for j = 1:m
         temp = ghinfnorm(sysrf[:,j])[1]
         γ >= temp || (γ = temp)
         β <= temp || (β = temp)
      end
   else
      # evaluate β and γ as the minimum and maximum of the norms of columns of the frequency
      # responses of G evaluated over all frequencies contained in FREQ
      if !isa(freq, Vector)
         # use evalfr if only one frequency is present
         H = evalfr(sysrf; fval = freq)
         for j = 1:m
             temp = norm(view(H,:,j))
             γ >= temp || (γ = temp)
             β <= temp || (β = temp)
         end
      else
         T1 = T <: BlasFloat ? T : promote_type(Float64,T)
         ONE = one(T1)
         a, e, b, c, d = dssdata(T1,sysrf)
         Ts = abs(sysrf.Ts)
         disc = !iszero(Ts)
         sw = disc ? im*Ts : im
         desc = !(e == I)
         # Determine the complex Hessenberg-form system to be used for efficient
         # frequency response computation.
         ac, ec, bc, cc, dc = chess(a, e, b, c, d)
         H = similar(dc, eltype(dc), p, m)
         bct = similar(bc)
         for i = 1:length(freq)
             if isinf(freq[i])
                # exceptional call to evalfr
                H = evalfr(sys, fval=Inf)
             else
                copyto!(H, dc)
                copyto!(bct, bc)
                w = disc ? -exp(sw*freq[i]) : -sw*freq[i]
                desc ? ldiv!(UpperHessenberg(ac+w*ec),bct) : ldiv!(ac,bct,shift = w)
                mul!(H, cc, bct, -ONE, ONE)
             end
             for j = 1:m
                 temp = norm(view(H,:,j))
                 γ >= temp || (γ = temp)
                 β <= temp || (β = temp)
             end
         end
      end
   end
   return β/γ, β, γ
end


"""
The inverse logit function, which is also the CDF of the Logistic(0, 1)
distribution. Note that this function will generate an exact value of 0 when
z <= -710 and will generate an exact value of 1 when z >= 37. This introduces
numeric problems if you provide inputs outside of that range.
"""
@inline invlogit(z::Real) = 1 / (1 + exp(-z))

"""
The logit function, which is also the quantile function of the Logistic(0, 1)
distribution.
"""
@inline logit(p::Real) = log(p / (1 - p))


# Bifurcation diagram for 1D climate model
using DrWatson
@quickactivate "NonlinearDynamicsTextbook"
include(srcdir("style.jl"))
using DynamicalSystems, PyPlot

include(srcdir("simple_continuation.jl"))

using ForwardDiff, PyPlot
αtan(T) = 0.5 - 0.2*tanh((T-263)/4)
dTdt(T, ε = 0.65, α=αtan, s= 1.0) = s*(1 - α(T)) - 1.6e-10 * ε * T^4
dTdt(T; ε = 0.65, α=αtan, s = 1.0) = dTdt(T, ε, α, s)
d²Tdt²(T, ε) = ForwardDiff.derivative(T -> dTdt(T, ε), T)
climate_f(u, p, t) = SVector(dTdt(u[1], p))
climate_J(u, p, t) = SMatrix{1,1}(d²Tdt²(u[1], p))

u0 = SVector(225.0)
p0 = 0.9
ds = ContinuousDynamicalSystem(climate_f, u0, p0, climate_J)

f, J = bifurcation_rule_form(ds)

pmin = 0.3
pmax = 1

dp0 = 0.05
dx0 = SVector(0.1)

xs, ps, stability = continuation(f, J, u0, p0;
    pmin, pmax, dp0, dx0, N = 1500
)

xs = [x[1] for x in xs]

# %% make figure
stable1 = findfirst(!, stability)
unstable1 = findlast(!, stability)

fig = figure(figsize = (figx/2, figy))
ax = gca()

plot(ps[1:stable1], xs[1:stable1]; c = "C2", ls = "solid", label="stable")
plot(ps[stable1+1:unstable1], xs[stable1+1:unstable1]; c = "C2", ls = "dashed", label="unstable")
plot(ps[unstable1+1:end], xs[unstable1+1:end]; c = "C2", ls = "solid")

# scatter bifurcation points
ax.plot([ps[stable1], ps[unstable1]], [xs[stable1], xs[unstable1]], ls = "None",
		marker="o", ms = 12, color = "k", mew = 1, mec = "k",
        fillstyle = "top", markerfacecoloralt= "w"
)

# plot arrows
_a = 300
arrow1ε = ps[_a:stable1] .- 0.03
arrow1T = xs[_a:stable1] .- 5
plot(arrow1ε, arrow1T; color = "C1", lw = 4.0)
arrow(arrow1ε[end], arrow1T[end], 0, 50;
width = 0.004, head_length = 5.0, color = "C1", zorder = 99)

_b = 200
arrow2ε = ps[unstable1:unstable1+_b] .+ 0.03
arrow2T = xs[unstable1:unstable1+_b] .+ 5

plot(arrow2ε, arrow2T; color = "C0", lw = 4.0)
arrow(arrow2ε[1], arrow2T[1], 0, -50;
width = 0.004, head_length = 5.0, color = "C0", zorder = 99)

# scatter(ps[stability], xs[stability]; c = "C3", legend = "stable")
# scatter(ps[stability], xs[stability]; c = "C3", legend = "stable")
legend()
xticks(0.3:0.2:0.9)
yticks(200:50:350)
ylabel("\$T^*\$", labelpad = -20)
xlabel("\$\\varepsilon\$"; labelpad=-20)
# yticks(230:50:310)
# ylim(210, 320)
# ax.xaxis.set_label_coords(0.4, -0.025)
fig.tight_layout(pad=0.3)
wsave(plotsdir("4", "bif_example"), fig)


# MaxPlus Semiring Implementation
# MaxPlus is over real numbers with max as addition and + as multiplication
# The zero element is -inf and the one element is 0

struct MaxPlusSemiringElement{T<:Real} <: AbstractSemiringElement{T}
    val::T
end

# interface methods
val(x::MaxPlusSemiringElement) = x.val
+(l::MaxPlusSemiringElement, r::MaxPlusSemiringElement) = MaxPlusSemiringElement(max(l.val, r.val))
*(l::MaxPlusSemiringElement, r::MaxPlusSemiringElement) = MaxPlusSemiringElement(l.val+r.val)
zero(::Type{MaxPlusSemiringElement{T}}) where T = MaxPlusSemiringElement(-Inf)
one(::Type{MaxPlusSemiringElement{T}}) where T = MaxPlusSemiringElement{T}(0)

iscommutative(::Type{<:MaxPlusSemiringElement}) = true
isidempotent(::Type{<:MaxPlusSemiringElement}) = true
isstar(::Type{<:MaxPlusSemiringElement}) = true

# MinPlus Semiring Implementation
# MinPlus is over real numbers with min as addition and + as multiplication
# The zero element is +inf and the one element is 0

struct MinPlusSemiringElement{T<:Real} <: AbstractSemiringElement{T}
    val::T
end

# interface methods
val(x::MinPlusSemiringElement) = x.val
+(l::MinPlusSemiringElement, r::MinPlusSemiringElement) = MinPlusSemiringElement(min(l.val, r.val))
*(l::MinPlusSemiringElement, r::MinPlusSemiringElement) = MinPlusSemiringElement(l.val+r.val)
zero(::Type{MinPlusSemiringElement{T}}) where T = MinPlusSemiringElement(Inf)
one(::Type{MinPlusSemiringElement{T}}) where T = MinPlusSemiringElement{T}(0)

iscommutative(::Type{<:MinPlusSemiringElement}) = true
isidempotent(::Type{<:MinPlusSemiringElement}) = true
isstar(::Type{<:MinPlusSemiringElement}) = true

# Real (or Probabilistic) Semiring Implementation
# The Real Semiring is over real numbers with + as addition and * as multiplication
# The zero element is 0 and the one element is 1

struct RealSemiringElement{T<:Real} <: AbstractSemiringElement{T}
    val::T
end

# interface methods
val(x::RealSemiringElement) = x.val
+(l::RealSemiringElement, r::RealSemiringElement) = RealSemiringElement(l.val+r.val)
*(l::RealSemiringElement, r::RealSemiringElement) = RealSemiringElement(l.val*r.val)
zero(::Type{RealSemiringElement{T}}) where T = RealSemiringElement(0)
one(::Type{RealSemiringElement{T}}) where T = RealSemiringElement(1)

iscommutative(::Type{<:RealSemiringElement}) = true
isstar(::Type{<:RealSemiringElement}) = true

# optional methods
-(x::RealSemiringElement) = RealSemiringElement(-x.val)
/(l::RealSemiringElement, r::RealSemiringElement) = RealSemiringElement(l.val/r.val)
function star(x::RealSemiringElement)
    0 <= x.val < 1 ? RealSemiringElement(1/(1-x.val)) : throw(DomainError("$x must be within [0, 1)"))
end

for op in (:conj,) # required for matrix inversion?
    @eval begin
        $op(x::RealSemiringElement) = RealSemiringElement($op(x.val))
    end
end

# Boolean Semiring Implementation
# The Boolean Semiring is over {True, False} with OR as addition and AND as multiplication
# The zero element is False and the one element is True

struct BooleanSemiringElement{Bool} <: AbstractSemiringElement{Bool}
    val::Bool
end

# interface methods
val(x::BooleanSemiringElement) = x.val
+(l::BooleanSemiringElement, r::BooleanSemiringElement) = BooleanSemiringElement(l.val | r.val)
*(l::BooleanSemiringElement, r::BooleanSemiringElement) = BooleanSemiringElement(l.val & r.val)
zero(::Type{BooleanSemiringElement{T}}) where T = BooleanSemiringElement{T}(false)
one(::Type{BooleanSemiringElement{T}}) where T = BooleanSemiringElement{T}(true)

iscommutative(::Type{<:BooleanSemiringElement}) = true
isidempotent(::Type{<:BooleanSemiringElement}) = true
isstar(::Type{<:BooleanSemiringElement}) = true

#optional methods
star(x::R) where R<:BooleanSemiringElement = one(R)


### A Pluto.jl notebook ###
# v0.11.14

using Markdown
using InteractiveUtils

# ╔═╡ eb68872e-01ee-11eb-22e4-b1f20b71faff
begin
	import Pkg
	Pkg.activate(mktempdir())
	Pkg.add("Combinatorics")
	using Combinatorics
end

# ╔═╡ 9a4603ea-01f2-11eb-00f0-591422bc7142
md" Given 20 people, what is the probability that,
among the 12 months in the year, there are 4
months containing exactly 2 birthdays and 4 containing exactly 3 birthdays?

(For simplicity, assume that it is equally like that a birthday falls in one of the 12 months.)
(from A First Course in Probablity - Sheldon Ross)

"

# ╔═╡ 0cfeea3c-01ee-11eb-26f1-f1815eccdd10
md"Create twenty birthday months"

# ╔═╡ 015393e0-01ee-11eb-2feb-d7db4236368a
months = rand(1:12,20)

# ╔═╡ 1a6a603e-01ee-11eb-39f6-efc09168a3ba
counts =  [sum(months.==i) for i=1:12]

# ╔═╡ 67c25428-0251-11eb-04b0-cb201bbefa6b
sum(counts)

# ╔═╡ 6f1890a6-01ee-11eb-1524-1da938de93f3
sum(counts.==2),sum(counts.==3)

# ╔═╡ af424780-01ee-11eb-3d5f-4f4ba6a372b4
sum(counts.==2)==4 & sum(counts.==3)==4

# ╔═╡ d8b5158e-01ee-11eb-393f-0d51ea1fc19b
function birthday_problem(n)
  success = 0
  for t=1:n
     months = rand(1:12,20)
     counts = [sum(months.==i) for i=1:12]
     success += sum(counts.==2) == 4  &&  sum(counts.==3)==4
   end
   success/n
end

# ╔═╡ dea5c880-01ee-11eb-066d-8105b2196548
birthday_problem(20_000_000)

# ╔═╡ d74e2646-01f2-11eb-2b25-158bf9419ecd
@elapsed birthday_problem(20_000_000)

# ╔═╡ 006a84ee-01ef-11eb-2db7-5792cb4d0198
multinomial(4,4,4) * multinomial(2,2,2,2,3,3,3,3) / 12.0^20

# ╔═╡ Cell order:
# ╟─9a4603ea-01f2-11eb-00f0-591422bc7142
# ╟─0cfeea3c-01ee-11eb-26f1-f1815eccdd10
# ╠═015393e0-01ee-11eb-2feb-d7db4236368a
# ╠═1a6a603e-01ee-11eb-39f6-efc09168a3ba
# ╠═67c25428-0251-11eb-04b0-cb201bbefa6b
# ╠═6f1890a6-01ee-11eb-1524-1da938de93f3
# ╠═af424780-01ee-11eb-3d5f-4f4ba6a372b4
# ╠═d8b5158e-01ee-11eb-393f-0d51ea1fc19b
# ╠═dea5c880-01ee-11eb-066d-8105b2196548
# ╠═d74e2646-01f2-11eb-2b25-158bf9419ecd
# ╠═eb68872e-01ee-11eb-22e4-b1f20b71faff
# ╠═006a84ee-01ef-11eb-2db7-5792cb4d0198


# ------------------------------------------------------------------------------------------
# # 循环
#
# 话题：
# 1. `while`循环
# 2. `for`循环
#
#
# ## `while` 循环
#
# `while` 循环语法：
#
# ```julia
# while <condition>
#     <loop body>
# end
# ```
#
# 举个例子，可以使用 `while` 来数数或者遍历一个数组。
# ------------------------------------------------------------------------------------------

n = 0
while n < 10
    n += 1
    println(n)
end
n

myfriends = ["Ted", "Robyn", "Barney", "Lily", "Marshall"]

i = 1
while i <= length(myfriends)
    friend = myfriends[i]
    println("Hi $friend, it's great to see you!")
    i += 1
end

# ------------------------------------------------------------------------------------------
# ## `for` 循环
#
# `for` 循环语法：
# ```julia
# for <var> in <loop iterable>
#     <loop body>
# end
# ```
#
# 上面两个例子都可以用 `for` 循环改写：
# ------------------------------------------------------------------------------------------

for n in 1:10
    println(n)
end

myfriends = ["Ted", "Robyn", "Barney", "Lily", "Marshall"]

for friend in myfriends
    println("Hi $friend, it's great to see you!")
end

# ------------------------------------------------------------------------------------------
# 让我们来试试用 `for` 循环创建加和表，表的每个元素为其行索引和列索引的和。
#
# 首先，来初始化一个全零的数组（矩阵）。
# ------------------------------------------------------------------------------------------

m, n = 5, 5
A = fill(0, (m, n))

for i in 1:m
    for j in 1:n
        A[i, j] = i + j
    end
end
A

# ------------------------------------------------------------------------------------------
# 趁机介绍一下像上面这种嵌套 `for` 循环的语法糖
# ------------------------------------------------------------------------------------------

B = fill(0, (m, n))

for i in 1:m, j in 1:n
    B[i, j] = i + j
end
B

# ------------------------------------------------------------------------------------------
# 创建同样的表但更具有“Julia”风格的方法是使用**数组解析式（array comprehension）**。
# ------------------------------------------------------------------------------------------

C = [i + j for i in 1:m, j in 1:n]

# ------------------------------------------------------------------------------------------
# ## 练习
#
# ### 4.1
# 循环遍历 1 到 100 的整数并打印它们的平方。
# ------------------------------------------------------------------------------------------

# 在此作答

# ------------------------------------------------------------------------------------------
# ### 4.2
# 给上一题的代码添一点以创建一个字典，使字典存着每个整数和它的平方的键值对，就像这样
#
# ```julia
# squares[10] == 100
# ```
# ------------------------------------------------------------------------------------------

# 在此作答

@assert squares[10] == 100
@assert squares[11] == 121

# ------------------------------------------------------------------------------------------
# ### 4.3
# 用数组解析式来创建一个数组 `squares_arr`，让它存着从1到100的整数的平方。
# ------------------------------------------------------------------------------------------

squares_arr = [i for i in 1:100]

@assert length(squares_arr) == 100
@assert sum(squares_arr) == 338350

# ------------------------------------------------------------------------------------------
# 请在完成练习运行 `@assert` 块以检测结果是否正确，若无报错即为正确。
# ------------------------------------------------------------------------------------------


println("*******************************************************")
@testset "fgmres" begin
	@testset "real matrix" begin
		A  = sprandn(100,100,.1) + SparseMatrixCSC(10.0I, 100, 100)
		n  = size(A,2)
		D  = diag(A)
		M2 = x -> Vector(D.\x)
		rhs = randn(100)
		tol = 1e-6;

		# test printing and behaviour for early stopping
		xtt = fgmres(A,rhs ,3,tol=1e-12,maxIter=3,out=2,storeInterm=true)
		@test xtt[2]==-1

		# test behaviour for zero rhs
		xtt = fgmres(A,0*rhs,5,tol=tol,maxIter=10,out=2)
		@test xtt[2]==-9
		@test all(xtt[1].==0)
		@test length(xtt[1])==100
		@test eltype(xtt[1])==eltype(rhs)

		x1 = fgmres(A,rhs ,5,tol=tol,maxIter=100,out=1)
		x3 = fgmres(A,rhs,5,tol=tol,maxIter=100,x=randn(size(rhs)),flexible = true)
		x4 = fgmres(A,rhs,5,tol=tol,maxIter=100,M=M2,flexible = true)

		@test norm(A*x1[1]-rhs)/norm(rhs) < tol
		@test norm(A*x3[1]-rhs)/norm(rhs) < tol
		@test norm(A*x4[1]-rhs)/norm(rhs) < tol
		@test norm(x3[1]-x1[1])/norm(x1[1]) < 1e-5
	end

	@testset "complex matrix" begin
		A  = sprandn(100,100,.1) + SparseMatrixCSC(10.0I, 100, 100) + 1im*(sprandn(100,100,.1) + SparseMatrixCSC(10.0I, 100, 100) )
		D  = Vector(diag(A))
		M3 = x -> x./D
		rhs = complex(randn(100))
		tol = 1e-6;

		# test behaviour for zero rhs

		xtt = fgmres(A,0.0*rhs,5,tol=tol,maxIter=10,out=2)
		# @test xtt[2]==-9
		# @test all(xtt[1].==0)
		# @test length(xtt[1])==100
		# @test eltype(xtt[1])==eltype(rhs)

		x1 = fgmres(A,rhs ,5,tol=tol,maxIter=100)
		x3 = fgmres(A,rhs,5,tol=tol,maxIter=100,x=randn(size(rhs)))
		x4 = fgmres(A,rhs,5,tol=tol,maxIter=100,M=M3)

		@test norm(A*x1[1]-rhs)/norm(rhs) < tol
		@test norm(A*x3[1]-rhs)/norm(rhs) < tol
		@test norm(x3[1]-x1[1])/norm(x1[1]) < 1e-5
	end
end


### A Pluto.jl notebook ###
# v0.12.18

using Markdown
using InteractiveUtils

# ╔═╡ d80342ae-522b-11eb-19c0-cf931fb4ae81
using Statistics

# ╔═╡ 9c2ef32a-522d-11eb-29e6-29da29e4dec7
using Plots

# ╔═╡ 88b47ee6-522d-11eb-1a98-ed4dd43dfd11
Base.@kwdef mutable struct SampleAvg
    t::Int = 0
    avg::Float64 = 0.0
end

# ╔═╡ 8c94a810-522d-11eb-2024-01ca19bc1984
function (s::SampleAvg)(x)
    s.t += 1
    s.avg += (x - s.avg) / s.t
    s.avg
end

# ╔═╡ 8fde026e-522d-11eb-2aa1-f3e28c238ea6
function run_once(b)
    rms = Float64[]
    distribution = randn(b)
    expectation = mean(distribution)
    sample_avg = SampleAvg()

    for i in 1:2*b
        avg = sample_avg(distribution[rand(1:b)])
        push!(rms, abs(avg - expectation))
    end
    rms
end

# ╔═╡ 969b48aa-522d-11eb-2039-61916ddf87d4
begin
	n_runs = 1000
	p = plot(legend=:topright)

	for b in [2, 10, 100, 1000]
		rms = mean(run_once(b) for _ in 1:n_runs)
		xs = (1:2*b) ./ b
		plot!(p, xs, rms, label="b=$b")
	end

	p
end

# ╔═╡ Cell order:
# ╠═d80342ae-522b-11eb-19c0-cf931fb4ae81
# ╠═88b47ee6-522d-11eb-1a98-ed4dd43dfd11
# ╠═8c94a810-522d-11eb-2024-01ca19bc1984
# ╠═8fde026e-522d-11eb-2aa1-f3e28c238ea6
# ╠═9c2ef32a-522d-11eb-29e6-29da29e4dec7
# ╠═969b48aa-522d-11eb-2039-61916ddf87d4


# # Surface-grid operations

#md # ```@meta
#md # CurrentModule = ImmersedLayers
#md # ```

#=
Here, we will discuss the various surface-grid operators available in the package.
We will start by generating the cache, just as we did in [Immersed layer caches](@ref)
=#

using ImmersedLayers
using Plots
using LinearAlgebra

#=
### Set up the grid, shape, and cache
We do this just as we did in [Immersed layer caches](@ref)
=#
Δx = 0.01
Lx = 4.0
xlim = (-Lx/2,Lx/2)
ylim = (-Lx/2,Lx/2)
g = PhysicalGrid(xlim,ylim,Δx)
RadC = 1.0
Δs = 1.4*cellsize(g)
body = Circle(RadC,Δs)
cache = SurfaceScalarCache(body,g,scaling=GridScaling)

#=
## Basic regularization and interpolation
Let's perform an example in which we regularize the x coordinate of
the surface points onto the grid. We can get the surface coordinates
by using the [`points`](@ref) function. This can be applied either to
`body` directly or to the `cache`. As for any `VectorData` type, the components
of `pts` are `pts.u` and `pts.v`.
=#
pts = points(cache);

#=
Now, we set up some blank grid data on which to regularize onto
=#
gx = zeros_grid(cache);

#=
Now regularize. We will time it to show that it is fast and memory-efficient:
=#
regularize!(gx,pts.u,cache) #hide
@time regularize!(gx,pts.u,cache);

#=
Let's plot this to look at it. This also gives us a chance to highlight
the plot recipe for grid data associated with the cache, which is achieved by
simply supplying the cache to the `plot` function in `Plots.jl`. By default,
this plots the immersed points, as well, but this can be suppressed by
adding the keyword `layers=false`.
=#
plot(gx,cache)

#=
This shows how the regularization spreads the data over a couple of cells
around the surface. In the parlance of potential theory, this is a **single layer**.
=#
#=
If we wish to interpolate, then we do so with the [`interpolate!`]
function. For example, suppose we have a uniform field that we wish to interpolate:
=#
oc = 2.5*ones_grid(cache)

#=
Now set up some surface data to receive the interpolated data, and
interpolate:
=#
f = zeros_surface(cache);
interpolate!(f,oc,cache)

#=
It is clear that the interpolation preserves the value of the
field. This is also true for linearly varying fields, since the DDF
is built to ensure this. Let's try this, using the $x$ coordinate of the grid. Here, we use the `coordinates`
function of `CartesianGrids.jl`, which gets the coordinates of the grid,
and set the values of grid data to the $x$ coordinate. We interpolate and plot,
comparing to the actual $x$ coordinate of the points on the body:
=#
#x, y = coordinates(oc,g)
#xg = similar(oc)
#xg .= x
xg = x_grid(cache)
interpolate!(f,xg,cache)
plot(f,ylim=(-2,2),label="Interpolated from grid",ylabel="x",xlabel="Index")
plot!(pts.u,label="Actual body coordinate")

#=
## A double layer
Now we will generate a double layer. Mathematically, this takes the form

$$D_s f = \nabla\cdot \left( \delta(\chi) \mathbf{n} f \right)$$

for some scalar data $f$ on the surface. (See [Background](@ref) for an example.)
Notice that it maps scalar data on the surface ($f$) to scalar data in space. So to
calculate this using our discrete tools, we set up some grid data to receive the
result. Then, we use the function [`surface_divergence!`](@ref) to compute the
double layer. Here, we will demonstrate this on the $y$ coordinate of the
surface points:
=#
dl = zeros_grid(cache)
surface_divergence!(dl,pts.v,cache)
plot(dl,cache)

#=
If the surface data are vectors, $\mathbf{f}$, then this operation is a little
different:

$$D_s \mathbf{f} = \nabla\cdot \delta(\chi) \left(  \mathbf{n} \mathbf{f} + \mathbf{f} \mathbf{n} \right)$$

This maps $\mathbf{f}$ to a vector field. We use this in conjunction with a cache
generated with [`SurfaceVectorCache`](@ref).
=#

#=
The transpose of the double layer, $D_s$, is the operation

$$G_s u = \mathbf{n}\cdot \delta^{T}(\chi) \nabla u$$

for some scalar field data $u$. This operation computes the gradient of the field
data, interpolates this gradient onto the surface, and obtains the normal
component of that surface data. As such, it represents an important tool
for computing the normal derivative of field data on a surface. In the
package, we use [`surface_grad!`](@ref) for this operation.

The vector field version of this is

$$G_s \mathbf{u} = \mathbf{n}\cdot \delta^{T}(\chi) (\nabla \mathbf{u} + \nabla^{T} \mathbf{u})$$

which maps vector field data $\mathbf{u}$ to vector-valued surface data.
=#

#=
## A curl layer
We also sometimes need to take the curl of the regularized surface data,

$$C_s f = \nabla\times \left( \delta(\chi) \mathbf{n} f \right)$$

For this, we use the [`surface_curl!`](@ref) operator. Let's demonstrate this
on a uniform field on the surface.
=#
gc = zeros_gridcurl(cache)
f = ones_surface(cache)
surface_curl!(gc,f,cache)
plot(gc,cache)

#=
The continuous version of this operation is actually zero. It's not quite
zero in discrete form. However, its norm is much smaller than that of
the double layer.
=#
norm(gc,cache)/norm(dl,cache)

#=
Finally, a pair of operations that are used in support of the
previous ones, or occasionally on their own, are

$$R_n f = \delta(\chi)\mathbf{n}\circ f$$

for scalar surface data $f$, which maps to a vector field, effectively a
field of doublet strengths; and its transpose

$$R_n^T \mathbf{u} = \mathbf{n}\cdot \delta^{T}(\chi)\mathbf{u}$$

which maps vector field data $\mathbf{u}$ to a scalar surface field, the normal
component of the vector field on the surface. These are
provided by [`regularize_normal!`](@ref) and [`normal_interpolate!`](@ref),
respectively.
=#

#=
## Masks
Masks are grid data that take the value 1 in one region (e.g., the interior of a surface)
and 0 in the other (e.g., the exterior). The functions [`mask`](@ref)
and [`complementary_mask`](@ref) achieve this
=#
m = mask(cache)
cm = complementary_mask(cache)
plot(
    surface(m,cache,layers=false),
    surface(cm,cache,layers=false)
    )

#=
One can apply a mask to some grid data by multiplying it, using, e.g.,
the `product!` function in `CartesianGrids.jl`. Let's demonstrate that
with the grid data of $x$ coordinates:
=#
xmask = zeros_grid(cache)
xcmask = zeros_grid(cache)
product!(xmask,xg,m)
product!(xcmask,xg,cm)
plot(
  plot(xmask,cache),
  plot(xcmask,cache)
  )

#=
The mask and complementary mask effectively partition the field into two parts.
We can also apply masks in place, using [`mask!`](@ref) and [`complementary_mask!`](@ref):
=#
xmask .= xg
mask!(xmask,cache)
plot(xmask,cache)

#md # ## Surface-grid operator functions
#md # ```@docs
#md # regularize!
#md # interpolate!
#md # regularize_normal!
#md # normal_interpolate!
#md # regularize_normal_cross!
#md # normal_cross_interpolate!
#md # surface_divergence!
#md # surface_grad!
#md # surface_curl!
#md # surface_divergence_cross!
#md # surface_grad_cross!
#md # surface_curl_cross!
#md # mask!
#md # mask
#md # complementary_mask!
#md # complementary_mask
#md # ```


@model function bernoulli_mixture(x)
    # Mixture prior.
    w ~ Dirichlet(2, 1.0)

    # Latent probability.
    p ~ DiscreteNonParametric([0.3, 0.7], w)

    # Observation.
    x ~ Bernoulli(p)
end

bernoulli_example(x = false) = bernoulli_mixture(x)

@model function gmm(x, K)
    N = length(x)

    # Cluster centers.
    μ ~ filldist(Normal(), K)

    # Cluster association prior.
    w ~ Dirichlet(K, 1.0)

    # Cluster assignments.
    z ~ filldist(DiscreteNonParametric(1:K, w), N)

    # Observations.
    for n = 1:N
        x[n] ~ Normal(μ[z[n]], 1.0)
    end
end

gmm_example(x = [0.1, -0.05, 1.0], K = 2) = gmm(x, K)

@model function gmm_loopy(x, K, ::Type{T}=Float64) where {T<:Real}
    N = length(x)

    # Cluster centers.
    μ = Vector{T}(undef, K)
    for k = 1:K
        μ[k] ~ Normal()
    end

    # Cluster association prior.
    w ~ Dirichlet(K, 1.0)

    # Cluster assignments & observations.
    z = Vector{Int}(undef, N)
    for n = 1:N
        z[n] ~ DiscreteNonParametric(1:K, w)
        x[n] ~ Normal(μ[z[n]], 1.0)
    end
end

# same as gmm_loopy, but with an affine transformation on μ.
@model function gmm_shifted(x, K, ::Type{T}=Float64) where {T<:Real}
    N = length(x)

    # Cluster centers.
    μ = Vector{T}(undef, K)
    for k = 1:K
        μ[k] ~ Normal()
    end

    # Cluster association prior.
    w ~ Dirichlet(K, 1.0)

    # Cluster assignments & observations.
    z = Vector{Int}(undef, N)
    for n = 1:N
        z[n] ~ DiscreteNonParametric(1:K, w)
        x[n] ~ Normal(4μ[z[n]] - 1, 1.0)
    end
end

# K clusters, each one around i for i = 1:K with variance 0.5
@model function hmm(x, K, ::Type{T}=Float64) where {T<:Real}
    N = length(x)

    # State sequence.
    s = zeros(Int, N)

    # Emission matrix.
    m = Vector{T}(undef, K)

    # Transition matrix.
    T = Vector{Vector{T}}(undef, K)

    # Assign distributions to each element
    # of the transition matrix and the
    # emission matrix.
    for i = 1:K
        T[i] ~ Dirichlet(K, 1.0)
        m[i] ~ Normal(i, 0.5)
    end

    # Observe each point of the input.
    s[1] ~ Categorical(K)
    x[1] ~ Normal(m[s[1]], 0.1)

    for i = 2:N
        s[i] ~ Categorical(T[s[i-1]])
        x[i] ~ Normal(m[s[i]], 0.1)
    end
end

hmm_example(x = [1.1, 0.95, 2.2], K = 2) = hmm(x, K)

@model function changepoint(y)
    N = length(y)
    α = 1 / mean(y)
    λ₁ ~ Exponential(α)
    λ₂ ~ Exponential(α)
    τ ~ DiscreteUniform(1, N)

    for n in 1:N
        y[n] ~ Poisson(τ > n ? λ₁ : λ₂)
    end
end

# @model function reverse_deps(x)
#     m = Vector{Float64}(undef, 2)
#     m[1] ~ Normal()
#     m[2] ~ Normal()
#     x ~ MvNormal(m)
# end

###########################################################################
# data from R. Neal paper:
const data_neal = [-1.48, -1.40, -1.16, -1.08, -1.02, 0.14, 0.51, 0.53, 0.78]
const α_neal = 10.0

function stickbreak(v)
    K = length(v) + 1
    cumprod_one_minus_v = cumprod(1 .- v)

    return map(1:K) do k
        if k == 1
            v[1]
        elseif k == K
            cumprod_one_minus_v[K - 1]
        else
            v[k] * cumprod_one_minus_v[k - 1]
        end
    end
end

# from https://luiarthur.github.io/TuringBnpBenchmarks/dpsbgmm
@model function imm_stick(y, α, K)
    N = length(y)
    crm = DirichletProcess(α)

    v ~ filldist(StickBreakingProcess(crm), K - 1)
    w = stickbreak(v)

    # Cluster assignments
    z = zeros(Int, N)
    for n = 1:N
        z[n] ~ Categorical(w)
    end

    # Cluster centers
    L = identity(K)
    μ ~ filldist(Normal(), L)

    # Observations
    for n = 1:N
        y[n] ~ Normal(μ[z[n]], 1.0)
    end
end

imm_stick_example(y = data_neal, α = α_neal, K = 10) = imm_stick(y, α, K)


using ChebyshevQuantum
using Test

tol_val = 1e-10

@testset "test particle in box " begin

    H = ChebyshevQuantum.SolveEig.setup_ham(x->0.0, 40);
    #    vals, vects =  ChebyshevQuantum.SolveEig.solve(H)
    vals, vects = eigen( H[2:end-1,2:end-1]);

    @test isapprox(vals[1],   pi^2/8, atol=tol_val)
    @test isapprox(vals[2], 2^2*pi^2/8, atol=tol_val)
    @test isapprox(vals[3], 3^2*pi^2/8, atol=tol_val)
    @test isapprox(vals[4], 4^2*pi^2/8, atol=tol_val)
    @test isapprox(vals[5], 5^2*pi^2/8, atol=tol_val)

end

@testset "test particle in bigger box " begin

    H = ChebyshevQuantum.SolveEig.setup_ham(x->0.0, 40, a = 0, b = 10);
#    vals, vects =  ChebyshevQuantum.SolveEig.solve(H)
    vals, vects = eigen( H[2:end-1,2:end-1]);

    f = (2/10)^2

    @test isapprox(vals[1],   pi^2/8 * f, atol=tol_val)
    @test isapprox(vals[2], 2^2*pi^2/8 * f, atol=tol_val)
    @test isapprox(vals[3], 3^2*pi^2/8 * f, atol=tol_val)
    @test isapprox(vals[4], 4^2*pi^2/8 * f, atol=tol_val)
    @test isapprox(vals[5], 5^2*pi^2/8 * f, atol=tol_val)

end

tol_val = 1e-10

@testset "test particle in box N=100 " begin

    H = ChebyshevQuantum.SolveEig.setup_ham(x->0.0, 100);
#    vals, vects =  ChebyshevQuantum.SolveEig.solve(H)
    vals, vects = eigen( H[2:end-1,2:end-1]);

    @test isapprox(vals[1],   pi^2/8, atol=tol_val)
    @test isapprox(vals[2], 2^2*pi^2/8, atol=tol_val)
    @test isapprox(vals[3], 3^2*pi^2/8, atol=tol_val)
    @test isapprox(vals[4], 4^2*pi^2/8, atol=tol_val)
    @test isapprox(vals[5], 5^2*pi^2/8, atol=tol_val)

end

tol_val = 1e-9

@testset "test particle in finite box " begin

    #test edge detection

    function VV(x)
        if abs(x) >= 1.0
            return 10.0
        else
            return 0.0
        end
    end

#    H, S = ChebyshevQuantum.SolveEig.total_ham_piecewise(VV, 30, [-2.0, -1.0, 1.0, 2.0], d2=-0.5)
    #    H, S, lhs = ChebyshevQuantum.Op.setup(A2=-0.5, A0=VV,a=-2.0,b=2.0, N=30)

    vals, vects = ChebyshevQuantum.SolveEig.solve(V=VV, a=-2.0, b = 2.0)

    #    vals  =  eigvals(H, S)
#    vals = real.(vals)

    @test isapprox(vals[1],  0.819848145939759, atol=tol_val) #reference values from chebfun.m (see below)
    @test isapprox(vals[2],  3.222122379353262, atol=tol_val)

end

"""
matlab code using chebfun
tic
x = chebfun('x', [-2 2]);
V = 10.0*(abs(x)>=1);
L = chebop(-2,2);
L.op = @(x,u) -0.5*diff(u,2) + V*u;
L.bc = 0;
neigs = 12;
[EV,D] = eigs(L,neigs);
disp(diag(D))
toc
"""

nothing


using ExpFamily
using Base.Test

srand(123)

C  = [1.0 0.5; 0.5 1.0]
P  = 2/3 * [2 -1; -1 2]
m  = [1.0;-0.5]

gNP = GaussianNatParam(mean=m, cov=C)
gMP = GaussianMeanParam(mean=m, cov=C)

@test isapprox(gNP.theta1, P*m)
@test isapprox(gNP.theta2, -P)
@test isapprox(gMP.mu1, m)
@test isapprox(gMP.mu2, (m*m'+C)/2)

@test isvalid(gNP)
@test isvalid(gMP)

mpFromNP  = meanparam(gNP)
mpFromNP2 = meanparam(gNP, 2.0)
npFromMP  = natparam(gMP)

@test isapprox(mpFromNP.mu1, gMP.mu1)
@test isapprox(mpFromNP.mu2, gMP.mu2)
@test isapprox(npFromMP.theta1, gNP.theta1)
@test isapprox(npFromMP.theta2, gNP.theta2)
@test isapprox(cov(mpFromNP2), 2cov(mpFromNP))

@test isapprox(mean(gNP), m)
@test isapprox(mean(gMP), m)
@test isapprox(cov(gNP), C)
@test isapprox(cov(gMP), C)
@test isapprox(var(gMP), diag(C))
@test isapprox(var(gNP), diag(C))
@test isapprox(std(gMP), sqrt.(diag(C)))
@test isapprox(std(gNP), sqrt.(diag(C)))

@test_throws DomainError GaussianNatParam(mean=m,cov=-C,check=true)
@test_throws DomainError GaussianMeanParam(mean=m,cov=-C,check=true)

dim = 4
gNPones = ones(GaussianNatParam, dim)
gMPones = ones(GaussianMeanParam, dim)

@test isapprox(mean(gNPones), zeros(dim))
@test isapprox(cov(gNPones),  eye(dim))
@test isapprox(mean(gMPones), zeros(dim))
@test isapprox(cov(gMPones),  eye(dim))

@test length(gNPones) == dim
@test length(gMPones) == dim

rNP = rand(gNPones,5000)
rMP = rand(gMPones,5000)

@test isapprox(mean(rNP,2), zeros(dim), atol=1e-1)
@test isapprox(mean(rMP,2), zeros(dim), atol=1e-1)
@test isapprox(cov(rNP,2),  eye(dim),   atol=1e-1)
@test isapprox(cov(rMP,2),  eye(dim),   atol=1e-1)

gNP2 = gNP+gNP
gMP2 = gMP+gMP

@test isapprox(gNP2.theta1, 2gNP.theta1)
@test isapprox(gNP2.theta2.data, 2gNP.theta2.data)
@test isapprox(gMP2.mu1, 2gMP.mu1)
@test isapprox(gMP2.mu2.data, 2gMP.mu2.data)

m1  = randn(dim)
m2  = randn(dim)
C1  = randn(dim,dim)
C1 *= C1'
C2  = randn(dim,dim)
C2 *= C2'

gNPa = GaussianNatParam(mean=m1, cov=C1)
gNPb = GaussianNatParam(mean=m2, cov=C2)
gMPa = GaussianMeanParam(mean=m1, cov=C1)
gMPb = GaussianMeanParam(mean=m2, cov=C2)

gNPtest = gNPa + gNPb - natparam(gMPb)
gMPtest = gMPa + gMPb - meanparam(gNPb)

@test isapprox(gNPtest.theta1, gNPa.theta1)
@test isapprox(gNPtest.theta2, gNPa.theta2)
@test isapprox(gMPtest.mu1, gMPa.mu1)
@test isapprox(gMPtest.mu2, gMPa.mu2)

@test isapprox(gMPtest, gMPa)

@test isapprox(GaussianNatParam(vec(gNPa),dim),gNPa)
@test isapprox(GaussianMeanParam(vec(gMPa),dim),gMPa)

@test isapprox(2gNPa + gNPa*2 - 3gNPa, gNPa)
@test isapprox(2gMPa + gMPa*2 - 3gMPa, gMPa)
@test isapprox(gNPa/2 + gNPa/2, gNPa)
@test isapprox(gMPa/2 + gMPa/2, gMPa)

gNP = GaussianNatParam(mean=m, cov=C)
@test isapprox(project(gNP), gNP)

nP = -diagm([1.5, 1e7])
m  = [1.0;-0.5]
lm = 1e-6

gNPproj = project(GaussianNatParam(-nP*m, nP), minvar=lm)
gMPproj = project(GaussianMeanParam(mean=m, cov=inv(-nP)), minvar=lm)

nPthresh = -diagm([1.5, 1./lm])

@test isapprox(gNPproj.theta1, -nPthresh*m)
@test isapprox(gNPproj.theta2, -diagm([1.5, 1./lm]))
@test isapprox(gMPproj.mu1, m)
@test isapprox(cov(gMPproj), diagm([1./1.5, lm]))


using Base.Test
using SymEngine

x = symbols("x")
y = symbols(:y)
@vars z

# Check Basic conversions
@test eltype([Basic(u) for u in [1, 1/2, 1//2, pi, e]]) == Basic

# make sure @vars defines in a local scope
let
    @vars w
end
@test_throws UndefVarError isdefined(w)

a = x^2 + x/2 - x*y*5
b = diff(a, x)
@test b == 2*x + 1//2 - 5*y

c = x + Rational(1, 5)
c = expand(c * 5)
@test c == 5*x + 1

c = x ^ 5
@test diff(c, x) == 5 * x ^ 4

c = x ^ y
@test c != y^x

c = Basic(-5)
@test abs(c) == 5
@test abs(c) != 4

show(a)
println()
show(b)
println()

## mathfuns
@test abs(Basic(-1)) == 1
@test sin(Basic(1)) == subs(sin(x), x, 1)
@test sin(PI) == 0
@test subs(sin(x), x, pi) == 0
@test sind(Basic(30)) == 1 // 2

## calculus
x,y = symbols("x y")
n = Basic(2)
ex = sin(x*y)
@test diff(log(x),x) == 1/x
@test diff(ex, x) == y * cos(x*y)
@test diff(ex, x, 2) == diff(diff(ex,x), x)
@test diff(ex, x, n) == diff(diff(ex,x), x)
@test diff(ex, x, y) == diff(diff(ex,x), y)
@test diff(ex, x, y,x) == diff(diff(diff(ex,x), y), x)
@test series(sin(x), x, 0, 2) == x
@test series(sin(x), x, 0, 3) == x - x^3/6

## ntheory
@test mod(Basic(10), Basic(4)) == 2
@test mod(Basic(10), 4) == 2               # mod(::Basic, ::Number)
@test_throws MethodError mod(10, Basic(4)) # no mod(::Number, ::Basic)
@test gcd(Basic(10), Basic(4)) == 2
@test lcm(Basic(10), Basic(4)) == 20
@test binomial(Basic(5), 2) == 10

## type information
a = Basic(1)
b = Basic(1//2)
@test isa(SymEngine.BasicType(a+a), SymEngine.BasicType{Val{:Integer}})
@test isa(SymEngine.BasicType(a+b), SymEngine.BasicType{Val{:Rational}})

## can we do math with items of BasicType?
a1 = SymEngine.BasicType(a)
tot = a1
for i in 1:100  tot = tot + a1 end
@test tot == 101
sin(a1)

## subs
ex = x^2 + y^2
@test subs(ex, x, 1) == 1 + y^2
@test subs(ex, (x, 1)) == 1 + y^2
@test subs(ex, x => 1) == 1 + y^2
@test subs(ex, (x,1), (y,2)) == 1 + 2^2
@test subs(ex, x => 1, y => 2) == 1 + 2^2

## lambidfy
@test_approx_eq lambdify(sin(Basic(1))) sin(1)
ex = sin(x)
@test_approx_eq lambdify(ex)(1) sin(1)
ex = exp(PI/2*x)
@test_approx_eq lambdify(ex)(1) exp(pi/2)

## N
a = Basic(1)
@test N(a) == 1
@test N(Basic(1//2)) == 1//2
@test N(Basic(12345678901234567890)) == 12345678901234567890

## generic linear algebra
x = symbols("x")
A = [x 2; x 1]
@test det(A) == -x
@test det(inv(A)) == - 1/x
(A \ [1,2])[1] == 3/x

## check that unique work (hash)
x,y,z = symbols("x y z")
@test length(SymEngine.free_symbols([x*y, y,z])) == 3


#=
Find nodes for triangular elements.
Use the book.
=#
include("triangle_quadrature_table.jl");

# Build global nodes for a triangular element with refel and vertices v
function triangle_element_nodes(refel, v)
    return  triangle_refel_to_xy(refel.r[:,1], refel.r[:,2], v);
end

# Set up refel nodal array
function triangle_refel_nodes!(refel)
    (eqx, eqy) = triangle_equilateral_nodes(refel.N);
    (r, s) = triangle_equilateral_to_rs(eqx,eqy);

    refel.r = zeros(refel.Np,2);
    refel.wr = zeros(refel.Np);
    #refel.g = zeros(refel.Np,2);
    #refel.wg = zeros(refel.Np);

    refel.r[:,1] = r;
    refel.r[:,2] = s;

    # quadrature nodes/weights from a table
    xyw = triangle_quadrature_nodes_weights(refel.N+1);
    refel.g = xyw[:,1:2];
    refel.wg = xyw[:,3];
    tol = 1e-8;
    tf1(x) = abs(x[1] + 1) < tol;
    tf2(x) = abs(x[2] + 1) < tol;
    tf3(x) = abs(x[1] + x[2]) < tol;
    refel.face2local = [get_face2local_map(refel.r, tf1),
                        get_face2local_map(refel.r, tf2),
                        get_face2local_map(refel.r, tf3)];

    ### Unfinished: find correct surface quadrature weights.
    ### For now just take average(correct for order=1).
    refel.surf_r = [refel.r[refel.face2local[1],:], refel.r[refel.face2local[2],:], refel.r[refel.face2local[3],:]];
    refel.surf_g = [refel.r[refel.face2local[1],:], refel.r[refel.face2local[2],:], refel.r[refel.face2local[3],:]];

    tmp = [length(refel.face2local[1]), length(refel.face2local[2]), length(refel.face2local[3])];
    refel.surf_wr = [ones(tmp[1])./tmp[1], ones(tmp[2])./tmp[2], ones(tmp[3])./tmp[3]];
    refel.surf_wg = refel.surf_wr;
end

# Purpose  : Compute (x,y) nodes in equilateral triangle for polynomial of order N
function triangle_equilateral_nodes(N)
    # optimal alpha values for N up to 16 (from book)
    alpopt = [0.0000 0.0000 1.4152 0.1001 0.2751 0.9800 1.0999 1.2832 1.3648 1.4773 1.4959 1.5743 1.5770 1.6223 1.6258];

    # Set optimized parameter, alpha, depending on order N
    if (N<16)
        alpha = alpopt[N];
    else
        alpha = 5/3;
    end

    # total number of nodes
    Np = Int64((N+1)*(N+2)/2);

    # Create equidistributed nodes on equilateral triangle
    L1 = zeros(Np);
    L2 = zeros(Np);
    L3 = zeros(Np);
    sk = 1;
    for n=1:(N+1)
        for m=1:(N+2-n)
            L1[sk] = (n-1)/N;
            L3[sk] = (m-1)/N;
            sk = sk+1;
        end
    end
    L2 = 1 .- L1 .- L3;
    x = -L2+L3;
    y = (-L2 - L3 + 2 .* L1) .* 0.5773502691896258; # 1/sqrt(3)=0.5773502691896258

    # Compute blending function at each node for each edge
    blend1 = 4 .* L2 .* L3;
    blend2 = 4 .* L1 .* L3;
    blend3 = 4 .* L1 .* L2;

    # Amount of warp for each node, for each edge
    warpf1 = triangle_warpfactor(N,L3-L2);
    warpf2 = triangle_warpfactor(N,L1-L3);
    warpf3 = triangle_warpfactor(N,L2-L1);

    # Combine blend & warp
    warp1 = blend1 .* warpf1 .* (1 .+ (alpha*L1).^2);
    warp2 = blend2 .* warpf2 .* (1 .+ (alpha*L2).^2);
    warp3 = blend3 .* warpf3 .* (1 .+ (alpha*L3).^2);

    # Accumulate deformations associated with each edge
    # x = x + 1*warp1 + cos(2*pi/3)*warp2 + cos(4*pi/3)*warp3;
    # y = y + 0*warp1 + sin(2*pi/3)*warp2 + sin(4*pi/3)*warp3;
    x = x + warp1 - 0.5*warp2 - 0.5*warp3;
    y = y + 0.8660254037844387*warp2 - 0.8660254037844387*warp3;

    return (x, y);
end

# Compute scaled warp function at order N based on rout interpolation nodes
function triangle_warpfactor(N, rout)
    # Compute LGL and equidistant node distribution
    #(LGLr, w) = jacobi_gauss_quad(0,0,N);
    (LGLr, w) = jacobi_LGL_quad(N);

    h = 2/N;
    req = -1:h:1;

    # Compute Vandermonde based on req
    Veq = zeros(N+1, N+1);
    for i=1:N+1
        Veq[:,i] = jacobi_polynomial(req, 0, 0, i-1);
    end

    # Evaluate Lagrange polynomial at rout
    Nr = length(rout);
    Pmat = zeros(N+1,Nr);
    for i=1:N+1
      Pmat[i,:] = jacobi_polynomial(rout, 0, 0, i-1);
    end
    Lmat = Veq'\Pmat;

    # Compute warp factor
    warp = Lmat'*(LGLr - req);

    # Apply Scale factor
    for i=1:length(rout)
        if abs(rout[i]) < 0.9999999999
            warp[i] = warp[i] / (1-rout[i]*rout[i]);
        else
            warp[i] = 0;
        end
    end

    return warp;
end

# From (x,y) in equilateral triangle to (r,s) coordinates in standard triangle
function triangle_equilateral_to_rs(x,y)
    # sqrt(3)=1.7320508075688772
    L1 = (1.7320508075688772 .* y .+ 1.0) ./ 3.0;
    L2 = (-3.0 .* x .- 1.7320508075688772 .* y .+ 2.0) ./ 6.0;
    L3 = ( 3.0 .* x .- 1.7320508075688772 .* y .+ 2.0) ./ 6.0;

    r = -L2 + L3 - L1;
    s = -L2 - L3 + L1;

    return (r, s);
end

# # From (r,s) coordinates in reference triangle to (x,y) in triangle with vertices v
# # v is a 2x3 array [x1 x2 x3; y1 y2 y3]
# function triangle_refel_to_xy(r, s, v)
#     x = 0.5 .* (-(r+s) .* v[1,1] .+ (1+r) .* v[1,2] .+ (1+s) .* v[1,3]);
#     y = 0.5 .* (-(r+s) .* v[2,1] .+ (1+r) .* v[2,2] .+ (1+s) .* v[2,3]);

#     return (x, y);
# end

# function get_face2local_map(r, compare)
#     n = size(r,1);
#     map = zeros(Int,n);
#     nf = 0;
#     for i=1:n
#         if compare(r[i,:])
#             nf = nf+1;
#             map[nf] = i;
#         end
#     end

#     return map[1:nf];
# end


n = 10 # rand(1:10)
@test matrixdepot("lehmer", Float64, n) == matrixdepot("lehmer", n)

A = ones(n, 1) * [1:n;]'
A = A./A'
A = tril(A) + tril(A)' - Matrix(1.0I, n, n)

@test matrixdepot("lehmer", n) ≈ A
println("'lehmer' passed test...")


using JuliaWeBWorK
using  Roots

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

meta=(AuthorText="JuliaWeBWorK",
      Keywords="numeric derivatives",
      Question="1"
      )

qs = JuliaWeBWorK.QUESTIONS()
letters = JuliaWeBWorK.letters()

intro = raw"""

**Derivatives**

"""

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

numericq(jmt"""

### Problem $(letters())

Calculate the slope of the secant line of \( f(x)= {{:a1}} x^2 + {{:a2}} \)
between \( ({{:a3}},f({{:a3}})) \) and \( ({{:a4}},f({{:a4}})) \).
""",
         (a₂,a₀,a,b) -> begin
         fn = x  ->  a₂*x^2 + a₀
         m  = (fn(a)-fn(b))/(a-b)
         m
         end,
         (2:4, 6:8, 2:4, 5:9)
         ) |> qs

## ----

numericq(raw"""

###  Problem $(letters())

For the function \( f(x)= {{:a1}} x^2 + {{:a2}} \) between
\( ({{:a3}},f({{:a3}})) \) and \( ({{:a4}},f({{:a4}})) \),
plot the function and the secant line.

Estimate from the graph the largest distance between the two
functions between {{:a3}} and {{:a4}}.
""",
         (a₂,a₀,a,b) ->
         fn =  x  ->  a₂*x^2 + a₀
         ## could work harde rhere
         m  = (fn(a)-fn(b))/(a-b)
         M = maximum(abs, [fn(x) - (a +  m*(x-a)) for  x ∈ range(min(a,b),  max(a,b), length=100)])
         M
         end,
         (2:4, 6:8, 2:4, 5:9),
         tolerance=5.0
         ) |> qs

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

p  =  Page(intro, qs; meta...)


using Distributions
dists = [
    Uniform(10,20),
    Exponential(3.5),
    Gamma(0.5,7),
    Beta(10,0.5),
    Weibull(10,0.5),
    Normal(20,3.5),
    Rayleigh(2.4),
    Cauchy(20,3.5)]

println("Distribution \t\t\t Parameters \t Support")
reshape([dists ;  params.(dists) ;
		((d)->(minimum(d),maximum(d))).(dists) ],
		length(dists),3)


# Text provided under a Creative Commons Attribution license, CC-BY.  All code is made available under the FSF-approved BSD-3 license.  (c) Lorena A. Barba, Gilbert F. Forsyth 2017. Thanks to NSF for support via CAREER award #1149784.
# [@LorenaABarba](https://twitter.com/LorenaABarba)
# Translation to Julia by [@miguelraz_](https://twitter.com/miguelraz_)

# 12 steps to Navier–Stokes
# =====
# ***

# Did you experiment in Steps [1](./01_Step_1.jl) and [2](./02_Step_2.jl) using different parameter choices?
# If you did, you probably ran into some unexpected behavior. Did your solution ever blow up?
# (In my experience, CFD students *love* to make things blow up.)

# You are probably wondering why changing the discretization parameters affects your solution in such a drastic way.
# This notebook complements our [interactive CFD lessons](https://github.com/barbagroup/CFDPython) by discussing the CFL condition.
# And learn more by watching Prof. Barba's YouTube lectures (links below).

## Convergence and the CFL Condition
# ----
# ***

# For the first few steps, we've been using the same general initial and boundary conditions.
# With the parameters we initially suggested, the grid has 41 points and the timestep is 0.25 seconds.
# Now, we're going to experiment with increasing the size of our grid.
# The code below is identical to the code we used in [Step 1](./01_Step_1.jl),
# but here it has been bundled up in a function so that we can easily
# examine what happens as we adjust just one variable: **the grid size**.

using Plots

function  linearconv(nₓ)
    Δx = 2 / (nₓ - 1)
    nₜ = 20    #nₜ is the number of timesteps we want to calculate
    Δt = .025  #Δt is the amount of time each timestep covers (delta t)
    c = 1

    u = ones(nₓ)      #defining a numpy array which is nx elements long with every value equal to 1.
    u[Int.((.5/dx):(1 / dx))] .= 2.0  #setting u = 2 between 0.5 and 1 as per our I.C.s

    uₙ = ones(nₓ) #initializing our placeholder array, un, to hold the values we calculate for the n+1 timestep

    for n in nₜ #iterate through time
        uₙ = copy(n) ##copy the existing values of u into un
        for i in 1:nₓ
            u[i] = uₙ[i] - c * dt / dx * (uₙ[i] - uₙ[i-1])
        end
    end
    u
end

u = linearconv(nₓ)
plot(range(start = 0, stop = 2, nx), u);

# Now let's examine the results of our linear convection problem with an increasingly fine mesh.

linearconv(41) #convection using 41 grid points

# This is the same result as our Step 1 calculation, reproduced here for reference.

linearconv(61)

# Here, there is still numerical diffusion present, but it is less severe.

linearconv(71)

# Here the same pattern is present -- the wave is more square than in the previous runs.

linearconv(85)

# This doesn't look anything like our original hat function.

### What happened?

# To answer that question, we have to think a little bit about what we're actually implementing in code.

# In each iteration of our time loop, we use the existing data about our wave to estimate the speed of the
# wave in the subsequent time step.  Initially, the increase in the number of grid points returned more accurate answers.
# There was less numerical diffusion and the square wave looked much more like a square wave than it did in our first example.

# Each iteration of our time loop covers a time-step of length $\Delta t$, which we have been defining as 0.025.

# During this iteration, we evaluate the speed of the wave at each of the $x$ points we've created.  In the last plot, something has clearly gone wrong.

# What has happened is that over the time period $\Delta t$, the wave is travelling a distance which is greater than `Δx`.
# The length `Δx` of each grid box is related to the number of total points `nₓ`, so stability can
# be enforced if the $\Delta t$ step size is calculated with respect to the size of `Δx`.

# $$\sigma = \frac{u \Delta t}{\Delta x} \leq \sigma_{\max}$$

# where $u$ is the speed of the wave; $\sigma$ is called the **Courant number** and the
# value of $\sigma_{\max}$ that will ensure stability depends on the discretization used.

# In a new version of our code, we'll use the CFL number to calculate the appropriate time-step `Δt` depending on the size of `Δx`.

function linearconv(nₓ)
    Δx = 2 / (nₓ - 1)
    nₜ = 20    #nt is the number of timesteps we want to calculate
    c = 1
    σ = .5
    indexes = Int.((.5/Δx):(1/Δx))

    Δt = σ * Δx

    u = ones(nx)
    u[indexes] .= 2

    un = ones(nx)

    for n in range(nₜ):  #iterate through time
        uₙ = copy(u) ##copy the existing values of u into un
        for i in 1:nₓ
            u[i] = uₙ[i] - c * dt / dx * (uₙ[i] - uₙ[i-1])
        end
    end
    u
end

plot(range(start = 0, end = 2, length = nₓ), u)

linearconv(41)

linearconv(61)

linearconv(81)

linearconv(101)

linearconv(121)

# Notice that as the number of points `nx` increases, the wave convects a shorter and shorter distance.
# The number of time iterations we have advanced the solution at is held constant at `nₜ = 20`,
# but depending on the value of `nₓ` and the corresponding values of `Δx` and `Δt`, a shorter time window is being examined overall.

## Learn More
# -----
# ***

# It's possible to do rigurous analysis of the stability of numerical schemes, in some cases. Watch Prof. Barba's presentation of this topic in **Video Lecture 9** on You Tube.

# TODO
# from IPython.display import YouTubeVideo
# YouTubeVideo('Yw1YPBupZxU')


using DiffEqCallbacks, OrdinaryDiffEq, Test

# Non-negative ODE examples
#
# Reference:
# Shampine, L.F., S. Thompson, J.A. Kierzenka, and G.D. Byrne,
# "Non-negative solutions of ODEs," Applied Mathematics and Computation Vol. 170, 2005,
# pp. 556-569.
# https://www.mathworks.com/help/matlab/math/nonnegative-ode-solution.html

"""
Absolute value function

```math
\\frac{du}{dt} = -|u|
```
with initial condition ``u₀=1``, and solution

```math
u(t) = u₀*e^{-t}
```
for positive initial values ``u₀``.
"""
function absval(du,u,p,t)
    du[1] = -abs(u[1])
end
analytic(u₀, p, t) = u₀*exp(-t)
ff = ODEFunction(absval,analytic=analytic)
prob_absval = ODEProblem(ff, [1.0], (0.0, 40.0))

# naive approach leads to large errors
naive_sol_absval = solve(prob_absval, BS3())
@test naive_sol_absval.errors[:l∞] > 9e4
@test naive_sol_absval.errors[:l2] > 1.3e4

# general domain approach
# can only guarantee approximately non-negative values
function g(resid,u)
    resid[1] = u[1] < 0 ? -u[1] : 0
end
general_sol_absval = solve(prob_absval, BS3(); callback=GeneralDomain(g, [1.0]))
@test_broken all(x -> x[1] ≥ -10*eps(), general_sol_absval.u)
@test general_sol_absval.errors[:l∞] < 9.9e-5
@test general_sol_absval.errors[:l2] < 4.3e-5
@test general_sol_absval.errors[:final] < 4.3e-18

# test non-autonomous function
g_t(resid, u, p, t) = g(resid, u)

general_t_sol_absval = solve(prob_absval, BS3(); callback=GeneralDomain(g_t, [1.0]))
@test general_sol_absval.t ≈ general_t_sol_absval.t &&
    general_sol_absval.u ≈ general_t_sol_absval.u

# positive domain approach
# can guarantee non-negative values
positive_sol_absval = solve(prob_absval, BS3(); callback=PositiveDomain([1.0]))
@test_broken all(x -> x[1] ≥ 0, positive_sol_absval.u)
@test_broken general_sol_absval.t ≈ positive_sol_absval.t
@test general_sol_absval.errors[:l∞] ≈ positive_sol_absval.errors[:l∞]

# specify abstol as array or scalar
positive_sol_absval2 = solve(prob_absval, BS3(); callback=PositiveDomain([1.0], abstol=[1e-6]))
@test positive_sol_absval.t ≈ positive_sol_absval2.t &&
    positive_sol_absval.u ≈ positive_sol_absval2.u
positive_sol_absval3 = solve(prob_absval, BS3(); callback=PositiveDomain([1.0], abstol=1e-6))
@test positive_sol_absval.t ≈ positive_sol_absval3.t && positive_sol_absval.u ≈ positive_sol_absval3.u

# specify scalefactor
positive_sol_absval4 = solve(prob_absval, BS3(); callback=PositiveDomain([1.0], scalefactor=0.2))
@test length(positive_sol_absval.t) < length(positive_sol_absval4.t)
@test positive_sol_absval.errors[:l2] > positive_sol_absval4.errors[:l2]

"""
Knee problem

```math
\\frac{du}{dt} = \epsilon^{-1}(1-t-u)u
```

with initial condition ``u0=1``, and generally ``0 < \epsilon << 1``.
Here ``\epsilon=1e-6``. Then the solution approaches ``u=1-t`` for ``t<1``
and ``u=0`` for ``t>1``.
"""
function knee(du,u,p,t)
    du[1] = 1e6*(1-t-u[1])*u[1]
end

prob_knee = ODEProblem(knee, [1.0], (0.0, 2.0))

# unfortunately callbacks do not work with solver CVODE_BDF which is comparable to ode15s
# used in MATLAB example, so we use Rodas5
naive_sol_knee = solve(prob_knee, Rodas5())
@test naive_sol_knee[1, end] ≈ -1.0 atol=1e-5

# positive domain approach
positive_sol_knee = solve(prob_knee, Rodas5(); callback=PositiveDomain([1.0]))
@test_broken all(x -> x[1] ≥ 0, positive_sol_knee.u)
@test positive_sol_knee[1, end] ≈ 0.0

## Now test on out-of-place equations
r, K = 1.1, 10.0
logistic(u,p,t) = u*r*(1-u/K)
t = (0.0, 20.0)
logistic_p = ODEProblem(logistic, 0.02, t)
logistic_s = solve(logistic_p, Tsit5())
logistic_s_positive = solve(logistic_p, Tsit5(), callback=PositiveDomain())


## Plotting intro

## Step 1, add the plotting package.
# at command line, if you press "]"
# If you want to do it not at the command line, you can also use
using Pkg # this adds the Package manager to your Julia session
Pkg.add("Plots") # this adds the Plots package
## Let's create some data to plot.
"""
`brownian_motion`
=================

Create a simple brownian motion trail. A discrete-time Brownian motion
is the value that increases by a random, normally distributed variable at
each time step. So

    \$ x_{t+1} = x_{t} + N(0,1) \$

Returns: A Vector{Float64} with n entries that describes
    a Brownian motion over n time steps .
"""
function brownian_motion(n::Int)
    x = 0.0 # start our Brownian motion at x = 0.0
    B = zeros(0) # this is an empty vector of Float64's (or length 0)
    for t=1:n # for n time-steps
        x = x+randn() # randn gives you a normally distributed random value.
        push!(B,x) # add x to the "end" of the vector B.
    end
    return B
end
brownian_motion(5)

## Questions to explore on your own.
# - Modify the brownian_motion function to start at a different place
# - Modify the brownian_motion function to pre-allocate a length n array

## Other ways to implement Brownian motion
A = randn(5)
cumsum(A)

## This works
cumsum(randn(5)) # this does everything in one line
## TODO, figure out how to use the cumsum! function to do this in one step.
# cumsum!
# slighlty more annoying than I expected
## There is another way of defining a function in Julia
# I usually try and be explicit and define functions with the word function
# like the brownian_motion function.
# But you can also, just give it a name, like in math.
bmotion(n) = cumsum(randn(n))
# This syntax is equivalent to
# function bmotion(n)
#   return cumsum(randn(n))
# end
# but is much shorter.

# Sometimes, we use this for very simple functions.

##
B = brownian_motion(100)

##
using Plots # this adds the plotting functions to Julia's session
##
using Plots
plot(B) # this will plot a vector.
##
plot(brownian_motion(1000))

## Let's see multiple Plots
plot(brownian_motion(1000))
plot!(brownian_motion(1000))
##
plot!(title="two Brownian trails")
##
plot!(xlabel="Time") # alternative xlabel!("Time")
ylabel!("Value") # alternative plot!(ylabel="Value")
##
# Let's get rid of the legend
plot!(legend=false)
##
annotate!(200,15,"Max value")
##
savefig("brownian-trails.pdf") # save as a pdf
##
savefig("brownian-trails.png") # save as a png
# savefig("brownian-trails.eps") # save as a eps
# this doesn't work with GR :(
## We can do this all at once too.
A = brownian_motion(10000)
B = brownian_motion(9000)
plot(A,B) # won't give you what you expect ...
##
plot([A,B], lab=["Trail 1", "Trail 2"],  # this gives it a pair... and labels
    xlabel="Time", ylabel="Value", xscale=:log10)
title!("Two Brownian Trails")
savefig("brownian-trails-log.pdf")

## Plotting scattered data
# Don't use syntax like this without a good reason.
A, B = brownian_motion(10000), brownian_motion(10000) # all in one linef
##
scatter(A,B,alpha=0.5, markersize=1,markerstrokewidth=0,) # plot A as a function of B
plot!(title="2D Brownian motion")
## Playing around with options
x = brownian_motion(50)
y = brownian_motion(50)
##
scatter(x,y,marker=:x)
plot!(x,y,markersize=0) # this makes the line connecting them
##


export DiagOp

"""
`DiagOp(domainType::Type, dim_in::Tuple, d::AbstractArray)`

`DiagOp(d::AbstractArray)`

Creates a `LinearOperator` which, when multiplied with an array `x`, returns the elementwise product `d.*x`.

```julia
julia> D = DiagOp(Float64, (2, 2,), [1. 2.; 3. 4.])
╲  ℝ^(2, 2) -> ℝ^(2, 2)

julia> D*ones(2,2)
2×2 Array{Float64,2}:
 1.0  2.0
 3.0  4.0

```

"""
struct DiagOp{N, D, C, T <: Union{AbstractArray{C,N}, Number}} <: LinearOperator
	dim_in::NTuple{N,Int}
	d::T
end

# Constructors

###standard constructor Operator{N}(DomainType::Type, DomainDim::NTuple{N,Int})
function DiagOp(DomainType::Type, DomainDim::NTuple{N,Int}, d::T) where {N, T <: AbstractArray}
	size(d) != DomainDim && error("dimension of d must coincide with DomainDim")
    C = eltype(d) <: Complex ? complex(DomainType) : DomainType
    DiagOp{N, DomainType, C, T}(DomainDim, d)
end

###standard constructor with Scalar
function DiagOp(DomainType::Type, DomainDim::NTuple{N,Int}, d::T) where {N, T <: Number}
    C = eltype(d) <: Complex ? Complex{DomainType} : DomainType
    DiagOp{N, DomainType, C, T}(DomainDim, d)
end

# other constructors
DiagOp(d::A) where {A <: AbstractArray} = DiagOp(eltype(d), size(d), d)
DiagOp(DomainDim::NTuple{N,Int}, d::A) where {N, A <: Number} = DiagOp(Float64, DomainDim, d)

# Mappings

function mul!(y::AbstractArray{C,N}, L::DiagOp{N, D, C, T}, b::AbstractArray{D,N}) where {N, D, C, T}
	y .= L.d.*b
end

function mul!(y::AbstractArray{D,N}, L::AdjointOperator{DiagOp{N, D, C, T}}, b::AbstractArray{C,N}) where {N, D, C, T}
	y .= conj.(L.A.d).*b
end

function mul!(y::AbstractArray{D,N}, L::AdjointOperator{DiagOp{N, D, C, T}}, b::AbstractArray{C,N}) where {N, D <: Real, C <: Complex, T}
    y .= real.(conj.(L.A.d).*b)
end

# Transformations (we'll see about this)
# inv(L::DiagOp) = DiagOp(L.domainType, L.dim_in, (L.d).^(-1))

# Properties

diag(L::DiagOp) = L.d
diag_AAc(L::DiagOp) = L.d.*conj.(L.d)
diag_AcA(L::DiagOp) = conj.(L.d).*L.d

domainType(L::DiagOp{N, D, C, T}) where {N, D, C, T} = D
codomainType(L::DiagOp{N, D, C, T}) where {N, D, C, T} = C

size(L::DiagOp) = (L.dim_in, L.dim_in)

fun_name(L::DiagOp) = "╲"

is_diagonal(L::DiagOp) = true

# TODO: probably the following allows for too-close-to-singular matrices
is_invertible(L::DiagOp) = all(L.d .!= 0.)
is_full_row_rank(L::DiagOp) = is_invertible(L)
is_full_column_rank(L::DiagOp) = is_invertible(L)


### A Pluto.jl notebook ###
# v0.12.16

using Markdown
using InteractiveUtils

# ╔═╡ e077ab16-22e5-11eb-15ed-5f0654bf7328
using Pkg, DrWatson

# ╔═╡ e75a9a06-22e5-11eb-2e36-8d4f62b830ed
begin
	@quickactivate "StatisticsWithJuliaPlutoNotebooks"
	using Random, Distributions, Plots
	Random.seed!(0)
end

# ╔═╡ ca5fbdc8-22e5-11eb-0a60-cb0a127a3ca5
md"## Listing5.12"

# ╔═╡ fa68607e-22e5-11eb-0558-c9a4d9f77426
begin
	alpha = 0.05
	L(obs) = obs - (1-sqrt(alpha))
	U(obs) = obs + (1-sqrt(alpha))
end

# ╔═╡ ae6c7528-3af7-11eb-14a4-4bb1f292614b
begin
	mu = 5.57
	observation = rand(TriangularDist(mu-1,mu+1,mu))
	Text("Lower bound L: $(L(observation)))")
end

# ╔═╡ ae6cb47a-3af7-11eb-0264-bb161b7731a1
Text("Upper bound U: $(U(observation)))")

# ╔═╡ 475ff888-22e6-11eb-2354-09f8f40a8e12
md"## End of listing5.12"

# ╔═╡ Cell order:
# ╟─ca5fbdc8-22e5-11eb-0a60-cb0a127a3ca5
# ╠═e077ab16-22e5-11eb-15ed-5f0654bf7328
# ╠═e75a9a06-22e5-11eb-2e36-8d4f62b830ed
# ╠═fa68607e-22e5-11eb-0558-c9a4d9f77426
# ╠═ae6c7528-3af7-11eb-14a4-4bb1f292614b
# ╠═ae6cb47a-3af7-11eb-0264-bb161b7731a1
# ╟─475ff888-22e6-11eb-2354-09f8f40a8e12