##############################################################################
##
## Solve hamilton-jacobi-bellman
##
##############################################################################

##############################
##
## FDSpec
##
##############################
"""
Use finite difference method to solve for HJB equation.

### INPUTS
- `fd`    : Finite Difference scheme
- `sol`   : holds solution information
- `twoap` : household problem description
"""
function solve_hjb!(twoap::TwoAssetsProb, fd::FDSpec, sol::SolutionExp;
                    maxit::Int = 400, crit::Float64 = 1e-8, monotonicity = true, verbose::Bool = true)

    bn, an, zn = length(fd.b), length(fd.a), length(fd.z)

    #== Initialize Vold ==#
    Vold = zeros(bn,an,zn)

    #== Construct functions ==#
    compute_pol!(solut, V_old) = comp_pol!(twoap, fd, solut, V_old)
    update_V!(solut, V_old)  = updateV!(twoap, fd, solut, V_old, test_mon = monotonicity)

    for it in 1:maxit
        copy!(Vold, vcat(sol.V...))
        # vcat is important in the FDExp case where fd.V is vector of vectors

        #== Compute the Optimal Policy ==#
        compute_pol!(sol, Vold)

        #== Update vⁿ ==#
        update_V!(sol, Vold)

        # check convergence
        distance = chebyshev(vcat(sol.V...), vec(Vold))
        # distance = norm(vec(fd.V)-vec(Vold), Inf)
        if distance < crit
            println("hjb solved : $(it) iterations")
            # println(out_text,"hjb solved : $(it) iterations")
            break
        else
            # update V using newV
            verbose && (it % 10 == 0) && @printf("  value function iteration %d, distance %.4f \n", it,distance)
            # verbose && (it % 10 == 0) && @printf(out_text,"  value function iteration %d, distance %.4f \n", it,distance)
        end
    end
end

function optconsump(twoap::TwoAssetsProb, ∂Vb::Float64, liqinc::Float64, ℓutil::Float64)

    invγ   = 1.0/twoap.γ

    # ∂Vb>=0 ? c = ∂Vb^(-invγ) : error("negative value not allowed")
    # REVIEW:110 trick to not get stuck in case of non-monotonicity
    ∂Vb>=0 ? c = ∂Vb^(-invγ) + ℓutil : c = liqinc
    sc = liqinc - c
    dV = utilfn(twoap, c, ℓutil) + sc*∂Vb

    return c, sc, dV
end

function comp_pol!(twoap::TwoAssetsProb, fd::FDSpec, sol::SolutionExp, V::Array{Float64,3})

    #== Construct fns ==#
    dχinv1(∂ratio, a) = dχinv(twoap, ∂ratio, a)
    χ1(d,a) = χ(twoap, d, a)
    utilfn1(c, ℓutil) = utilfn(twoap, c, ℓutil)

    #== Check if function is changing ==#
    # println(dχinv1(1.1, 1.0))
    # println(χ1(1e-5, 100.0))

    #== Unpack from FD ==#
    z = fd.z; a = fd.a ; b = fd.b;
    Δbgrid = fd.Δbgrid
    Δagrid = fd.Δagrid
    netbinc = fd.netbinc

    ℓutilgrid = fd.ℓutilgrid

    # precompute the lengths
    bn = length(b); an = length(a); zn = length(z)

    invdb = 1.0./Δbgrid
    invda = 1.0./Δagrid

    for zi in 1:zn, ai in 1:an, bi in 1:bn

        # == Derivative w.r.t B == #
        bi<bn && ( ∂Vbᶠ = ( V[bi+1, ai, zi] - V[bi, ai, zi] ) * invdb[bi] )
        bi>1  && ( ∂Vbᴮ = ( V[bi, ai, zi] - V[bi-1, ai, zi] ) * invdb[bi-1] )

        # == Derivative w.r.t A == #
        ai<an && ( ∂Vaᶠ = ( V[bi, ai+1, zi] - V[bi, ai, zi] ) * invda[ai] )
        ai>1  && ( ∂Vaᴮ = ( V[bi, ai, zi] - V[bi, ai-1, zi] ) * invda[ai-1] )

        ### CONSUMPTION decision ###
        liqinc = netbinc[bi,zi]
        ℓutil  = ℓutilgrid[zi]

        # == Foward == #
        if bi<bn
            cᶠ, scᶠ, Hcᶠ  = optconsump(twoap, ∂Vbᶠ, liqinc, ℓutil)
        else
            scᶠ, dVᶠ = 0.0, -1.0e12
        end
        scᶠ>0.0 ? (validᶠ = 1 ) : (validᶠ = 0 )

        # == Backward == #
        if bi>1
            cᴮ, scᴮ, Hcᴮ  = optconsump(twoap, ∂Vbᴮ, liqinc, ℓutil)
        else
            scᴮ, Hcᴮ = 0.0, -1.0e12
        end
        scᴮ<0.0 ? (validᴮ = 1) : (validᴮ = 0)

        #== Catch Other cases ==#
        c⁰  = liqinc
        sc⁰ = 0.0
        Hc⁰ = utilfn1(c⁰, ℓutil)

        if validᶠ==1   && ( validᴮ==0 || Hcᶠ>=Hcᴮ ) && ( Hcᶠ>=Hc⁰ )
            sol.c[bi, ai, zi]  = cᶠ
            sol.sc[bi, ai, zi] = scᶠ
        elseif validᴮ==1 && ( validᶠ==0 || Hcᴮ>=Hcᶠ ) && ( Hcᴮ>=Hc⁰ )
            sol.c[bi, ai, zi]  = cᴮ
            sol.sc[bi, ai, zi] = scᴮ
        else
            sol.c[bi, ai, zi]  = c⁰
            sol.sc[bi, ai, zi] = sc⁰
        end

        #== DEPOSIT decision ==#
        #== b backward, a forward ==#
        if bi>1 && ai<an
            dᴮᶠ = dχinv1(∂Vaᶠ/∂Vbᴮ, a[ai])
            # change in utility
            Hdᴮᶠ = ∂Vaᶠ * dᴮᶠ - ∂Vbᴮ*( dᴮᶠ + χ1(dᴮᶠ,a[ai]) )

            (dᴮᶠ>0.0 && Hdᴮᶠ>=0.0) ? ( validᴮᶠ=1 ) : ( validᴮᶠ=0 )
        else
            validᴮᶠ = 0; Hdᴮᶠ = -1.0e-12
        end

        #== b forward, a backward ==#
        if bi < bn && ai > 1
            dᶠᴮ = dχinv1(∂Vaᴮ/∂Vbᶠ, a[ai])
            # change in utility
            Hdᶠᴮ = ∂Vaᴮ * dᶠᴮ - ∂Vbᶠ*( dᶠᴮ + χ1(dᶠᴮ,a[ai]) )

            (dᶠᴮ<-χ1(dᶠᴮ,a[ai]) && Hdᶠᴮ>=0.0) ? ( validᶠᴮ=1 ) : ( validᶠᴮ=0 )
        else
            validᶠᴮ = 0; Hdᶠᴮ = -1.0e-12
        end

        #== b backward, a backward ==#
        if bi>1 && ai>1
            dᴮᴮ = dχinv1(∂Vaᴮ/∂Vbᴮ, a[ai])
            # change in utility
            Hdᴮᴮ = ∂Vaᴮ * dᴮᴮ - ∂Vbᴮ*( dᴮᴮ + χ1(dᴮᴮ,a[ai]) )

            (dᴮᴮ>-χ1(dᴮᴮ, a[ai]) && dᴮᴮ<=0 && Hdᴮᴮ>=0.0) ? ( validᴮᴮ=1 ) : ( validᴮᴮ=0 )

            # REVIEW:80 BOUNDARY Adjustment 01
            #        force the use of ∂Vbᴮ in case bi==bn
            # bi == bn && (dᴮᴮ<=0 && Hdᴮᴮ>=0.0) && (validᴮᴮ = 1)
        else
            validᴮᴮ = 0; Hdᴮᴮ = -1.0e-12
        end

        #== Check and assign deposit ==#
        if    validᴮᶠ==1                 && (validᶠᴮ==0 || Hdᴮᶠ>=Hdᶠᴮ)   && (validᴮᴮ==0 || Hdᴮᶠ>=Hdᴮᴮ)

            sol.d[bi, ai, zi] = dᴮᶠ
        # ------------------------------------------------------------------------------------------ #
        elseif (validᴮᶠ==0 || Hdᶠᴮ>=Hdᴮᶠ) &&  validᶠᴮ==1                  && (validᴮᴮ==0 || Hdᶠᴮ>=Hdᴮᴮ)

            sol.d[bi, ai, zi] = dᶠᴮ
        # ------------------------------------------------------------------------------------------ #
        elseif (validᴮᶠ==0 || Hdᴮᴮ>=Hdᴮᶠ) && (validᶠᴮ==0 || Hdᴮᴮ>=Hdᶠᴮ)   &&  validᴮᴮ==1

            sol.d[bi, ai, zi] = dᴮᴮ
        # ------------------------------------------------------------------------------------------ #
        elseif validᴮᶠ==0 && validᶠᴮ==0 && validᴮᴮ==0

            sol.d[bi, ai, zi] = 0.0
        end

        # u[bi, ai, zi] = utilfn1( c[bi, ai, zi] )
        # bdot[bi, ai, zi] = sc[bi, ai, zi] - d[bi, ai, zi] - χ1(d, a[ai])
    end

    Void
end

function updateV!(twoap::TwoAssetsProb, fd::FDExp, sol::SolutionExp, V::Array{Float64,3}; test_mon = true)

    #== Construct function ==#
    χ1(d,a) = χ(twoap, d, a)
    utilfn1(c, ℓutil) = utilfn(twoap, c, ℓutil)
    #== Idiosyncratic component ==#
    λdiag = fd.λdiag
    λtoff = fd.λoff'

    #== HOUSEHOLD Parameters ==#
    γ, ρ, _, σ, ψ = _unpackparams(twoap)

    #== Optimal policies ==#
    c = sol.c; sc = sol.sc; d = sol.d

    #== Extract GRIDS ==#
    b = fd.b; a = fd.a; z = fd.z
    Δbgrid, Δagrid = fd.Δbgrid, fd.Δagrid
    netainc = fd.netainc;
    ℓutilgrid = fd.ℓutilgrid

    invΔ = fd.invΔ

    #== SOLUTION Matrices to fill ==#
    Vnew = sol.V
    A    = sol.A
    Au   = sol.Au

    #== Storage matrices ==#
    B  = fd.B
    b̃  = fd.b̃

    # precompute the lengths
    bn = length(b); an = length(a); zn = length(z); abn = bn*an

    invdb = 1.0./Δbgrid
    invda = 1.0./Δagrid
    inv1γ = 1.0/(1.0-γ)

    for zi in 1:zn
        #== set values to zero ==#
        fill!(nonzeros(A[zi]), zero(Float64))
        fill!(nonzeros(Au[zi]), zero(Float64))

        for abi in 1:abn
        # NOTE:90 Positions of A that I need to fill up
            # A[zi][abi-bn, abi] - v(i  ,j-1,k) --> - invda*[ d⁻ ]                  (SHOULD BE POSITIVE)
            # A[zi][abi-1 , abi] - v(i+1,j  ,k) --> - invdb*[ (scᴮ)⁻ + (sdᴮ)⁻ ]     (SHOULD BE POSITIVE)
            # A[zi][abi   , abi] - v(i  ,j  ,k) --> + invdb*[ (scᴮ)⁻ + (sdᴮ)⁻ - ( (scᶠ)⁺ - (sdᶠ)⁺ ) ]
            #                                       + invda*[ d⁻ - (d⁺ + ξwz + rᵃa) ]
            # A[zi][abi+1 , abi] - v(i-1,j  ,k) --> + invdb*[ (scᶠ)⁺ + (sdᶠ)⁺ ]     (SHOULD BE POSITIVE)
            # A[zi][abi+bn, abi] - v(i  .j+1,k) --> + invda*[ d⁺ + ξwx + rᵃa ]      (SHOULD BE POSITIVE)

            ai = afromab(abi)
            bi = bfromab(abi)
            ℓutil  = ℓutilgrid[zi]

            #== RHS of the Bellman equation==#
            b̃[zi][abi] = utilfn1(c[bi, ai, zi], ℓutil) + V[bi, ai, zi] * invΔ + dot(λtoff[:,zi] ,V[bi,ai,:][:] )

            ## COMPUTE THE ENTRIES ##
            # ==================================================================
            #== bdrift ==#
            bdriftᶠ = max( sc[bi, ai, zi], 0.0 ) + max( -d[bi, ai, zi] -χ1(d[bi, ai, zi], a[ai]) , 0.0 )
            bdriftᴮ = min( sc[bi, ai, zi], 0.0 ) + min( -d[bi, ai, zi] -χ1(d[bi, ai, zi], a[ai]) , 0.0 )

            bi <  bn && ( budriftᶠ = max( sc[bi, ai, zi] - d[bi, ai, zi] - χ1(d[bi, ai, zi], a[ai]) , 0.0 ) )
            bi <  bn && ( budriftᴮ = min( sc[bi, ai, zi] - d[bi, ai, zi] - χ1(d[bi, ai, zi], a[ai]) , 0.0 ) )
            bi == bn && ( budriftᶠ = max( sc[bi, ai, zi] - d[bi-1, ai, zi] - χ1(d[bi-1, ai, zi], a[ai]) , 0.0 ) )
            bi == bn && ( budriftᴮ = min( sc[bi, ai, zi] - d[bi-1, ai, zi] - χ1(d[bi-1, ai, zi], a[ai]) , 0.0 ) )

            #== adrift ==#
            adriftᶠ = max( d[bi, ai, zi], 0.0 ) + netainc[ai,zi]
            adriftᴮ = min( d[bi, ai, zi], 0.0 )

            bi <  bn && ( audriftᶠ = max( d[bi, ai, zi] + netainc[ai,zi], 0.0 ) )
            bi <  bn && ( audriftᴮ = min( d[bi, ai, zi] + netainc[ai,zi], 0.0 ) )
            bi == bn && ( audriftᶠ = max( d[bi-1, ai, zi] + netainc[ai,zi], 0.0 ) )
            bi == bn && ( audriftᴮ = min( d[bi-1, ai, zi] + netainc[ai,zi], 0.0 ) )

            #== BOUNDARY Adjustment 02 ==#
            # bi <  bn && ( bdriftᴮ = min( sc[bi, ai, zi], 0.0 ) + min( -d[bi, ai, zi] -χ1(d[bi, ai, zi], a[ai]) , 0.0 ) )
            # bi == bn && ( bdriftᴮ = min( sc[bi, ai, zi], 0.0 ) -d[bi, ai, zi] -χ1(d[bi, ai, zi], a[ai]) )
            #== BOUNDARY Adjustment 03 ==#
            # ai <  an && ( adriftᴮ = min( d[bi, ai, zi], 0.0 ) )
            # ai == an && ( adriftᴮ = d[bi, ai, zi] + netainc[ai,zi] )

            # v[i, j- 1, k]
            if adriftᴮ != 0.0 && ai>1
                row = abfromab(ai-1,bi)
                A[zi][row, abi] = -adriftᴮ * invda[ai-1]
            end

            # v[i, j- 1, k] KF
            if audriftᴮ != 0.0 && ai>1
                row = abfromab(ai-1,bi)
                Au[zi][row, abi] = -audriftᴮ * invda[ai-1]
            end

            # v[i-1, j, k]
            if bdriftᴮ != 0.0 && bi>1
                row = abfromab(ai,bi-1)
                A[zi][row, abi] = -bdriftᴮ * invdb[bi-1]
            end

            # v[i-1, j, k] KF
            if budriftᴮ != 0.0 && bi>1
                row = abfromab(ai,bi-1)
                Au[zi][row, abi] = -budriftᴮ * invdb[bi-1]
            end

            # v[i, j, k]
            val = 0.0
            ( bdriftᴮ != 0.0 && bi > 1 ) && ( val += bdriftᴮ * invdb[bi-1] )
            ( bdriftᶠ != 0.0 && bi < bn) && ( val -= bdriftᶠ * invdb[bi]   )
            ( adriftᴮ != 0.0 && ai > 1)  && ( val += adriftᴮ * invda[ai-1] )
            ( adriftᶠ != 0.0 && ai < an) && ( val -= adriftᶠ * invda[ai]   )
            val != 0.0 && ( A[zi][abi, abi] = val )

            # v[i, j, k] KF
            val = 0.0
            ( budriftᴮ != 0.0 && bi > 1 ) && ( val += budriftᴮ * invdb[bi-1] )
            ( budriftᶠ != 0.0 && bi < bn) && ( val -= budriftᶠ * invdb[bi]   )
            ( audriftᴮ != 0.0 && ai > 1)  && ( val += audriftᴮ * invda[ai-1] )
            ( audriftᶠ != 0.0 && ai < an) && ( val -= audriftᶠ * invda[ai]   )
            val != 0.0 && ( Au[zi][abi, abi] = val )

            # v[i-1, j, k]
            if bdriftᶠ != 0.0 && bi<bn
                row = abfromab(ai,bi+1)
                A[zi][row, abi] = bdriftᶠ * invdb[bi]
            end

            # v[i-1, j, k] KF
            if budriftᶠ != 0.0 && bi<bn
                row = abfromab(ai,bi+1)
                Au[zi][row, abi] = budriftᶠ * invdb[bi]
            end

            # v[i, j+1, k]
            if adriftᶠ != 0.0 && ai<an
                row = abfromab(ai+1,bi)
                A[zi][row, abi] = adriftᶠ * invda[ai]
            end

            # v[i, j+1, k]
            if audriftᶠ != 0.0 && ai<an
                row = abfromab(ai+1,bi)
                Au[zi][row, abi] = audriftᶠ * invda[ai]
            end

        end
    end

    ## Set B = diag(invΔ + ρ) - A
    # TODO: think of a better way to do this part
    @inbounds for zi in 1:zn
        Avals = nonzeros(A[zi])
        Bvals = nonzeros(B[zi])
        #== clean the sotorage matrix before assigning values ==#
        fill!(Bvals, zero(Float64))
        Brows = rowvals(B[zi])
        for abi in 1:abn
            for k in nzrange(B[zi], abi)
                # loop over elements in the column ij
                row = Brows[k]
                Bvals[k] = - Avals[k]
                row == abi && (Bvals[k] += invΔ + ρ - λdiag[zi])
                #NOTE:120 `λdiag` is entering only in this ste
            end
        end
        Vnew[zi][:] = B[zi]' \ b̃[zi]
    end

    #== Check for monotonicity ==#
    if test_mon
        for zi in 1:zn, abi in 1:abn
            ai = afromab(abi)
            bi = bfromab(abi)
            #= Check mon in the b dimension =#
            if bi>1
                Vnew[zi][abi]<Vnew[zi][abfromab(ai,bi-1)] &&
                @printf("   non-mon on b dimension at = (%d,%d)  \n", bi,ai)
            end

            #= Check mon in the a dimension =#
            if ai>1
                Vnew[zi][abi]<Vnew[zi][abfromab(ai-1,bi)] &&
                @printf("   non-mon on a dimension at = (%d,%d)  \n", bi,ai)
            end
        end
    end
    Void
end

"""
Update value function based on Fortran code
"""
function updateV!(fd::FDImp, sol::SolutionImp, twoap::TwoAssetsProb, V::Array{Float64,3}; test_mon = true)

    #== HOUSEHOLD Parameters ==#
    γ, ρ, ξ, _, _ = _unpackparams(twoap)

    #== Optimal policies ==#
    c = sol.c; sc = sol.sc; d = sol.d

    #== Extract grids and matrices ==#
    b = fd.b; a = fd.a; z = fd.z
    Δbgrid, Δagrid = fd.Δbgrid, fd.Δagrid
    netainc = fd.netainc;
    ℓutilgrid = fd.ℓutilgrid
    invΔ = fd.invΔ

    #== Solution matrices to fill ==#
    V  = sol.V
    A  = sol.A
    Au = sol.Au

    #== Storage matrices to fill ==#
    b̃ = fd.b̃
    Λ = fd.Λ
    B = fd.B

    # precompute the lengths
    bn = length(b); an = length(a); zn = length(z); abn = bn*an

    invdb = 1.0./Δbgrid
    invda = 1.0./Δagrid
    inv1γ = 1.0/(1.0-γ)

    ijk = zero(Int)

    #== set past values to zero ==#
    fill!(nonzeros(Au), zero(Float64))
    fill!(nonzeros(A), zero(Float64))

    for zi in 1:zn, abi in 1:abn
    # NOTE:190 Positions of A that I need to fill
        # A[ijk-bn, ijk] - v(i  ,j-1,k) --> - invda*[ d⁻ ]
        # A[ijk- 1, ijk] - v(i+1,j  ,k) --> - invdb*[ (scᴮ)⁻ + (sdᴮ)⁻ ]
        # A[ijk   , ijk] - v(i  ,j  ,k) -->   invdb*[ (scᴮ)⁻ + (sdᴮ)⁻ -(scᶠ)⁺ -(sdᶠ)⁺] + invda*[ d⁻ -(d⁺ + ξwz + rᵃa) ]
        # A[ijk+ 1, ijk] - v(i-1,j  ,k) -->   invdb*[ (scᶠ)⁺ + (sdᶠ)⁺ ]
        # A[ijk+bn, ijk] - v(i  .j+1,k) -->   invda*[ d⁺ + ξwx + rᵃa ]

        ijk +=1

        ai = afromab(abi)
        bi = bfromab(abi)
        ℓutil  = ℓutilgrid[zi]

        #== RHS of the Bellman equation==#
        b̃[ijk] =  utilfn1(c[bi, ai, zi], ℓutil)  + V[bi, ai, zi] * invΔ
        # ==================================================================

        ### COMPUTE THE ENTRIES ###

        #== bdrift ==#
        bdriftᶠ = max( sc[bi, ai, zi], 0.0 ) + max( -d[bi, ai, zi] -χ1(d[bi, ai, zi], a[ai]) , 0.0 )
        bdriftᴮ = min( sc[bi, ai, zi], 0.0 ) + min( -d[bi, ai, zi] -χ1(d[bi, ai, zi], a[ai]) , 0.0 )
        #== bdrift for KF ==#
        bi <  bn && ( budriftᶠ = max( sc[bi, ai, zi] - d[bi, ai, zi] - χ1(d[bi, ai, zi], a[ai]) , 0.0 ) )
        bi <  bn && ( budriftᴮ = min( sc[bi, ai, zi] - d[bi, ai, zi] - χ1(d[bi, ai, zi], a[ai]) , 0.0 ) )
        bi == bn && ( budriftᶠ = max( sc[bi, ai, zi] - d[bi-1, ai, zi] - χ1(d[bi-1, ai, zi], a[ai]) , 0.0 ) )
        bi == bn && ( budriftᴮ = min( sc[bi, ai, zi] - d[bi-1, ai, zi] - χ1(d[bi-1, ai, zi], a[ai]) , 0.0 ) )

        #== adrift ==#
        adriftᶠ = max( d[bi, ai, zi], 0.0 ) + netainc[ai,zi]
        adriftᴮ = min( d[bi, ai, zi], 0.0 )
        # for KF
        bi <  bn && ( audriftᶠ = max( d[bi, ai, zi] + netainc[ai,zi], 0.0 ) )
        bi <  bn && ( audriftᴮ = min( d[bi, ai, zi] + netainc[ai,zi], 0.0 ) )
        bi == bn && ( audriftᶠ = max( d[bi-1, ai, zi] + netainc[ai,zi], 0.0 ) )
        bi == bn && ( audriftᴮ = min( d[bi-1, ai, zi] + netainc[ai,zi], 0.0 ) )

        #== REVIEW:100 BOUNDARY ADJUSTMENT 02 ==#
        # bi <  bn && ( bdriftᴮ = min( sc[bi, ai, zi], 0.0 ) + min( -d[bi, ai, zi] -χ1(d[bi, ai, zi], a[ai]) , 0.0 ) )
        # bi == bn && ( bdriftᴮ = min( sc[bi, ai, zi], 0.0 ) -d[bi, ai, zi] -χ1(d[bi, ai, zi], a[ai]) )
        #== REVIEW:100 BOUNDARY ADJUSTMENT 03 ==#
        # ai <  an && ( adriftᴮ = min( d[bi, ai, zi], 0.0 ) )
        # ai == an && ( adriftᴮ = d[bi, ai, zi] + netainc[ai,zi] )

        # v[i, j- 1, k]
        if adriftᴮ != 0.0 && ai>1
            row = abzfromabz(ai-1,bi,zi)
            A[row, ijk] = -adriftᴮ * invda[ai-1]
        end

        # v[i, j- 1, k] KF
        if audriftᴮ != 0.0 && ai>1
            row = abzfromabz(ai-1,bi,zi)
            Au[row, ijk] = -audriftᴮ * invda[ai-1]
        end

        # v[i-1, j, k]
        if bdriftᴮ != 0.0 && bi>1
            row = abzfromabz(ai,bi-1,zi)
            A[row,ijk] = -bdriftᴮ*invdb[bi-1]
        end

        # v[i-1, j, k] KF
        if budriftᴮ != 0.0 && bi>1
            row = abzfromabz(ai,bi-1,zi)
            Au[row,ijk] = -budriftᴮ*invdb[bi-1]
        end

        # v[i, j, k]
        val = 0.0
        ( bdriftᴮ != 0.0 && bi > 1 ) && ( val += bdriftᴮ * invdb[bi-1] )
        ( bdriftᶠ != 0.0 && bi < bn) && ( val -= bdriftᶠ * invdb[bi]   )
        ( adriftᴮ != 0.0 && ai > 1)  && ( val += adriftᴮ * invda[ai-1] )
        ( adriftᶠ != 0.0 && ai < an) && ( val -= adriftᶠ * invda[ai]   )
        val != 0.0 && ( A[ijk, ijk] = val )

        # v[i, j, k] KF
        val = 0.0
        ( budriftᴮ != 0.0 && bi > 1 ) && ( val += budriftᴮ * invdb[bi-1] )
        ( budriftᶠ != 0.0 && bi < bn) && ( val -= budriftᶠ * invdb[bi]   )
        ( audriftᴮ != 0.0 && ai > 1)  && ( val += audriftᴮ * invda[ai-1] )
        ( audriftᶠ != 0.0 && ai < an) && ( val -= audriftᶠ * invda[ai]   )
        val != 0.0 && ( Au[ijk, ijk] = val )

        # v[i-1, j, k]
        if bdriftᶠ != 0.0 && bi<bn
            row = abzfromabz(ai,bi+1,zi)
            A[row, ijk] = bdriftᶠ*invdb[bi]
        end

        # v[i-1, j, k] KF
        if budriftᶠ != 0.0 && bi<bn
            row = abzfromabz(ai,bi+1,zi)
            Au[row, ijk] = budriftᶠ*invdb[bi]
        end

        # v[i, j+1, k]
        if adriftᶠ != 0.0 && ai<an
            row = abzfromabz(ai+1,bi,zi)
            A[row, ijk] = adriftᶠ * invda[ai]
        end

        # v[i, j+1, k] KF
        if audriftᶠ != 0.0 && ai<an
            row = abzfromabz(ai+1,bi,zi)
            Au[row, ijk] = audriftᶠ * invda[ai]
        end

    end
    #== set A = A + Λ ==#
    # NOTE:70 Notice that all matrices are initialized transposed
    #        they are going to be transposed at the END
    Λvals = nonzeros(Λ)
    Avals = nonzeros(A)
    copy!(Avals, Avals + Λvals)

    #== Set B = diag(invΔ + ρ) - A ==#
    # TODO: think of a better way to do this part
    Bvals = nonzeros(B)
    fill!(Bvals, zero(Float64))
    Brows = rowvals(B)
    ijk = zero(Int)
    @inbounds for zi in 1:zn, ai in 1:an, bi in 1:bn
        ijk += 1
        for k in nzrange(B, ijk)
            # loop over elements in the column ijk
            row = Brows[k]
            Bvals[k] = - Avals[k]
            if row == ijk
                Bvals[k] += invΔ + ρ
            end
        end
    end

    # solve for new V (NOTE:280 the B transposed)
    V[:] = B' \ b̃

    return Void
end

##############################
##
## TwoAssetsFD
##
# soldated
##############################
# """
# Use finite difference method to solve for HJB equation.
# """
# function solve_hjb!(fd::TwoAssetsFD;
#                     maxit::Int = 100,
#                     crit::Float64 = 1e-5,
#                     verbose::Bool = true)
#
#     for iter in 1:maxit
#
#         ## UPDATE newV
#         updateV!(fd)
#
#         # check convergence
#         # distance = chebyshev(vec(fd.newV), vec(fd.V))
#         distance = norm(vec(fd.newV)-vec(fd.V),Inf)
#         if distance < crit
#             if verbose
#                 println("hjb solved : $(iter) iterations")
#             end
#             break
#         else
#             # update V using newV
#             fd.V = fd.newV
#             @printf("Value function iteration %d, distance %.4f \n", iter,distance)
#         end
#     end
# end
#
# """
#
# Update value function based on MATLAB code
# """
# function updateV!(fd::TwoAssetsFD)
#
#     #== Parameters ==#
#     γ = fd.twoap.γ; ρ = fd.twoap.ρ; ξ = fd.twoap.ξ; rᴬ = fd.twoap.rᴬ;
#     rᴮ⁺ = fd.twoap.rᴮ; w = fd.twoap.w; χ₀ = fd.twoap.χ₀; χ₁ = fd.twoap.χ₁
#     rᴮ⁻ = fd.twoap.rᴮ + fd.twoap.wedge
#
#     #== FD method ==#
#     invΔ = fd.invΔ;
#     z = fd.z; a = fd.a ; b = fd.b
#     Δb = fd.Δb; Δa = fd.Δa
#     u = fd.u;
#     B = fd.B; A = fd.A; Λ = fd.Λ; V = fd.V
#
#     # precompute the lengths
#     bn = length(b)
#     an = length(a)
#     zn = length(z)
#     N = an*bn*zn
#
#     invdb = 1.0/(Δb)
#     invda = 1.0/(Δa)
#     invγ  = 1/γ
#     inv1γ = 1.0/(1.0-γ)
#
#     V = reshape(V, bn, an, zn)
#
#     # set A = Λ
#     # NOTE:20 Notice that all matrices are initialized transposed
#     #        they are going to be transposed at the END
#     Λvals = nonzeros(Λ)
#     Avals = nonzeros(A)
#     copy!(Avals, Λvals)
#
#     # update A and u
#     ij = zero(Int)
#     @inbounds for zi in 1:zn, ai in 1:an, bi in 1:bn
#         ij += 1
#
#         rᴮ = (b[bi]>0)*rᴮ⁺ + ( 1-(b[bi]>0) )*rᴮ⁻ # set the interest rate
#
#         ## == Derivative w.r.t B == ##
#         bi<bn ? ( ∂Vbᶠ = ( V[bi+1, ai, zi] - V[bi, ai, zi] ) * invdb ) : ( ∂Vbᶠ = ( ( 1-ξ ) * w * z[zi] + rᴮ⁺ * b[end] )^(-γ) )
#         bi>1  ? ( ∂Vbᴮ = ( V[bi, ai, zi] - V[bi-1, ai, zi] ) * invdb ) : ( ∂Vbᴮ = ( ( 1-ξ ) * w * z[zi] + rᴮ⁻ * b[1]   )^(-γ) )
#
#         ## == Derivative w.r.t A == ##
#         ai<an ? ( ∂Vaᶠ = ( V[bi, ai+1, zi] - V[bi, ai, zi] ) * invda ) : ∂Vaᶠ = zero(Float64) # make sure below that don't use ∂Vaᶠ
#         ai>1  ? ( ∂Vaᴮ = ( V[bi, ai, zi] - V[bi, ai-1, zi] ) * invda ) : ∂Vaᴮ = zero(Float64)
#
#         ## CONSUMPTION decision ##
#         cᶠ = ∂Vbᶠ^(-invγ); cᴮ = ∂Vbᴮ^(-invγ)
#         scᴮ = (1-ξ) * z[zi] * w + b[bi] * rᴮ - cᴮ
#         scᶠ = (1-ξ) * z[zi] * w + b[bi] * rᴮ - cᶠ
#
#         ## DEPOSIT decision ##
#         dᴮᴮ = foc_dep( ∂Vbᴮ, ∂Vaᴮ, a[ai] )
#         dᴮᶠ = foc_dep( ∂Vbᴮ, ∂Vaᶠ, a[ai] )
#         dᶠᴮ = foc_dep( ∂Vbᶠ, ∂Vaᴮ, a[ai] )
#         dᶠᶠ = foc_dep( ∂Vbᶠ, ∂Vaᶠ, a[ai] )
#
#     # DOUBT: having a hard time to understand these constraints...
#         #    Although the deposits are computed for bi=1,bn, later on
#         #    these are set to zero when doing upwind for b
#
#         #== eq (13),(14) + set the boundaries in a ==#
#         dᴮ = min( dᴮᴮ, 0.0 ) + max(dᴮᶠ, 0.0)     # eq (13)/ note dᴮᴮ > dᴮᶠ
#         ai==1  && ( dᴮ = (dᴮᶠ>0)*dᴮᶠ )              # Moll: make sure d>=0 at amin, don't use VaB(:,1,:)
#         ai==an && ( dᴮ = (dᴮᴮ<0)*dᴮᴮ )              # Moll: make sure d<=0 at amax, don't use VaF(:,J,:)
#
#         dᶠ = min( dᶠᴮ,0.0 ) + max(dᶠᶠ, 0.0)         # eq (14)
#         ai==1  && ( dᶠ = (dᶠᶠ>0)*dᶠᶠ )
#         ai==an && ( dᶠ = (dᶠᴮ<0)*dᶠᴮ )
#
#         sdᴮ = - dᴮ - χ(dᴮ,a[ai])
#         sdᶠ = - dᶠ - χ(dᶠ,a[ai])
#         bi==bn && ( sdᶠ = min(sdᶠ,0.0) )
#
#         # == consumption UPWIND SCHEME == #
#         c₀ = (1-ξ) * z[zi] * w + b[bi] * rᴮ
#         Icᴮ = (scᴮ<0) ; Icᶠ = (scᶠ>0 && scᴮ>=0); Ic⁰ = 1 - Icᴮ - Icᶠ
#
#         c  = Icᴮ*cᴮ + Icᶠ*cᶠ + Ic⁰*c₀
#         u[ij] = inv1γ * c^(1-γ) + invΔ * V[ij]  # notice that V[ij] still valid
#
#         # == dep UPWIND SCHEME for Vb == #
#         Idᶠ = (sdᶠ>0); Idᴮ = (sdᴮ<0 && sdᶠ<=0)
#
#         # NOTE:30 Boundary constraints of b
#         bi==1  && ( Idᴮ =0  )               # make sure don't use ∂Vbᴮ  if ib=1  for deposit decision
#         bi==bn && ((Idᴮ,Idᶠ)=(1,0))         # make sure don't use ∂Vbᶠ  if ib=bn for deposit decision
#                                             # and FORCE ∂Vbᴮ is used at bn for deposit decision
#         Id⁰ = 1 - Idᶠ - Idᴮ                 # INACTION Region
#
#         #== dep UPWINDE SCHEMEq for Va- eq (17) ==#
#         d⁻ = min(Idᴮ * dᴮᴮ + Idᶠ * dᶠᴮ,0.0)       # use backward for ∂Va
#         d⁺ = max(Idᴮ * dᴮᶠ + Idᶠ * dᶠᶠ,0.0)       # use forward for ∂Va
#         mf = d⁺ + ξ*w*z[zi] + rᴬ*a[ai]
#
#         # NOTE:40 Boundary constraints of a
#         #   Aggregate all terms in ∂Vaᴮ at an
#         ai==an && (mf = 0; d⁻ = ξ*w*z[zi] + rᴬ*a[an] + d⁻)
#
#         # NOTE:50 Positions of A that I need to fill
#             # A[ij-bn, ij] - v(i  ,j-1,k) --> - invda*[ d⁻ ]
#             # A[ij- 1, ij] - v(i+1,j  ,k) --> - invdb*[(scᴮ)⁻ + (sdᴮ)⁻ ]
#             # A[ij   , ij] - v(i  ,j  ,k) --> invdb*[ (scᴮ)⁻ - (scᶠ)⁺ + (sdᴮ)⁻ - (sdᶠ)⁺ ] + invda*[ d⁻ - (d⁺ + ξwz + rᵃa) ]
#             # A[ij+ 1, ij] - v(i-1,j  ,k) --> invdb*[ (scᶠ)⁺ + (sdᶠ)⁺ ]
#             # A[ij+bn, ij] - v(i  .j+1,k) --> invda*[ d⁺ + ξwx + rᵃa ]
#             #
#             # the matrix a is changing its entrances... Note that code from aigary only changed when
#             # it was nonzero...
#
#         current = Array(Float64,0)
#         zi == 2   && push!(current, 0.0)
#
#         ij>bn     && ( ai>1  ? push!(current, - invda * d⁻ ) : push!(current, 0.0) )
#         ij>1      && ( bi>1  ? push!(current, - invdb *( Icᴮ*scᴮ + Idᴮ*sdᴮ ) ) : push!(current, 0.0) )
#         push!(current, invdb *( Icᴮ*scᴮ + Idᴮ*sdᴮ - Icᶠ*scᶠ - Idᶠ*sdᶠ ) + invda *( d⁻ - mf ) )
#         ij<N      && ( bi<bn ? push!(current, invdb *( Icᶠ*scᶠ + Idᶠ*sdᶠ ) ) : push!(current,0.0) )
#         ij<=N-bn  && ( ai<an ? push!(current, invda * mf ) : push!(current,0.0) )
#
#         zi == 1   && push!(current, 0.0)
#
#         Avals[nzrange(A, ij)] += current
#     end
#
#     Bvals = nonzeros(B)
#     Brows = rowvals(B)
#     ij = zero(Int)
#     @inbounds for zi in 1:zn, ai in 1:an, bi in 1:bn
#         ij += 1
#         for k in nzrange(B, ij)
#             # loop over elements in the column ij
#             row = Brows[k]
#             Bvals[k] = - Avals[k]
#             if row == ij
#                 Bvals[k] += invΔ + ρ
#             end
#         end
#     end
#
#     #== Solve for new V ==#
#     #NOTE:60 the B transposed
#     fd.newV = B' \ u
#
#     return Void
# end


export Plane

"""
    Plane(𝗻̂::Vec{Dim,T}, d::T)
    Plane(P₁::Point{3}, P₂::Point{3}, P₃::Point{3})

A `Plane` in 3-dimensional space that satisfies X ⋅𝗻̂ = d, where
X is a `Dim`-dimensional point and 𝗻̂ is the unit normal to the plane.
"""
struct Plane{T}
    𝗻̂::Vec{3,T}
    d::T
end

function Plane(P₁::Point{3}, P₂::Point{3}, P₃::Point{3})
    𝗻̂ = normalize((P₂ - P₁) × (P₃ - P₁))
    return Plane(𝗻̂, coordinates(P₁) ⋅ 𝗻̂)
end


import StaticArrays.SVector
import GeometryBasics: Point
import Rotations.RotMatrix

export Point, Path, MultiPath, unitvec, point_matrix

"""
The unit vector (0, 1).

Useful for creating n-gons, e.g.:

```jldoctest
using PenPlots

expect = [
    [-1, 0],
    [0, -1],
    [1, 0],
    [0, 1],
]

isapprox(expect, [frac_rotation(i/4) * unitvec for i in 1:4])

# output

true
```
"""
const unitvec = Point(0, 1)

"""
A 1-dimensional array of `Point`s.
"""
struct Path <: AbstractVector{Point}
    points::Vector{Point}
end

function Base.size(path::Path)
    size(path.points)
end

function Base.getindex(path::Path, i::Int64)
    getindex(path.points, i)
end

"""
An array of paths. Corresponds one-to-one with a layer in a plot.
"""
const MultiPath = Vector{Path}

# Rotation by a matrix.

function Base.:*(rot::RotMatrix, path::Path)
    Path([rot * point for point in path.points])
end

function Base.:*(rot::RotMatrix, paths::MultiPath)
    MultiPath([rot * path for path in paths])
end

# Multiplication by a point.

function Base.:*(point::Point, path::Path)
    Path([point * x for x in path])
end

function Base.:*(point::Point, paths::MultiPath)
    MultiPath([point * path for path in paths])
end

# Addition with a point.

function Base.:+(point::Point, path::Path)
    Path([point + x for x in path])
end

function Base.:+(point::Point, paths::MultiPath)
    MultiPath([point + path for path in paths])
end

"""
    point_matrix(xs, ys)

Generate a matrix by producing `Point`s for the cartesian product of the vectors
of `x` and `y` values given. The resulting matrix will have as many columns as
entries of `xs`, and as many rows as values of `ys`.
"""
function point_matrix(xs, ys)
	xxs = xs' .* ones(length(ys))
	yys = ys .* ones(length(xs))'
	Point.(xxs, yys)
end


D8fields = [x^8 - 3*x^7 + 4*x^6 - 3*x^5 + 3*x^4 - 3*x^3 + 4*x^2 - 3*x + 1,
x^8 - x^7 + 2*x^6 + 3*x^5 - x^4 + 3*x^3 + 2*x^2 - x + 1,
x^8 - 2*x^7 + 2*x^6 + 2*x^5 - 2*x^4 + 2*x^3 + 2*x^2 - 2*x + 1,
x^8 - 2*x^7 + 2*x^6 - 2*x^5 + 7*x^4 - 10*x^3 + 8*x^2 - 4*x + 1,
x^8 + 6*x^4 + 1,
x^8 - 3*x^7 + 9*x^6 - 13*x^5 + 18*x^4 - 11*x^3 + 11*x^2 - 4*x + 1,
x^8 + 4*x^6 + x^4 - 6*x^2 + 4,
x^8 - 4*x^6 + 8*x^4 - 4*x^2 + 1,
x^8 + 5*x^6 + 4*x^4 + 5*x^2 + 1,
x^8 - 2*x^7 + 2*x^6 + 4*x^5 - x^4 + 4*x^3 + 2*x^2 - 2*x + 1,
x^8 + 3*x^4 + 1,
x^8 + 3*x^4 + 9,
x^8 - 4*x^7 + 8*x^6 - 10*x^5 + 11*x^4 - 10*x^3 + 2*x^2 + 2*x + 1,
x^8 - x^7 + 8*x^6 + 7*x^5 + 19*x^4 + 7*x^3 + 8*x^2 - x + 1,
x^8 - 3*x^7 + x^6 + 6*x^5 + 6*x^4 + 6*x^3 + x^2 - 3*x + 1,
x^8 - 4*x^7 + 8*x^6 - 6*x^5 + 7*x^4 - 10*x^3 + 2*x^2 - 6*x + 9,
x^8 - x^7 + 2*x^6 - 3*x^5 - x^4 + 6*x^3 + 8*x^2 + 8*x + 16,
x^8 - 5*x^6 + 28*x^4 + 15*x^2 + 9,
x^8 + 2*x^6 + 10*x^4 + 2*x^2 + 1,
x^8 - 2*x^7 + 2*x^6 - 6*x^5 + 14*x^4 - 30*x^3 + 50*x^2 - 50*x + 25]

q8fields = [ x^8 + 12*x^6 + 36*x^4 + 36*x^2 + 9,
  x^8 - 12*x^6 + 36*x^4 - 36*x^2 + 9,
  x^8 - x^7 - 34*x^6 + 29*x^5 + 361*x^4 - 305*x^3 - 1090*x^2 + 1345*x - 395,
  x^8 - 3*x^7 + 22*x^6 - 60*x^5 + 201*x^4 - 450*x^3 + 1528*x^2 - 3069*x + 4561,
  x^8 + 60*x^6 + 900*x^4 + 4500*x^2 + 5625,
  x^8 - 60*x^6 + 900*x^4 - 4500*x^2 + 5625,
  x^8 - x^7 - 47*x^6 + 40*x^5 + 581*x^4 - 220*x^3 - 2038*x^2 - 932*x - 109,
  x^8 + 84*x^6 + 1764*x^4 + 12348*x^2 + 21609,
  x^8 - 84*x^6 + 1764*x^4 - 12348*x^2 + 21609,
  x^8 + 44*x^6 + 308*x^4 + 484*x^2 + 121,
  x^8 - 44*x^6 + 308*x^4 - 484*x^2 + 121,
  x^8 + 60*x^6 + 1170*x^4 + 9000*x^2 + 22500,
  x^8 + 60*x^6 + 810*x^4 + 1800*x^2 + 900,
  x^8 - 60*x^6 + 1170*x^4 - 9000*x^2 + 22500,
  x^8 - 60*x^6 + 810*x^4 - 1800*x^2 + 900,
  x^8 - 3*x^7 - 63*x^6 + 90*x^5 + 1311*x^4 - 20*x^3 - 7702*x^2 - 5524*x - 1009,
  x^8 + 68*x^6 + 986*x^4 + 4624*x^2 + 4624,
  x^8 - 68*x^6 + 986*x^4 - 4624*x^2 + 4624,
  x^8 - 3*x^7 - 67*x^6 - 16*x^5 + 863*x^4 + 1276*x^3 + 392*x^2 - 54*x + 1,
  x^8 - x^7 + 50*x^6 + 71*x^5 + 529*x^4 + 2173*x^3 + 842*x^2 + 5545*x + 17623]


# # data file for gnash1.mod
# # Original AMPL coding by Sven Leyffer, University of Dundee
#
# # An MPEC from F. Facchinei, H. Jiang and L. Qi, A smoothing method for
# # mathematical programs with equilibrium constraints, Universita di Roma
# # Technical report, 03.96. Problem number 8
#
# param	:	c	K	b	:=
# 	1	10	5	1.2
# 	2	 8	5	1.1
# 	3	 6	5	1.0
# 	4	 4	5	0.9
# 	5	 2	5	0.8	;
#
# param L := 20;
# param g := 1.7;
#
# let x := 75;

c = [10, 8, 6, 4, 2]
K = [5, 5, 5, 5, 5]
b = [1.2, 1.1, 1.0, 0.9, 0.8]
L = 20
g = 1.7
# x = 75

# # gnash1m.mod	QQR2-MN-10-8
# # Original AMPL coding by Sven Leyffer, University of Dundee
#
# # Formulation with mixed complementarity of ...
# #
# # An MPEC from F. Facchinei, H. Jiang and L. Qi, A smoothing method for
# # mathematical programs with equilibrium constraints, Universita di Roma
# # Technical report, 03.96. Problem number 8
# #
# # Arises from Gournot Nash equilibrium , 10 instances are available
# # (see gnash1i.dat, i=0,1,...,9)
#
# # Number of variables:   10
# # Number of constraints:  8
#
# # ... parameters for each firm/company
# param c{1..5};			# c_i
# param K{1..5};			# K_i
# param b{1..5};			# \beta_i
#
# # ... parameters for each problem instance
# param L;			# L
# param g;			# \gamma
#
# # ... computed constants
# param gg := 5000^(1/g);
#
# var x >= 0, <= L;
# var y{1..4};
# var l{1..4};   		        # Multipliers
# var Q = x+y[1]+y[2]+y[3]+y[4];	# defined variable Q
#
# minimize f: c[1]*x + b[1]/(b[1]+1)*K[1]^(-1/b[1])*x^((1+b[1])/b[1])
# 		- x*( gg*Q^(-1/g) );
#
# subject to
#
#    F1: 0 = ( c[2] + K[2]^(-1/b[2])*y[1] ) - ( gg*Q^(-1/g) )
# 				- y[1]*( -1/g*gg*Q^(-1-1/g) ) - l[1];
#    F2: 0 = ( c[3] + K[3]^(-1/b[3])*y[2] ) - ( gg*Q^(-1/g) )
# 				- y[2]*( -1/g*gg*Q^(-1-1/g) ) - l[2];
#    F3: 0 = ( c[4] + K[4]^(-1/b[4])*y[3] ) - ( gg*Q^(-1/g) )
# 				- y[3]*( -1/g*gg*Q^(-1-1/g) ) - l[3];
#    F4: 0 = ( c[5] + K[5]^(-1/b[5])*y[4] ) - ( gg*Q^(-1/g) )
# 				- y[4]*( -1/g*gg*Q^(-1-1/g) ) - l[4];
#
#    g1: 0 <= y[1] <= L  complements   l[1];
#    g3: 0 <= y[2] <= L  complements   l[2];
#    g5: 0 <= y[3] <= L  complements   l[3];
#    g7: 0 <= y[4] <= L  complements   l[4];

using JuMP, Complementarity
using Ipopt
using Test

@testset "gnash1m.jl" begin

    gg = 5000^(1/g)

    gnash1m = Model(with_optimizer(Ipopt.Optimizer))

    @variable(gnash1m, 0 <= x <= L)
    @variable(gnash1m, y[1:4])
    @variable(gnash1m, l[1:4])
    @variable(gnash1m, Q >= 0)
    @constraint(gnash1m, Q == x+y[1]+y[2]+y[3]+y[4])
    @NLobjective(gnash1m, Min, c[1]*x + b[1]/(b[1]+1)*K[1]^(-1/b[1])*x^((1+b[1])/b[1])
     		             - x*( gg*Q^(-1/g) ) )

    @NLconstraint(gnash1m, cnstr[i=1:4], 0 == ( c[i+1] + K[i+1]^(-1/b[i+1])*y[i] ) - ( gg*Q^(-1/g) )
                                      - y[i]*( -1/g*gg*Q^(-1-1/g) ) - l[i] )

    for i in 1:4
        @complements(gnash1m, l[i], 0 <= y[i] <= L, smooth)
    end

    # Initial solutions to help reaching the optimality
    set_start_value.(y, [11, 19, 20, 20])
    set_start_value.(l, [0, 0, -3, -5])
    set_start_value(x, 5)
    set_start_value(Q, 70)

    JuMP.optimize!(gnash1m)

    @show JuMP.value.(y)
    @show JuMP.value.(l)
    @show JuMP.value(x)
    @show JuMP.value(Q)

    # JuMP.value.(y) = [11.3327, 19.5003, 20.0, 20.0]
    # JuMP.value.(l) = [-1.35686e-10, -9.7489e-9, -2.51077, -5.18082]
    # JuMP.value(x) = 6.369341692612341
    # JuMP.value(Q) = 77.20235592852553
    # JuMP.objective_value(gnash1m) = -6.116708234438121

    @show JuMP.objective_value(gnash1m)
    @test isapprox(JuMP.objective_value(gnash1m), -6.11671, atol=1e-4)
    @test isapprox( JuMP.value.(l)' * (L .- JuMP.value.(y)), 0, atol=1e-5)
end


# Multivariate Normal distribution

immutable MultivariateNormal{Cov<:AbstractPDMat} <: ContinuousMultivariateDistribution
    dim::Int
    zeromean::Bool
    μ::Vector{Float64}
    Σ::Cov
end

function MultivariateNormal{Cov<:AbstractPDMat}(μ::Vector{Float64}, Σ::Cov, zmean::Bool)
    d = length(μ)
    if dim(Σ) != d
        throw(ArgumentError("The dimensions of μ and Σ are inconsistent."))
    end
    MultivariateNormal{Cov}(d, zmean, μ, Σ)
end

function MultivariateNormal{Cov<:AbstractPDMat}(μ::Vector{Float64}, Σ::Cov)
    d = length(μ)
    if dim(Σ) != d
        throw(ArgumentError("The dimensions of μ and Σ are inconsistent."))
    end
    zmean::Bool = true
    for i = 1:d
        if μ[i] != 0.
            zmean = false
            break
        end
    end
    MultivariateNormal{Cov}(d, zmean, μ, Σ)
end

function MultivariateNormal{Cov<:AbstractPDMat}(Σ::Cov)
    d = dim(Σ)
    MultivariateNormal{Cov}(d, true, zeros(d), Σ)
end

MultivariateNormal(μ::Vector{Float64}, σ::Float64) = MultivariateNormal(μ, ScalMat(length(μ), abs2(σ)))
MultivariateNormal(μ::Vector{Float64}, σ::Vector{Float64}) = MultivariateNormal(μ, PDiagMat(abs2(σ)))
MultivariateNormal(μ::Vector{Float64}, Σ::Matrix{Float64}) = MultivariateNormal(μ, PDMat(Σ))

MultivariateNormal(d::Int, σ::Float64) = MultivariateNormal(ScalMat(d, abs2(σ)))
MultivariateNormal(Σ::Matrix{Float64}) = MultivariateNormal(PDMat(Σ))

const MvNormal = MultivariateNormal

# Basic statistics

dim(d::MvNormal) = d.dim

mean(d::MvNormal) = d.μ

var(d::MvNormal) = diag(d.Σ)

cov(d::MvNormal) = full(d.Σ)

logdet_cov(d::MvNormal) = logdet(d.Σ)

mode(d::MvNormal) = d.μ

modes(d::MvNormal) = [mode(d)]

entropy(d::MvNormal) = 0.5 * (dim(d) * (float64(log2π) + 1.0) + logdet_cov(d))

# PDF evaluation

_gauss_c0(g::MvNormal) = -0.5 * (dim(g) * float64(log2π) + logdet_cov(g))

function sqmahal(d::MvNormal, x::Vector{Float64})
    z::Vector{Float64} = d.zeromean ? x : x - d.μ
    invquad(d.Σ, z)
end

function sqmahal!(r::Array{Float64}, d::MvNormal, x::Matrix{Float64})
    if !(size(x, 1) == dim(d) && size(x, 2) == length(r))
        throw(ArgumentError("Inconsistent argument dimensions."))
    end
    z::Matrix{Float64} = d.zeromean ? x : bsubtract(x, d.μ, 1)
    invquad!(r, d.Σ, z)
end

sqmahal(d::MvNormal, x::Matrix{Float64}) = sqmahal!(Array(Float64, size(x, 2)), d, x)

logpdf(d::MvNormal, x::Vector{Float64}) = _gauss_c0(d) - 0.5 * sqmahal(d, x)

function logpdf!(r::Array{Float64}, d::MvNormal, x::Matrix{Float64})
    sqmahal!(r, d, x)
    c0::Float64 = _gauss_c0(d)
    for i = 1:size(x, 2)
        r[i] = c0 - 0.5 * r[i]
    end
    r
end

# Sampling

function rand!(d::MvNormal, x::Vector{Float64})
    unwhiten!(d.Σ, randn!(x))
    if !d.zeromean
        add!(x, d.μ)
    end
    x
end

function rand!(d::MvNormal, x::Matrix{Float64})
    unwhiten!(d.Σ, randn!(x))
    if !d.zeromean
        badd!(x, d.μ, 1)
    end
    x
end

# Maximum Likelihood Estimation
#
# Specialized algorithms are respectively implemented for
# each kind of covariance
#

function fit_mle(::Type{MvNormal{PDMat}}, x::Matrix{Float64})
    n = size(x, 2)
    mu = vec(mean(x, 2))
    z = bsubtract(x, mu, 1)
    C = Base.LinAlg.BLAS.gemm('N', 'T', 1.0/n, z, z)
    MvNormal(mu, PDMat(C))
end

function fit_mle(::Type{MvNormal{PDMat}}, x::Matrix{Float64}, w::Vector{Float64})
    m = size(x, 1)
    n = size(x, 2)
    if length(w) != n
        throw(ArgumentError("Inconsistent argument dimensions"))
    end

    inv_sw = 1.0 / sum(w)
    mu = Base.LinAlg.BLAS.gemv('N', inv_sw, x, w)

    z = Array(Float64, m, n)
    for j = 1:n
        cj = sqrt(w[j])
        for i = 1:m
            @inbounds z[i,j] = (x[i,j] - mu[i]) * cj
        end
    end
    C = Base.LinAlg.BLAS.gemm('N', 'T', inv_sw, z, z)

    MvNormal(mu, PDMat(C))
end

function fit_mle(::Type{MvNormal{PDiagMat}}, x::Matrix{Float64})
    m = size(x, 1)
    n = size(x, 2)

    mu = vec(mean(x, 2))
    va = zeros(Float64, m)
    for j = 1:n
        for i = 1:m
            @inbounds va[i] += abs2(x[i,j] - mu[i])
        end
    end
    multiply!(va, inv(n))

    MvNormal(mu, PDiagMat(va))
end

function fit_mle(::Type{MvNormal{PDiagMat}}, x::Matrix{Float64}, w::Vector{Float64})
    m = size(x, 1)
    n = size(x, 2)
    if length(w) != n
        throw(ArgumentError("Inconsistent argument dimensions"))
    end

    inv_sw = 1.0 / sum(w)
    mu = Base.LinAlg.BLAS.gemv('N', inv_sw, x, w)

    va = zeros(Float64, m)
    for j = 1:n
        @inbounds wj = w[j]
        for i = 1:m
            @inbounds va[i] += abs2(x[i,j] - mu[i]) * wj
        end
    end
    multiply!(va, inv_sw)

    MvNormal(mu, PDiagMat(va))
end

function fit_mle(::Type{MvNormal{ScalMat}}, x::Matrix{Float64})
    m = size(x, 1)
    n = size(x, 2)

    mu = vec(mean(x, 2))
    va = 0.
    for j = 1:n
        va_j = 0.
        for i = 1:m
            @inbounds va_j += abs2(x[i,j] - mu[i])
        end
        va += va_j
    end

    MvNormal(mu, ScalMat(m, va / (m * n)))
end

function fit_mle(::Type{MvNormal{ScalMat}}, x::Matrix{Float64}, w::Vector{Float64})
    m = size(x, 1)
    n = size(x, 2)
    if length(w) != n
        throw(ArgumentError("Inconsistent argument dimensions"))
    end

    sw = sum(w)
    inv_sw = 1.0 / sw
    mu = Base.LinAlg.BLAS.gemv('N', inv_sw, x, w)

    va = 0.
    for j = 1:n
        @inbounds wj = w[j]
        va_j = 0.
        for i = 1:m
            @inbounds va_j += abs2(x[i,j] - mu[i]) * wj
        end
        va += va_j
    end

    MvNormal(mu, ScalMat(m, va / (m * sw)))
end

fit_mle(dty::Type{MvNormal}, x::Matrix{Float64}) = fit_mle(MvNormal{PDMat}, x)
fit_mle(dty::Type{MvNormal}, x::Matrix{Float64}, w::Vector{Float64}) = fit_mle(MvNormal{PDMat}, x, w)


using Flux, Distributions, Random
using Flux: onehotbatch, Zygote
using AdvancedHMC
using LinearAlgebra
using Plots
include("turing_data.jl")

struct BayesNNet{N, P, RE}
    nnet::N
    p::P
    re::RE

    function BayesNNet(nnet;p = nothing)
        _p, re = Flux.destructure(nnet)
        if p === nothing
            p = _p
        end
        new{typeof(nnet), typeof(p), typeof(re)}(
            nnet, p, re)
    end
end

(b::BayesNNet)(x, p=b.p) = b.re(p)(x)

net_topology = Chain(Dense(2,3,tanh), Dense(3,2,tanh), Dense(2,2,σ))
bnn = BayesNNet(net_topology)

x_train = Flux.normalise(X,dims=2)
y_train = onehotbatch(Y, 0:1)

# log posterior density
alpha = 0.09
prior_sigma = sqrt(1/alpha)
function ℒ(p)
    -sum(Flux.binarycrossentropy.(bnn(x_train, p), y_train)) -
        0.5*norm(p)^2 / prior_sigma^2
end

# Set the number of samples to draw and warmup iterations
n_samples, n_adapts = 6_000, 2_000

# Define a Hamiltonian system
D = length(bnn.p)
metric = DiagEuclideanMetric(D)
hamiltonian = Hamiltonian(metric, ℒ, Zygote)

Random.seed!(1)
initial_p = randn(D)

# Define a leapfrog solver, with initial step size chosen heuristically
initial_ϵ = find_good_stepsize(hamiltonian, initial_p)
integrator = Leapfrog(initial_ϵ)

# Define an HMC sampler, with the following components
#   - multinomial sampling scheme,
#   - generalised No-U-Turn criteria, and
#   - windowed adaption for step-size and diagonal mass matrix
proposal = NUTS{MultinomialTS, GeneralisedNoUTurn}(integrator)
adaptor = StanHMCAdaptor(MassMatrixAdaptor(metric), StepSizeAdaptor(initial_ϵ, integrator))

# Run the sampler to draw samples from the specified Gaussian, where
#   - `samples` will store the samples
#   - `stats` will store diagnostic statistics for each sample
Random.seed!(123)
samples, stats = sample(hamiltonian, proposal, initial_p, n_samples, adaptor, n_adapts;
                        drop_warmup=true, progress=false)

# define test set
x_range = collect(range(-6,stop=6,length=25))
y_range = collect(range(-6,stop=6,length=25))
n_end = size(samples, 1)
anim = @gif for i=1:20:n_end
    plot_data();
    Z = [bnn([x, y], samples[i])[1] for x=x_range, y=y_range]
    contour!(x_range, y_range, Z, title="Iteration $i", clim = (0,1));
end every 30;


"""Returns the coriolis parameter on the q-grid for beta plane approximation."""
function beta_plane()

    if bc_x == "periodic"
        # points on the right edge needed too
        xx_q,yy_q = meshgrid(x_q_halo[3:end-1],y_q_halo[3:end-2])
    else
        xx_q,yy_q = meshgrid(x_q,y_q)
    end

    f₀ = 2*ω*sind(ϕ)      # Coriolis parameter [s^-1] at central latitude
    β = 2*ω/R*cosd(ϕ)     # Coriolis derivative wrt latitude [(ms)^-1]

    # for non-dimensional gradient operators f contains the grid spacing Δ
    f_q = Numtype.(Δ*(f₀ .+ β*yy_q))

    return f_q
end


#=
   multiplication by known constant
=#

# sigdigits = precision(Float64)-2
function find_consts_for_mulbyconst(c::BigFloat)
    b = round(c, sigdigits=51, base=2)
    chi = Double64(b)
    clo = Double64(c - chi)
    return chi, clo
end

function find_consts_for_divbyconst(c::BigFloat)
    return find_consts_for_mulbyconst(inv(c))
end

function mulbyconst(x::Double64, chi::Double64, clo::Double64)
    lo = clo * x
    hi = chi * x
    return hi + lo
end

divbyconst(x::Double64, chi::Double64, clo::Double64) = mulbyconst(x, chi, clo)


@compat abstract type Link end     # Link types define linkfun!, linkinv!, and mueta!

type CauchitLink <: Link end
type CloglogLink  <: Link end
type IdentityLink <: Link end
type InverseLink  <: Link end
type LogitLink <: Link end
type LogLink <: Link end
type ProbitLink <: Link end
type SqrtLink <: Link end

linkfun(::CauchitLink, μ) = tan(pi * (μ - oftype(μ, 0.5)))
linkinv(::CauchitLink, η) = oftype(η, 0.5) + atan(η) / pi
mueta(::CauchitLink, η) = one(η) / (pi * (one(η) + abs2(η)))

linkfun(::CloglogLink, μ) = log(-log1p(-μ))
linkinv(::CloglogLink, η) = -expm1(-exp(η))
mueta(::CloglogLink, η) = exp(η) * exp(-exp(η))

linkfun(::IdentityLink, μ) = μ
linkinv(::IdentityLink, η) = η
mueta(::IdentityLink, η) = 1

linkfun(::InverseLink, μ) = inv(μ)
linkinv(::InverseLink, η) = inv(η)
mueta(::InverseLink, η) = -inv(abs2(η))

linkfun(::LogitLink, μ) = logit(μ)
linkinv(::LogitLink, η) = logistic(η)
function mueta(::LogitLink, η)
    e = exp(-abs(η))
    f = one(η) + e
    return e / (f * f)
end

linkfun(::LogLink, μ) = log(μ)
linkinv(::LogLink, η) = exp(η)
mueta(::LogLink, η) = exp(η)

linkfun(::ProbitLink, μ) = -sqrt2 * erfcinv(2μ)
linkinv(::ProbitLink, η) = erfc(-η / sqrt2) / 2
mueta(::ProbitLink, η) = exp(-abs2(η) / 2) / sqrt2π

linkfun(::SqrtLink, μ) = sqrt(μ)
linkinv(::SqrtLink, η) = abs2(η)
mueta(::SqrtLink, η) = 2η

canonicallink(::Bernoulli) = LogitLink()
canonicallink(::Binomial) = LogitLink()
canonicallink(::Gamma) = InverseLink()
canonicallink(::Normal) = IdentityLink()
canonicallink(::Poisson) = LogLink()

# For the "odd" link functions we evaluate the linear predictor such that mu is closest to zero where the precision is higher
function glmvar(::Union{Bernoulli,Binomial}, link::Union{CauchitLink,InverseLink,LogitLink,ProbitLink}, μ, η)
    μ = linkinv(link, ifelse(η < 0, η, -η))
    μ * (1 - μ)
end
glmvar(::Union{Bernoulli,Binomial}, ::Link, μ, η) = μ * (1 - μ)
glmvar(::Gamma, ::Link, μ, η) = abs2(μ)
glmvar(::Normal, ::Link, μ, η) = 1
glmvar(::Poisson, ::Link, μ, η) = μ

mustart(::Bernoulli, y, wt) = (y + oftype(y, 0.5)) / 2
mustart(::Binomial, y, wt) = (wt * y + oftype(y, 0.5)) / (wt + one(y))
mustart(::Gamma, y, wt) = y == 0 ? oftype(y, 0.1) : y
mustart(::Normal, y, wt) = y
mustart(::Poisson, y, wt) = y + oftype(y, 0.1)

function devresid(::Bernoulli, y, μ)
    if y == 1
        return -2 * log(μ)
    elseif y == 0
        return -2 * log1p(-μ)
    end
    throw(ArgumentError("y should be 0 or 1 (got $y)"))
end
function devresid(::Binomial, y, μ)
    if y == 1
        return -2 * log(μ)
    elseif y == 0
        return -2 * log1p(-μ)
    else
        return 2 * (y * (log(y) - log(μ)) + (1 - y)*(log1p(-y) - log1p(-μ)))
    end
end
devresid(::Gamma, y, μ) = -2 * (log(y / μ) - (y - μ) / μ)
devresid(::Normal, y, μ) = abs2(y - μ)
devresid(::Poisson, y, μ) = 2 * (xlogy(y, y / μ) - (y - μ))

# Whether a dispersion parameter has to be estimated for a distribution
dispersion_parameter(::Union{Bernoulli, Binomial, Poisson}) = false
dispersion_parameter(::Union{Gamma, Normal, InverseGaussian}) = true

# Log-likelihood for an observation
loglik_obs(::Bernoulli, y, μ, wt, ϕ) = wt*logpdf(Bernoulli(μ), y)
loglik_obs(::Binomial, y, μ, wt, ϕ) = logpdf(Binomial(Int(wt), μ), Int(y*wt))
loglik_obs(::Gamma, y, μ, wt, ϕ) = wt*logpdf(Gamma(1/ϕ, μ*ϕ), y)
loglik_obs(::Normal, y, μ, wt, ϕ) = wt*logpdf(Normal(μ, sqrt(ϕ)), y)
loglik_obs(::Poisson, y, μ, wt, ϕ) = wt*logpdf(Poisson(μ), y)


#=
Implements an online version of Layer Normalization from https://arxiv.org/abs/1607.06450:
    Ba et al (2016): Layer Normalization

The input a is normalized ((a .- μ) ./ σ), and then re-scaled and shifted
with learnable params gain (g) and bias (b).
=#

#=
Input is a, output is y (same size as a):
    y = g ∘ z + b
    z = (a - μₜ) / σₜ
=#
type LayerNorm{T,P<:Params,W<:Weight} <: Learnable
    n::Int  # maps n --> n
    m::T  # this period's mean(aₜ)
    μ::T  # running average of m = mean(aₜ)
    σ::T  # running average of s = std(aₜ)
    σ̂::T  # the clamped version or σ
    αₜ::T
    wgt::W
    z::Vector{T}
    input::SumNode{T,1}    # a
    output::OutputNode{T,1}  # y
    params::P                  # holds g and b
end

function LayerNorm{T}(::Type{T}, n::Int,
                    θ::AbstractVector = zeros(T, 2n),
                    ∇::AbstractVector = zeros(T, 2n);
                    lookback::Int = 100,
                    α::Float64 = NaN,
                    wgt::Weight =  BoundedEqualWeight(isnan(α) ? lookback : α)
                    )
    input = InputNode(T, n)
    output = OutputNode(T, n)
    params = Params(θ, ∇, ((n,), (n,)))
    g, b = params.views
    fill!(g, one(T))
    fill!(b, zero(T))
    LayerNorm(n, zero(T), zero(T),
        one(T), one(T), zero(T),
        wgt, zeros(T,n),
        input, output, params
    )
end
LayerNorm(n::Int, args...; kw...) = LayerNorm(Float64, n, args...; kw...)

function Base.show(io::IO, t::LayerNorm)
    print(io, "LayerNorm{n=$(t.n), mu=$(t.μ), sigma=$(t.σ)}")
end

# compute output = wx + b
function transform!{T}(layer::LayerNorm{T})
    g, b = layer.params.views
    a = layer.input.val
    y = layer.output.val

    # update the αₜ weighting
    OnlineStats.updatecounter!(layer.wgt)
    layer.αₜ = weight(layer.wgt)

    # normalize the layer
    # TODO: μ/σ should be weighted avgs using BoundedEqualWeight
    layer.m = mean(a)
    layer.μ = layer.αₜ * layer.m + (one(T) - layer.αₜ) * layer.μ
    if isnan(layer.μ)
        warn("layer.μ is NaN:")
        @show layer a g b map(extrema, (a, g, b)) layer.wgt.nobs
    end
    layer.σ = layer.αₜ * std(a) + (one(T) - layer.αₜ) * layer.μ
    layer.σ̂ = if layer.σ == 0
        # @show layer.μ, layer.σ, layer.σ̂
        1e-10
    else
        layer.σ
    end

    # mult by g and add b
    @inbounds for o=1:layer.n
        layer.z[o] = (a[o] - layer.μ) / layer.σ̂
        y[o] = g[o] * layer.z[o] + b[o]
    end
    y
end

function grad!{T}(layer::LayerNorm{T})
    g, b = layer.params.views
    ∇g, ∇b = layer.params.∇_views
    a = layer.input.val
    y = layer.output.val
    ∇a = layer.input.∇
    ∇y = layer.output.∇

    D = layer.n
    Dinv = one(T) / D
    Dm1 = D - one(T)
    # @show layer.μ, layer.σ, layer.σ̂

    # TODO: ∇a, ∇g
    fill!(∇a, zero(T))
    for o=1:layer.n
        # ∇g
        ∇g[o] = ∇y[o] * layer.z[o]

        # ∇a
        base_dzₒdaₚ = Dinv - layer.αₜ * layer.z[o] * (layer.μ - layer.m) / layer.σ̂
        for p=1:layer.n
            dzₒdaₚ = base_dzₒdaₚ + layer.z[o] * layer.z[p] / Dm1
            dzₒdaₚ = (o==p ? one(T) : zero(T)) - layer.αₜ * dzₒdaₚ
            dzₒdaₚ /= layer.σ̂
            ∇a[p] += ∇y[o] * g[o] * dzₒdaₚ
        end
    end

    # ∇b
    copy!(∇b, ∇y)
    return
end


using Test
using PolynomialBasis
using ImplicitDomainQuadrature
# using Revise
using CutCell
include("useful_routines.jl")

function bulk_modulus(l, m)
    return l + 2m / 3
end

function lame_lambda(k, m)
    return k - 2m / 3
end

function analytical_coefficient_matrix(inradius, outradius, ls, ms, lc, mc)
    a = zeros(3, 3)
    a[1, 1] = inradius
    a[1, 2] = -inradius
    a[1, 3] = -1.0 / inradius
    a[2, 1] = 2 * (lc + mc)
    a[2, 2] = -2 * (ls + ms)
    a[2, 3] = 2ms / inradius^2
    a[3, 2] = 2(ls + ms)
    a[3, 3] = -2ms / outradius^2
    return a
end

function analytical_coefficient_rhs(ls, ms, theta0)
    r = zeros(3)
    Ks = bulk_modulus(ls, ms)
    r[2] = -Ks * theta0
    r[3] = Ks * theta0
    return r
end

struct AnalyticalSolution
    inradius::Any
    outradius::Any
    center::Any
    A1c::Any
    A1s::Any
    A2s::Any
    ls::Any
    ms::Any
    lc::Any
    mc::Any
    theta0::Any
    function AnalyticalSolution(inradius, outradius, center, ls, ms, lc, mc, theta0)
        a = analytical_coefficient_matrix(inradius, outradius, ls, ms, lc, mc)
        r = analytical_coefficient_rhs(ls, ms, theta0)
        coeffs = a \ r
        new(
            inradius,
            outradius,
            center,
            coeffs[1],
            coeffs[2],
            coeffs[3],
            ls,
            ms,
            lc,
            mc,
            theta0,
        )
    end
end

function radial_displacement(A::AnalyticalSolution, r)
    if r <= A.inradius
        return A.A1c * r
    else
        return A.A1s * r + A.A2s / r
    end
end

function (A::AnalyticalSolution)(x)
    relpos = x - A.center
    r = sqrt(relpos' * relpos)
    ur = radial_displacement(A, r)
    if ur ≈ 0.0
        [0.0, 0.0]
    else
        costheta = (x[1] - A.center[1]) / r
        sintheta = (x[2] - A.center[2]) / r
        u1 = ur * costheta
        u2 = ur * sintheta
        return [u1, u2]
    end
end

function shell_radial_stress(ls, ms, theta0, A1, A2, r)
    return (ls + 2ms) * (A1 - A2 / r^2) + ls * (A1 + A2 / r^2) - (ls + 2ms / 3) * theta0
end

function shell_circumferential_stress(ls, ms, theta0, A1, A2, r)
    return ls * (A1 - A2 / r^2) + (ls + 2ms) * (A1 + A2 / r^2) - (ls + 2ms / 3) * theta0
end

function shell_out_of_plane_stress(ls, ms, A1, theta0)
    return 2 * ls * A1 - (ls + 2ms / 3) * theta0
end

function core_in_plane_stress(lc, mc, A1)
    return (lc + 2mc) * A1 + lc * A1
end

function core_out_of_plane_stress(lc, A1)
    return 2 * lc * A1
end

function rotation_matrix(x, r)
    costheta = x[1] / r
    sintheta = x[2] / r
    Q = [
        costheta -sintheta
        sintheta costheta
    ]
    return Q
end

function shell_stress(A::AnalyticalSolution, x)
    relpos = x - A.center
    r = sqrt(relpos' * relpos)
    Q = rotation_matrix(relpos, r)

    srr = shell_radial_stress(A.ls, A.ms, A.theta0, A.A1s, A.A2s, r)
    stt = shell_circumferential_stress(A.ls, A.ms, A.theta0, A.A1s, A.A2s, r)

    cylstress = [
        srr 0.0
        0.0 stt
    ]

    cartstress = Q * cylstress * Q'
    s11 = cartstress[1, 1]
    s22 = cartstress[2, 2]
    s12 = cartstress[1, 2]
    s33 = shell_out_of_plane_stress(A.ls, A.ms, A.A1s, A.theta0)

    return [s11, s22, s12, s33]
end

function core_stress(A::AnalyticalSolution)
    s11 = core_in_plane_stress(A.lc, A.mc, A.A1c)
    s33 = core_out_of_plane_stress(A.lc, A.A1c)
    return [s11, s11, 0.0, s33]
end

function exact_stress(A::AnalyticalSolution, x)
    relpos = x - A.center
    r = sqrt(relpos' * relpos)
    if r < A.inradius
        return core_stress(A)
    else
        return shell_stress(A, x)
    end
end

function onleftboundary(x, L, W)
    return x[1] ≈ 0.0
end

function onbottomboundary(x, L, W)
    return x[2] ≈ 0.0
end

function onrightboundary(x, L, W)
    return x[1] ≈ L
end

function ontopboundary(x, L, W)
    return x[2] ≈ W
end

function onboundary(x, L, W)
    return onbottomboundary(x, L, W) ||
           onrightboundary(x, L, W) ||
           ontopboundary(x, L, W) ||
           onleftboundary(x, L, W)
end

function solve_and_compute_displacement_error(
    width,
    center,
    inradius,
    outradius,
    stiffness,
    theta0,
    nelmts,
    polyorder,
    numqp,
    penaltyfactor,
)

    lambda1, mu1 = CutCell.lame_coefficients(stiffness, +1)
    lambda2, mu2 = CutCell.lame_coefficients(stiffness, -1)
    transfstress = CutCell.plane_strain_transformation_stress(lambda1, mu1, theta0)

    analyticalsolution =
        AnalyticalSolution(inradius, outradius, center, lambda1, mu1, lambda2, mu2, theta0)

    dx = width / nelmts
    meanmoduli = 0.5 * (lambda1 + lambda2 + mu1 + mu2)
    penalty = penaltyfactor / dx * meanmoduli

    basis = TensorProductBasis(2, polyorder)
    mesh = CutCell.Mesh([0.0, 0.0], [width, width], [nelmts, nelmts], basis)
    levelset = InterpolatingPolynomial(1, basis)
    levelsetcoeffs = CutCell.levelset_coefficients(
        x -> -circle_distance_function(x, center, inradius),
        mesh,
    )

    cutmesh = CutCell.CutMesh(levelset, levelsetcoeffs, mesh)
    cellquads = CutCell.CellQuadratures(levelset, levelsetcoeffs, cutmesh, numqp)
    interfacequads = CutCell.InterfaceQuadratures(levelset, levelsetcoeffs, cutmesh, numqp)
    facequads = CutCell.FaceQuadratures(levelset, levelsetcoeffs, cutmesh, numqp)

    bilinearforms = CutCell.BilinearForms(basis, cellquads, stiffness, cutmesh)
    interfacecondition =
        CutCell.incoherent_interface_condition(basis, interfacequads, stiffness, cutmesh, penalty)

    leftdisplacementbc = CutCell.DisplacementComponentCondition(
        x->analyticalsolution(x)[1],
        basis,
        facequads,
        stiffness,
        cutmesh,
        x -> onleftboundary(x, width, width),
        [1.0, 0.0],
        penalty,
    )
    bottomdisplacementbc = CutCell.DisplacementCondition(
        analyticalsolution,
        basis,
        facequads,
        stiffness,
        cutmesh,
        x -> onbottomboundary(x, width, width),
        penalty,
    )
    topdisplacementbc = CutCell.DisplacementCondition(
        analyticalsolution,
        basis,
        facequads,
        stiffness,
        cutmesh,
        x -> ontopboundary(x, width, width),
        penalty,
    )

    sysmatrix = CutCell.SystemMatrix()
    sysrhs = CutCell.SystemRHS()

    CutCell.assemble_bilinear_form!(sysmatrix, bilinearforms, cutmesh)
    CutCell.assemble_interface_condition!(sysmatrix, interfacecondition, cutmesh)
    CutCell.assemble_bulk_transformation_linear_form!(
        sysrhs,
        transfstress,
        basis,
        cellquads,
        cutmesh,
    )
    CutCell.assemble_incoherent_interface_transformation_rhs!(
        sysrhs,
        transfstress,
        basis,
        interfacequads,
        cutmesh,
    )
    CutCell.assemble_traction_force_linear_form!(
        sysrhs,
        x -> exact_stress(analyticalsolution, x)[[1, 3]],
        basis,
        facequads,
        cutmesh,
        x -> onrightboundary(x, width, width),
    )
    CutCell.assemble_traction_force_component_linear_form!(
        sysrhs,
        x -> -exact_stress(analyticalsolution, x)[3],
        basis,
        facequads,
        cutmesh,
        x -> onleftboundary(x, width, width),
        [0.0, 1.0],
    )
    CutCell.assemble_penalty_displacement_bc!(
        sysmatrix,
        sysrhs,
        leftdisplacementbc,
        cutmesh,
    )
    CutCell.assemble_penalty_displacement_bc!(
        sysmatrix,
        sysrhs,
        bottomdisplacementbc,
        cutmesh,
    )
    CutCell.assemble_penalty_displacement_bc!(sysmatrix, sysrhs, topdisplacementbc, cutmesh)
    CutCell.assemble_penalty_displacement_component_transformation_rhs!(
        sysrhs,
        transfstress,
        basis,
        facequads,
        cutmesh,
        x -> onleftboundary(x, width, width),
        [1.0, 0.0],
    )
    CutCell.assemble_penalty_displacement_transformation_rhs!(
        sysrhs,
        transfstress,
        basis,
        facequads,
        cutmesh,
        x -> onbottomboundary(x, width, width),
    )
    CutCell.assemble_penalty_displacement_transformation_rhs!(
        sysrhs,
        transfstress,
        basis,
        facequads,
        cutmesh,
        x -> ontopboundary(x, width, width),
    )

    matrix = CutCell.make_sparse(sysmatrix, cutmesh)
    rhs = CutCell.rhs(sysrhs, cutmesh)

    sol = matrix \ rhs
    disp = reshape(sol, 2, :)

    err = mesh_L2_error(disp, analyticalsolution, basis, cellquads, cutmesh)
    den = integral_norm_on_cut_mesh(analyticalsolution, cellquads, cutmesh, 2)

    return err ./ den
end

lambda1, mu1 = 100.0, 80.0
lambda2, mu2 = 80.0, 60.0
theta0 = -0.067
stiffness = CutCell.HookeStiffness(lambda1, mu1, lambda2, mu2)

width = 1.0
penaltyfactor = 1e2

polyorder = 2
numqp = required_quadrature_order(polyorder) + 2

center = [width / 2, width / 2]
inradius = width / 4
outradius = width

powers = [3,4,5]
nelmts = [2^p + 1 for p in powers]

err = [
    solve_and_compute_displacement_error(
        width,
        center,
        inradius,
        outradius,
        stiffness,
        theta0,
        ne,
        polyorder,
        numqp,
        penaltyfactor,
    ) for ne in nelmts
]

dx = 1.0 ./ nelmts
u1err = [er[1] for er in err]
u2err = [er[2] for er in err]

u1rate = convergence_rate(dx,u1err)
u2rate = convergence_rate(dx,u2err)

@test allapprox(u1rate,repeat([3.0],length(u1rate)),0.15)
@test allapprox(u2rate,repeat([3.0],length(u2rate)),0.15)


export GaussianKernel

"Gaussian smoothing kernel."
struct GaussianKernel{N} <: KernelFormula{N}
    support::Float
    σ::Float

    GaussianKernel{1}(; support::Number = 3.0) =
        @check(support > 0) && new(support, 1 / √π)
    GaussianKernel{2}(; support::Number = 3.0) =
        @check(support > 0) && new(support, 1 / π)
    GaussianKernel{3}(; support::Number = 3.0) =
        @check(support > 0) && new(support, 1 / (π * √π))
end

Base.show(io::IO, ::Type{GaussianKernel{N}}) where {N} =
    print(io, "GaussianKernel")

"""
    (W::GaussianKernel)(q::Number) -> Float

Evaluate the kernel with unit width at relative distance `q`.
"""
(W::GaussianKernel)(q::Number) = @dbgassert(q ≥ 0) && W.σ * exp(-q^2)

"""
    ddr(W::GaussianKernel, q::Number) -> Float

Evaluate the derivative of the kernel with unit width at relative distance `q`.
The derivative is with respect to `r`.
"""
ddr(W::GaussianKernel, q::Number) = -2q * W(q)

"""
    ddh(W::GaussianKernel, q::Number) -> Float

Evaluate the derivative of the kernel with unit width at relative distance `q`.
The derivative is with respect to `h`.
"""
ddh(W::GaussianKernel{N}, q::Number) where {N} = (2 * q^2 - N) * W(q)

get_support(W::GaussianKernel) = W.support


using Roots

# Unoptimized reference implementation
# NB: ϕ is scaled by √2π compared to standard definition
function SD_reference(α, X::AbstractArray{T}) where T
	# (n(n-1))⁻¹α⁻⁵∑ᵢ∑ⱼϕⁱᵛ(α⁻¹(Xᵢ-Xⱼ))
	n = length(X)
	C = 1/(n*(n-1)*α^5)
	Y = X./α
	s = zero(T)
	@inbounds for i=1:n
		for j=i+1:n
			s += 2*ϕ4(Y[i]-Y[j])
		end
	end
	s += n*ϕ4(0)
	C * s
end

# NB: ϕ is scaled by √2π compared to standard definition
function TD_reference(b, X::AbstractArray{T}) where T
	# -(n(n-1))⁻¹b⁻⁷∑ᵢ∑ⱼϕᵛⁱ(b⁻¹(Xᵢ-Xⱼ))
	n = length(X)
	C = -1/(n*(n-1)*b^7)
	Y = X./b
	s = zero(T)
	@inbounds for i=1:n
		for j=i+1:n
			s += 2*ϕ6(Y[i]-Y[j])
		end
	end
	s += n*ϕ6(0)
	C * s
end

function objective_reference(h, X, α2Constant)
	# [ R(K) / (nσ⁴ₖSD(α₂(h))) ]^(1/5) - h = 0
	# R(K) = 1/(2√π)
	# σ⁴ₖ = 1
	n = length(X)
	α2 = α2Constant*h^(5/7) # 1.357[SD(a)/TD(b)]^(1/7) * h^(5/7)
	1/(√2*n*h^5) - SD_reference(α2,X)
end

function _bwsj_reference(X; rtol=0.1)
	n = length(X)
	q25,q75 = quantile(X,(0.25,0.75))
	λ = min(q75-q25, std(X)*1.349) # Givivng the same scale estimate as e.g. the KernSmooth R package.
	a = 0.920λ*n^(-1/7)
	b = 0.912λ*n^(-1/9)
	α2Constant = 1.357*(SD_reference(a,X)/TD_reference(b,X))^(1/7)

	# Figure out lower and upper bounds for search.
	# The Bandwidth scales roughly with n^(-1/5). TODO: Use this fact!

	lower = 0.1*λ
	for i=1:11
		i==11 && error("Failed to find lower bound")
		objective_reference(lower,X,α2Constant)>0 && break
		lower/=10
	end
	upper = 1*λ
	for i=1:11
		i==11 && error("Failed to find upper bound")
		objective_reference(upper,X,α2Constant)<0 && break
		upper*=10
	end

	find_zero(h->objective_reference(h,X,α2Constant), (lower,upper), Bisection(), xrtol=rtol)
end

bwsj_reference(X; kwargs...) = _bwsj_reference(issorted(X) ? X : sort(X); kwargs...)

function _smooth_reference(x,y,c,xeval,w)
	w .= -c.*(x.-xeval).^2
	w .= exp.(w .- maximum(w)) # multiply all weights with a factor for improved precision
	y'w / sum(w)
end

function smooth_reference(x::AbstractVector{<:T}, y::AbstractVector{<:S}, bandwidth, xeval) where {T<:Real,S<:Real}
	@assert length(x)==length(y)
	c = 1 ./(2 .* bandwidth.^2)
	w = similar(x, promote_type(T,eltype(bandwidth),eltype(xeval)))
	_smooth_reference.(Ref(x),Ref(y),c,xeval,Ref(w))
end

function _density_reference(x,bandwidth,c,xeval,w)
	w .= exp.(-c.*(x.-xeval).^2)
	sum(w) / (sqrt(2π)*bandwidth*length(x))
end

function density_reference(x::AbstractVector{<:T}, bandwidth, xeval) where {T<:Real}
	c = 1 ./(2 .* bandwidth.^2)
	w = similar(x, promote_type(T,eltype(bandwidth),eltype(xeval)))
	_density_reference.(Ref(x),bandwidth,c,xeval,Ref(w))
end


using Flux
using Flux: @epochs
using Flux: params
using Flux: throttle
using Plots
using Distributions
using DataFrames
using CSV

@inline phi(x) = @. x/(1 + exp(-x))
@inline f(x) = 2*π^2 .*sum(cos.(π*x[i,:]') for i=1:2, dims=1)

#build a block
block_layer = SkipConnection(Chain(Dense(10,10, phi),Dense(10,10, phi)), +)

#construct network skeleton
M = Chain(Dense(2,10),
block_layer,
block_layer,
block_layer,
block_layer,
Dense(10,1))

#check number of params
params(M)

#inner loss function
function inner_loss(N, space)
    a = rand(Uniform(0,1),2,N)
    loss = 0
    l1 = [1, 0]
    l2 = [0, 1]

    @inbounds loss += (sum(0.5*((M(a .+ space.*l1) .- M(a))./space).^2)/N)[1]
    @inbounds loss += (sum(0.5*((M(a .+ space.*l2) .- M(a))./space).^2)/N)[1]
    @inbounds loss += (sum(0.5*(π^2 * (M(a)).^2))/N)[1]
    @inbounds loss += (sum(-f(a) .*M(a))/N)[1]

    return loss
end

function bc_loss(β, N = 1000)

    d = (rand(1)/100)
    point = zeros(2,1000)
    p = rand(1,100)
    point = hcat(hcat(zeros(250), rand(250))', hcat(ones(250), rand(250))', hcat(rand(250), zeros(250))', hcat(rand(250), ones(250))')
    loss = 0
    for i=1:2
        if i == 1
            l1 = [1, 0]
        else
            l1 = [0, 1]
        end
        @inbounds loss += β*(sum(((1/2*M(point[:,500*(i-1)+1:500*i] .+ d.*l1) .- 1/2*M(point[:,500*(i-1)+1:500*i] .- d.*l1))./d).^2)/100)[1]
    end
    return loss
end

#add losses
function total_loss(N, space = 0.01, β = 50)
    return inner_loss(N, space) + bc_loss(β)
end

function loss_true(xx=1,yy=1,xx1=1,xx2=2)
  points=rand(2,2000)
  F_true(x)= sum(cos.(pi*x[i,:]') for i=1:2, dims=1)
  errors=sqrt(sum((M(points)-F_true(points)).^2)/2000)
  return errors
end

#train model
function train_model(N = 100, train_time = 20)
    df = DataFrame(time = Int[], in_loss = Float64[],bound_loss = Float64[],true_loss = Float64[])
    c = 1
    evalcb = function()
        @show(total_loss(N), loss_true())
        in_loss = inner_loss(N, 0.01)
        bound_loss = bc_loss(50)
        true_loss = loss_true()
        push!(df, [c, in_loss, bound_loss, true_loss])
        c += 1
        if c == train_time
            CSV.write("$(c)_2d_neumann.csv",df)
        end
    end
    opt = ADAM()
    dataset = [(N) for i=1:train_time]
    Flux.train!(total_loss, params(M), zip(dataset), opt, cb=throttle(evalcb, 50))
end

@epochs 1000 train_model(4000)

#plot the solution over a mesh in the domain
F_true(x)= sum(cos.(pi*x[i,:]') for i=1:2, dims=1)
function plot_surface()
    data = zeros(100,100)
    grid_space = [i for i in 1:100]
    for i in grid_space
        for j in grid_space
            point = [i/100 j/100]
            data[i,j] = M(point')[1]
        end
    end

    gr()
    heatmap(1:size(data,1),
        1:size(data,2), data',
        c = cgrad([:blue, :white,:red, :yellow]),
        xlabel = "X", ylabel = "Y",
        title = "2D Neumann Fitted Solution")

    savefig("2d_neumann_fitted.png")

    for i in grid_space
        for j in grid_space
            point = [i/100 j/100]
            data[i,j] = F_true(point')[1]
        end
    end

    gr()
    heatmap(1:size(data,1),
        1:size(data,2), data',
        c = cgrad([:blue, :white,:red, :yellow]),
        xlabel = "X", ylabel = "Y",
        title = "2D Neumann True Solution")

    savefig("2d_neumann_true.png")
end

plot_surface()


struct StrangSplitting{T,R,V} <: AbstractIntegrator{T}
    uv::Tuple{V,V}
    E::Vector{<:AbstractExponentiator{T}}
    function StrangSplitting(u::V, E::AbstractExponentiator{T}...) where {T,V}
        foreach(e -> checkorder(e,StrangSplitting), E)
        new{T,real(T),V}((u,similar(u)), Vector{AbstractExponentiator{T}}(E...))
    end
end

order(::StrangSplitting) = 2
order(::Type{StrangSplitting}) = 2

function step!(strang::StrangSplitting{T,R,V}, t::R, Δt::R) where {T,R,V}
    # Integrate all operators over time-step, to second order since
    # the Strang splitting itself introduces a second-order splitting
    # error.
    foreach(e -> midpoint!(operator(e), t, Δt), strang.E)

    a,b = 1,2

    # When the operator is L = A + B + C + D + ..., the Strang
    # splitting is ... (D/2) (C/2) (B/2) A (B/2) (C/2) (D/2) ...

    # ... (D/2) (C/2) (B/2)
    for E in reverse(strang.E[2:end])
        expv!(strang.uv[b], E, Δt/2, strang.uv[a], -1)
        inplace(E) || @swap! a b
    end

    # A
    expv!(strang.uv[b], E₁, Δt, strang.uv[a], 0)
    inplace(strang.E[1]) || @swap! a b

    # (B/2) (C/2) (D/2) ...
    for E in strang.E[2:end]
        expv!(strang.uv[b], E, Δt/2, strang.uv[a], 1)
        inplace(E) || @swap! a b
    end

    # If the propagated solution is in the v vector, copy it to the u
    # vector.
    a == 2 && copyto!(strang.uv[1], strang.uv[2])
end


using DataStructures

function main()
    N, M = parse.(Int, split(readline()))
    As, Bs = zeros(Int, M), zeros(Int, M)
    for i in 1:M
        As[i], Bs[i] = parse.(Int, split(readline()))
    end

    ans = 0
    for i in 1:M
        ds = IntDisjointSets(N)
        for j in 1:M
            if j ≠ i
                union!(ds, As[j], Bs[j])
            end
        end

        cnt = 0
        for j in 1:N
            if j == find_root(ds, j)
                cnt += 1
            end
        end

        if cnt > 1
            ans += 1
        end
    end

    println(ans)
end

main()


using LinearAlgebra

# one-body matrix, nearest neighbour, PBC
function kinetic_matrix(self::Lattice)
    K_matrix = zeros(self.nsites, self.nsites)
    for i =  1 : self.nsites
        for j = 1 : self.nsites
            if abs(i - j) == 1 || abs(i - j) == self.nsites - 1
                K_matrix[i, j] = -self.t
            end
        end
    end
    return K_matrix
end

# σ-dependent generator
function auxiliary_field_matrix_stepwise(σ::Array{Int64,1}, γ::Float64, Δτ::Float64, U::Float64)
    AF_matrix_SpinUp = diagm(0 => exp.(2 * γ * σ .- Δτ * U / 2))
    AF_matrix_SpinDn = diagm(0 => exp.(-2 * γ * σ .- Δτ * U / 2))
    return AF_matrix_SpinUp, AF_matrix_SpinDn
end

function matrix_product_stepwise(σlist::Array{Int64,2}, lattice::Lattice, qmc::QMC)
    MatProdUp = Matrix{Float64}(I, lattice.nsites, lattice.nsites)
    MatProdDn = Matrix{Float64}(I, lattice.nsites, lattice.nsites)
    for i = 1 : qmc.L
        σ = σlist[:, i]
        AF_matrix_SpinUp, AF_matrix_SpinDn = auxiliary_field_matrix_stepwise(σ, qmc.γ, qmc.Δτ, lattice.U)
        MatProdUp = Symmetric(qmc.Kmatrix * AF_matrix_SpinUp * qmc.Kmatrix) * MatProdUp
        MatProdDn = Symmetric(qmc.Kmatrix * AF_matrix_SpinDn * qmc.Kmatrix) * MatProdDn
    end
    EigensUp = eigen(MatProdUp)
    EigensDn = eigen(MatProdDn)
    return MatProdUp, MatProdDn, EigensUp, EigensDn
end

# Use QR decomposition to stablize the propogation
function matrix_product_QR(σlist::Array{Int64,2}, lattice::Lattice, qmc::QMC, n_stab_interval::Int64)
    # Initialize QR matrices
    QRUp_Q = Matrix{Float64}(I, lattice.nsites, lattice.nsites)
    QRUp_R = Matrix{Float64}(I, lattice.nsites, lattice.nsites)
    QRDn_Q = Matrix{Float64}(I, lattice.nsites, lattice.nsites)
    QRDn_R = Matrix{Float64}(I, lattice.nsites, lattice.nsites)
    # Final Product Matrix
    MatProdUp = Matrix{Float64}(I, lattice.nsites, lattice.nsites)
    MatProdDn = Matrix{Float64}(I, lattice.nsites, lattice.nsites)
    for i = 1 : div(qmc.L, n_stab_interval)
        matprodUp = Matrix{Float64}(I, lattice.nsites, lattice.nsites)
        matprodDn = Matrix{Float64}(I, lattice.nsites, lattice.nsites)
        for j = 1 : n_stab_interval
            σ = σlist[:, (i - 1) * n_stab_interval + j]
            AF_matrix_SpinUp, AF_matrix_SpinDn = auxiliary_field_matrix_stepwise(σ, qmc.γ, qmc.Δτ, lattice.U)
            matprodUp = Symmetric(qmc.Kmatrix * AF_matrix_SpinUp * qmc.Kmatrix) * matprodUp
            matprodDn = Symmetric(qmc.Kmatrix * AF_matrix_SpinDn * qmc.Kmatrix) * matprodDn
        end
        # QR decomposition of the matrix product
        # leave here to see if can be replaced with rmul!
        QRUp_Q = matprodUp * QRUp_Q
        DecompUp = qr(QRUp_Q)
        QRDn_Q = matprodDn * QRDn_Q
        DecompDn = qr(QRDn_Q)
        # Remaining R_i goes to multiply R_1
        QRUp_Q = DecompUp.Q
        QRUp_R = DecompUp.R * QRUp_R
        QRDn_Q = DecompDn.Q
        QRDn_R = DecompDn.R * QRDn_R
    end
    mul!(MatProdUp, QRUp_Q, QRUp_R)
    mul!(MatProdDn, QRDn_Q, QRDn_R)
    EigensUp = eigen(MatProdUp)
    EigensDn = eigen(MatProdDn)
    return MatProdUp, MatProdDn, EigensUp, EigensDn
end

# generate overlap "tensor" / column major
function overlap_coefficient(nsites::Int64, Uup::Array{T1,2}, Udn::Array{T2,2}) where {T1<:FloatType, T2<:FloatType}
    Uup_inv, Udn_inv = inv(Uup), inv(Udn)
    Uup_array = complex(zeros(nsites, nsites, nsites))
    Udn_array = complex(zeros(nsites, nsites, nsites))
    for i = 1 : nsites
        for j = 1 : nsites
            for k = 1 : nsites
                Uup_array[k ,i, j] = Uup_inv[k, i] * Uup[j, k]
                Udn_array[k, i, j] = Udn_inv[k, i] * Udn[j, k]
            end
        end
    end
    real(Uup_array), real(Udn_array)
end


"""
    SeisEnvelope(d)

Calculate the envelope attribute of an input trace

# Arguments
* `d`: Input data.

# Output
* `out`: Envelope of input data

# Example
```julia
julia> d = SeisLinearEvents(d); SeisPlot(d)
julia> out = SeisEnvelope(d); SeisPlot(d_dec)
"""

function SeisEnvelope(d)

	D = fft(d,1)
	D[1:Int(floor(size(d,1)/2)),:] = 0.0
	return 2*abs.(ifft(D,1))

end

function SeisEnvelope(A,h::Array{Header,1})
	for itrace = 1 : size(A[:,:],2)
		A[:,itrace] = SeisEnvelope(A[:,itrace])
	end
	return A,h
end

function SeisEnvelope(in::AbstractString,out::AbstractString)
	@compat parameters = Dict()
	SeisProcess(in,out,[SeisEnvelope],[parameters];group="some",ntrace=100000)
end


include("diagonal_sbp.jl")
include("deriv_ops_new.jl")
include("split_matrix_free_dev.jl")
include("split_matrix_free.jl")

# using CUDAdrv
# CUDAdrv.CuDevice(0)

# Solving Poisson Equation
# Δ u(x,y) = f(x,y)
# Manufactured Solution: u(x,y) = sin(πx .+ πy)

# using CuArrays, CUDAnative
using LinearAlgebra
using SparseArrays
using Plots
using CUDA
using IterativeSolvers
using BenchmarkTools
using MAT

function e(i,n)
    # A = Matrix{Float64}(I,n,n)
    # return A[:,i]
    out = spzeros(n)
    out[i] = 1.0
    return out
end

function eyes(n)
    # return Matrix{Float64}(I,n,n)
    out = spzeros(n,n)
    for i in 1:n
        out[i,i] = 1.0
    end
    return out
end

function u(x,y)
    return sin.(π*x .+ π*y)
end

function Diag(A)
    # Self defined function that is similar to Matlab Diag
    return Diagonal(A[:])
end

function Operators_2d(i, j, p=2, h_list_x = ([1/2^1, 1/2^2, 1/2^3, 1/2^4, 1/2^5, 1/2^6, 1/2^7, 1/2^8, 1/2^9, 1/2^10, 1/2^11, 1/2^12, 1/2^13]),
			 h_list_y = ([1/2^1, 1/2^2, 1/2^3, 1/2^4, 1/2^5, 1/2^6, 1/2^7, 1/2^8, 1/2^9, 1/2^10, 1/2^11, 1/2^12, 1/2^13])
			 )
    hx = h_list_x[i];
    hy = h_list_y[j];

    x = range(0,step=hx,1);
    y = range(0,step=hy,1);
    m_list = 1 ./h_list_x;
    n_list = 1 ./h_list_y;

    # Matrix Size
    N_x = Integer(m_list[i]);
    N_y = Integer(n_list[j]);

    (D1x, HIx, H1x, r1x) = diagonal_sbp_D1(p,N_x,xc=(0,1));
    (D2x, S0x, SNx, HI2x, H2x, r2x) = diagonal_sbp_D2(p,N_x,xc=(0,1));

    (D1y, HIy, H1y, r1y) = diagonal_sbp_D1(p,N_y,xc=(0,1));
    (D2y, S0y, SNy, HI2y, H2y, r2y) = diagonal_sbp_D2(p,N_y,xc=(0,1));

    # BSx = sparse(SNx - S0x);
    # BSy = sparse(SNy - S0y);
    BSx = SNx - S0x
    BSy = SNy - S0y

    # Forming 2d Operators
    # e_1x = sparse(e(1,N_x+1));
    # e_Nx = sparse(e(N_x+1,N_x+1));
    # e_1y = sparse(e(1,N_y+1));
    # e_Ny = sparse(e(N_y+1,N_y+1));
    e_1x = e(1,N_x+1);
    e_Nx = e(N_x+1,N_x+1);
    e_1y = e(1,N_x+1);
    e_Ny = e(N_y+1,N_y+1);

    # I_Nx = sparse(eyes(N_x+1));
    # I_Ny = sparse(eyes(N_y+1));
    I_Nx = eyes(N_x+1);
    I_Ny = eyes(N_y+1);

    e_E = kron(e_Nx,I_Ny);
    e_W = kron(e_1x,I_Ny);
    e_S = kron(I_Nx,e_1y);
    e_N = kron(I_Nx,e_Ny);

    E_E = kron(sparse(Diag(e_Nx)),I_Ny);   # E_E = e_E * e_E'
    E_W = kron(sparse(Diag(e_1x)),I_Ny);
    E_S = kron(I_Nx,sparse(Diag(e_1y)));
    E_N = sparse(kron(I_Nx,sparse(Diag(e_Ny))));

    D1_x = kron(D1x,I_Ny);
    D1_y = kron(I_Nx,D1y);

    D2_x = kron(D2x,I_Ny);
    D2_y = kron(I_Nx,D2y);
    D2 = D2_x + D2_y

    HI_x = kron(HIx,I_Ny);
    HI_y = kron(I_Nx,HIy);

    H_x = kron(H1x,I_Ny);
    H_y = kron(I_Nx,H1y);

    BS_x = kron(BSx,I_Ny);
    BS_y = kron(I_Nx,BSy);

    HI_tilde = kron(HIx,HIx);
    H_tilde = kron(H1x,H1y);

    return (D1_x, D1_y, D2_x, D2_y, D2, HI_x, HI_y, BS_x, BS_y, HI_tilde, H_tilde, I_Nx, I_Ny, e_E, e_W, e_S, e_N, E_E, E_W, E_S, E_N)
end

h_list_x = [1/2^1, 1/2^2, 1/2^3, 1/2^4, 1/2^5, 1/2^6, 1/2^7, 1/2^8, 1/2^9, 1/2^10, 1/2^11, 1/2^12, 1/2^13]
h_list_y = [1/2^1, 1/2^2, 1/2^3, 1/2^4, 1/2^5, 1/2^6, 1/2^7, 1/2^8, 1/2^9, 1/2^10, 1/2^11, 1/2^12, 1/2^13]

rel_errs = []
iter_errs = []
# for k in 1:length(h_list_x)
println("################### BEGIN TEST #########################")
for k in 6:13

    println()
    i = j  = k
    println("##########   Starting Test for k = ", k, "   ######################")

    hx = h_list_x[i];
    hy = h_list_y[j];

    x = range(0,step=hx,1);
    y = range(0,step=hy,1);
    m_list = 1 ./h_list_x;
    n_list = 1 ./h_list_y;

    # Matrix Size
    N_x = Integer(m_list[i]);
    N_y = Integer(n_list[j]);

    Nx = N_x + 1;
    Ny = N_y + 1;

    println("Nx = $Nx, Ny=$Ny")

    # 2D operators
    (D1_x, D1_y, D2_x, D2_y, D2, HI_x, HI_y, BS_x, BS_y, HI_tilde, H_tilde, I_Nx, I_Ny, e_E, e_W, e_S, e_N, E_E, E_W, E_S, E_N) = Operators_2d(i,j);

     # Analytical Solutions
    analy_sol = u(x,y');

    # Penalty Parameters
    tau_E = -13/hx;
    tau_W = -13/hx;
    tau_N = -1;
    tau_S = -1;

    beta = 1;

    # Forming SAT terms

    ## Formulation 1
    SAT_W = tau_W*HI_x*E_W + beta*HI_x*BS_x'*E_W;
    SAT_E = tau_E*HI_x*E_E + beta*HI_x*BS_x'*E_E;

    # SAT_S = tau_S*HI_y*E_S*D1_y
    # SAT_N = tau_N*HI_y*E_N*D1_y

    SAT_S = tau_S*HI_y*E_S*BS_y;
    SAT_N = tau_N*HI_y*E_N*BS_y;

    SAT_W_r = tau_W*HI_x*E_W*e_W + beta*HI_x*BS_x'*E_W*e_W;
    SAT_E_r = tau_E*HI_x*E_E*e_E + beta*HI_x*BS_x'*E_E*e_E;
    SAT_S_r = tau_S*HI_y*E_S*e_S;
    SAT_N_r = tau_N*HI_y*E_N*e_N;

    (alpha1,alpha2,alpha3,alpha4,beta) = (tau_N,tau_S,tau_W,tau_E,beta);

    g_W = sin.(π*y);
    g_E = -sin.(π*y);
    g_S = -π*cos.(π*x);
    # g_N = -π*cos.(π*x)
    g_N = π*cos.(π*x .+ π);

    # Solving with CPU
    A = D2 + SAT_W + SAT_E + SAT_S + SAT_N;

    b = -2π^2*u(x,y')[:] + SAT_W_r*g_W + SAT_E_r*g_E + SAT_S_r*g_S + SAT_N_r*g_N;

    A = - H_tilde*A;
    b = - H_tilde*b;

    @show nnz(A) * sizeof(Float64)

    @show Base.summarysize(A)
    @show Base.summarysize(b)
    @show Nx * Ny * sizeof(Float64)

    @show Base.summarysize(A) / Base.summarysize(b)

    println()

    A_d = CUDA.CUSPARSE.CuSparseMatrixCSC(A);
    b_d = CuArray(b);

    # @show Base.summarysize(A_d)
    # @show Base.summarysize(A_d)

    # testing matrix_split method
    b_reshaped = reshape(b,Nx,Ny);
    b_reshaped_GPU = CuArray(b_reshaped);

    x_GPU = similar(b_reshaped_GPU)
    x_full_GPU = similar(b_reshaped_GPU)
    A_d * b_d
    matrix_free_A(b_reshaped_GPU,x_GPU)
    matrix_free_A_full_GPU(b_reshaped_GPU,x_full_GPU)

    odata = similar(b_reshaped_GPU);
    r_GPU = similar(b_reshaped_GPU);
    p_GPU = similar(b_reshaped_GPU);
    Ap_GPU = similar(b_reshaped_GPU);

    CG_GPU_v2(b_reshaped_GPU,x_GPU,odata,r_GPU,p_GPU,Ap_GPU)

    ## SpMV test

    println("TESTING SpMV vs Matrix_FREE")
    iter_times = Nx + Ny

    time_matrix_free = @elapsed begin
        for _ in 1:iter_times
            matrix_free_A(b_reshaped_GPU,x_GPU)
        end
    end
    time_matrix_free = time_matrix_free * 1000 / iter_times
    @show time_matrix_free

    time_matrix_free_full_GPU = @elapsed begin
        for _ in 1:iter_times
            matrix_free_A_full_GPU(b_reshaped_GPU,x_full_GPU)
        end
    end
    time_matrix_free_full_GPU = time_matrix_free_full_GPU * 1000 / iter_times
    @show time_matrix_free_full_GPU

    time_CUBLAS = @elapsed begin
        for _ in 1:iter_times
            A_d * b_d
        end
    end
    time_CUBLAS = time_CUBLAS * 1000 / iter_times
    @show time_CUBLAS
    println()

    ## End SpMV test

    x_GPU = CuArray(zeros(Nx,Ny));
    (iter_steps,rsold_GPU) = CG_GPU(b_reshaped_GPU,x_GPU);
    println("Iteration steps till convergence for x_GPU: $iter_steps")
    @show rsold_GPU
    matrix_free_GPU_time_in_CG = (time_matrix_free * iter_steps)

    x_full_GPU = CuArray(zeros(Nx,Ny));
    (iter_steps,rsold_full_GPU) = CG_full_GPU(b_reshaped_GPU,x_full_GPU);
    matrix_free_full_GPU_time_in_CG = (time_matrix_free_full_GPU * iter_steps)
    println("Iteration steps till convergence for x_full_GPU: $iter_steps")
    @show rsold_full_GPU

    _,history= cg(A_d,b_d,log=true)
    @show history
    @show history.data[:resnorm][end]
    @show history.data[:resnorm][end-1]

    # x_GPU = CuArray(zeros(Nx,Ny));
    #         # CG_GPU_dev(b_reshaped_GPU,x_GPU)
    # CG_GPU(b_reshaped_GPU,x_GPU)

    iter_times = min(5,max(13-k,1))
    println()
    println("Starting Timing, results in ms")

    println("Time for matrix_free SpMV in CG: $matrix_free_GPU_time_in_CG")
    println("Time for matrix_free_full_GPU SpMV in CG: $matrix_free_full_GPU_time_in_CG")
    CUBLAS_SpMV_time_in_CG = (time_CUBLAS * history.iters)
    println("Time for CUBLAS SpMV in CG_IterativeSolvers: $CUBLAS_SpMV_time_in_CG\n")

    println()

    t_CG_GPU = @elapsed begin
        for i in 1:iter_times
            x_GPU = CuArray(zeros(Nx,Ny))
            # CG_GPU_dev(b_reshaped_GPU,x_GPU)
            CG_GPU(b_reshaped_GPU,x_GPU)
        end
    end
    t_CG_GPU = t_CG_GPU * 1000 / iter_times
    @show t_CG_GPU

    # t_CG_GPU_v2 = @elapsed begin # testing preallocation
    #     for i in 1:iter_times
    #         x_GPU = CuArray(zeros(Nx,Ny))
    #         # CG_GPU_dev(b_reshaped_GPU,x_GPU)
    #         CG_GPU_v2(b_reshaped_GPU,x_GPU,odata,r_GPU,p_GPU,Ap_GPU)
    #     end
    # end
    # t_CG_GPU_v2 = t_CG_GPU_v2 * 1000 / iter_times
    # @show t_CG_GPU_v2

    t_CG_full_GPU = @elapsed begin
        for i in 1:iter_times
            x_full_GPU = CuArray(zeros(Nx,Ny))
            # CG_GPU_dev(b_reshaped_GPU,x_GPU)
            CG_full_GPU(b_reshaped_GPU,x_full_GPU)
        end
    end
    t_CG_full_GPU = t_CG_full_GPU * 1000 / iter_times
    @show t_CG_full_GPU

    t_CG_GPU_IterativeSolvers = @elapsed begin
        for i in 1:iter_times
            cg(A_d,b_d,log=true)
        end
    end
    t_CG_GPU_IterativeSolvers = t_CG_GPU_IterativeSolvers * 1000 / iter_times
    @show t_CG_GPU_IterativeSolvers

    println("End time comparison\n")

    ## Compare Efficiency

    overhead_CG_matrix_free =  t_CG_GPU - matrix_free_GPU_time_in_CG
    overhead_CG_matrix_free_full_GPU = t_CG_full_GPU - matrix_free_full_GPU_time_in_CG
    overhead_CG_IterativeSolvrs = t_CG_GPU_IterativeSolvers - CUBLAS_SpMV_time_in_CG

    @show overhead_CG_matrix_free
    @show overhead_CG_matrix_free_full_GPU
    @show overhead_CG_IterativeSolvrs

    println()

    efficiency_CG_matrix_free = matrix_free_GPU_time_in_CG / t_CG_GPU
    efficiency_CG_matrix_free_full_GPU = matrix_free_full_GPU_time_in_CG / t_CG_full_GPU
    efficiency_CG_IterativeSolvrs = CUBLAS_SpMV_time_in_CG / t_CG_GPU_IterativeSolvers

    @show efficiency_CG_matrix_free
    @show efficiency_CG_matrix_free_full_GPU
    @show efficiency_CG_IterativeSolvrs

    println()

    ## End Compare Efficiency
    println("##########   Ending Test for k = ", k, "   ######################\n")

end

function check_richardson(A,ω)
    return opnorm(Matrix(I,size(A)) - ω*A)
end

function modified_richardson(A,b;ω=0.15)
    x = zeros(length(b))
    # ω = -0.15
    count = 0
    norms = []
    for _ in 1:100
        iter_norm = norm(A*x-b)
        append!(norms,iter_norm)
        if iter_norm >= 1e-6 * norms[1]
            count += 1
            x = x + ω*(b- A*x)
        end
    end
    return (x,count,norms)
end

function jacobi_diy(A,b)
    x = zeros(length(b))
    D = Diagonal(A)
    LU = A - D
    # ω = -0.15
    count = 0
    norms = []
    for _ in 1:100
        iter_norm = norm(A*x-b)
        append!(norms,iter_norm)
        if iter_norm >= 1e-6 * norms[1]
            count += 1
            x = D\(b - LU*x)
        end
    end
    return (x,count,norms)
end

function test(A,b;ω=1/13)
    (x_j,count_j,norms_j) = jacobi_diy(A,b)
    (x_r,count_r,norms_r) = modified_richardson(A,b,ω=ω)
    plot(norms_j)
    plot!(norms_r)
end


# Construction of independent buckets of containment boxes

# Given as input a list `randomLineArray` of pairs of 2D points, the function
# `containment2DBoxes` returns, in the same order, the list of _containment boxes_ of the
# input lines. A _containment box_ of a geometric object of dimension $d$ is defined as the
# minimal $d$-cuboid, equioriented with the reference frame, that contains the object. For a
# 2D line it is given by the tuple $(x1,y1,x2,y2)$, where $(x1,y1)$ is the point of minimal
# coordinates, and $(x2,y2)$ is the point of maximal  coordinates.

using LAR

larview(randomLines(1000,0.1)...)

randomLineArray = randomLines(200,0.2);
lines = lar2hpc(randomLineArray...);
yellow = p.COLOR(p.YELLOW)(lines);
boxes = containment2DBoxes(randomLineArray...);
rects = [box2rect(box) for box in boxes];
polylines = p.AA(polyline)(rects);
cyan = p.COLOR(p.CYAN)(p.STRUCT(polylines));
p.VIEW(p.STRUCT([yellow,cyan]))


## Cartpole
import TrajectoryOptimization
traj_folder = joinpath(dirname(pathof(TrajectoryOptimization)),"..")
urdf_folder = joinpath(traj_folder, "dynamics/urdf")
urdf_cartpole = joinpath(urdf_folder, "cartpole.urdf")

cartpole_urdf = Model(urdf_cartpole,[1.;0.]) # underactuated, only control of slider

function cartpole_dynamics!(ẋ::AbstractVector{T}, x::AbstractVector{T}, u::AbstractVector{T}) where T
    mc = 1.0  # mass of the cart in kg (10)
    mp = 0.2    # mass of the pole (point mass at the end) in kg
    l = 0.5   # length of the pole in m
    g = 9.81  # gravity m/s^2

    q = x[1:2]
    qd = x[3:4]

    if isfinite(q[2])
        s = sin(q[2])
        c = cos(q[2])
    else
        s = Inf
        c = Inf
    end

    H = [mc+mp mp*l*c; mp*l*c mp*l^2]
    C = [0 -mp*qd[2]*l*s; 0 0]
    G = [0; mp*g*l*s]
    B = [1; 0]

    qdd = -H\(C*qd + G - B*u[1])

    ẋ[1:2] = qd
    ẋ[3:4] = qdd
    return nothing
end

n,m = 4,1
cartpole = Model(cartpole_dynamics!,n,m)

# unknown friction
function cartpole_dynamics_uncertain!(ẋ, x, u, w)
    mc = 1.0  # mass of the cart in kg (10)
    mp = 0.2    # mass of the pole (point mass at the end) in kg
    l = 0.5   # length of the pole in m
    g = 9.81  # gravity m/s^2
    mu = 0.1

    q = x[1:2]
    qd = x[3:4]

    if isfinite(q[2])
        s = sin(q[2])
        c = cos(q[2])
    else
        s = Inf
        c = Inf
    end

    H = [mc+mp mp*l*c; mp*l*c mp*l^2]
    C = [0 -mp*qd[2]*l*s; 0 0]
    G = [0; mp*g*l*s]
    B = [1; 0]

    qdd = -H\(C*qd + G - B*u[1] - B*w[1])

    ẋ[1:2] = qd
    ẋ[3:4] = qdd

    return nothing
end

n = 4; m = 1; r = 1
cartpole_uncertain = UncertainModel(cartpole_dynamics_uncertain!,n,m,r)


import Statistics: mean
import Base.Threads: @threads, @sync, @spawn, nthreads, threadid

export CloudInCell

struct CloudInCell end

export compute_rho

function compute_rho(p, kernel::CloudInCell, m :: TwoDGrid)

    nx, ny = m.nx, m.ny
    dx, dy = m.dx, m.dy
    rho = zeros(nx+1, ny+1)
    nbpart = size(p.array)[2]

    @inbounds for ipart = 1:nbpart
        xp = p.array[1, ipart] / dx
        yp = p.array[2, ipart] / dy

        i = trunc(Int, xp) + 1
        j = trunc(Int, yp) + 1

        dxp = xp - i + 1
        dyp = yp - j + 1

        a1 = (1 - dxp) * (1 - dyp)
        a2 = dxp * (1 - dyp)
        a3 = dxp * dyp
        a4 = (1 - dxp) * dyp

        ip1 = mod1(i + 1, nx)
        jp1 = mod1(j + 1, ny)

        w = p.array[5, ipart]

        rho[i, j] += a1 * w
        rho[ip1, j] += a2 * w
        rho[ip1, jp1] += a3 * w
        rho[i, jp1] += a4 * w

    end

    for i = 1:nx
       rho[i,1] += rho[i,ny+1]
    end
    for j = 1:ny
       rho[1,j] += rho[nx+1,j]
    end

    rho ./= ( dx * dy )

    return rho[1:nx,1:ny] .- mean(rho[1:nx,1:ny])

end

export compute_current!

function compute_current!(jx, jy, m::TwoDGrid, kernel::CloudInCell, p)

    nbpart = size(p.array)[2]
    nx, ny = m.nx, m.ny
    dx, dy = m.dx, m.dy

    fill!(jx, 0)
    fill!(jy, 0)

    scaling = 1 / ( dx * dy )

    ntid = nthreads()
    jxloc = [zero(jx) for _ in 1:ntid]
    jyloc = [zero(jy) for _ in 1:ntid]

    chunks = Iterators.partition(1:nbpart, nbpart÷ntid)

    @sync for chunk in chunks
        @spawn begin
            tid = threadid()
            @inbounds for ipart in chunk
                xp = p.array[1, ipart] / dx
                yp = p.array[2, ipart] / dy

                i = floor(Int, xp) + 1
                j = floor(Int, yp) + 1

                dxp = xp - i + 1
                dyp = yp - j + 1

                a1 = (1 - dxp) * (1 - dyp)
                a2 = dxp * (1 - dyp)
                a3 = dxp * dyp
                a4 = (1 - dxp) * dyp

                w = p.array[5, ipart]

                w1 = p.array[3, ipart] * scaling * w
                w2 = p.array[4, ipart] * scaling * w

                jxloc[tid][i, j] += a1 * w1
                jyloc[tid][i, j] += a1 * w2

                jxloc[tid][i+1, j] += a2 * w1
                jyloc[tid][i+1, j] += a2 * w2

                jxloc[tid][i+1, j+1] += a3 * w1
                jyloc[tid][i+1, j+1] += a3 * w2

                jxloc[tid][i, j+1] += a4 * w1
                jyloc[tid][i, j+1] += a4 * w2
            end
        end
    end

    jx .= reduce(+,jxloc)
    jy .= reduce(+,jyloc)

    for i = 1:nx+1
        jx[i, 1] += jx[i, ny+1]
        jx[i, ny+1] = jx[i, 1]
        jy[i, 1] += jy[i, ny+1]
        jy[i, ny+1] = jy[i, 1]
    end
    for j = 1:ny+1
        jx[1, j] += jx[nx+1, j]
        jx[nx+1, j] = jx[1, j]
        jy[1, j] += jy[nx+1, j]
        jy[nx+1, j] = jy[1, j]
    end

end


#=
   Fair coin from a biased coin.

   http://cplint.eu/e/von_neumann_trick.swinb
   """
   If you have a biased coin, how to obtain a fair coin, i.e., from a coin that lands heads with
   probability p with p≠0.5, how to generate samples from {heads,tails} with P(heads)=P(tails)=0.5?

   John von Neuamnn gave the following solution

    1) Toss the coin twice.
    2) If you get HT return H, if you get TH return T
    3) Otherwise, return to 1 and start over

   In fact, if p is the probability of the biased coin landing heads, then the outcomes HT and TH are
   equiprobable with probability p(1−p). However, with probability pp+(1−p)(1−p)
   these results are not obtained. In this case, we simply repeat the process. The probability that we
   never get HT or TH is 0 so this procedure ends with certainty.

   See
   - https://en.wikipedia.org/wiki/Fair_coin#Fair_results_from_a_biased_coin
   - von Neumann, John (1951). "Various techniques used in connection with random digits".
     National Bureau of Standards Applied Math Series. 12: 36.
   """

   Distributions of variable biased_coin1 (num:0)
   1.00000 =>    7992  (0.799200)
   2.00000 =>    2008  (0.200800)

   Distributions of variable coin1 (num:0)
   1.00000 =>    5004  (0.500400)
   2.00000 =>    4996  (0.499600)

   Cf ~/blog/fair_coin_from_a_biased_coin.blog
      ~/webppl/fair_coin_from_a_biased_coin.wppl
=#

using Turing, StatsPlots, DataFrames
include("jl_utils.jl")

@model function fair_coin_from_a_biased_coin()
    head = 1
    tail = 2

    n = 3
    function biased_coin(i)
        Categorical([0.8,0.2]) # [head,tail]
    end

    function coin(i)
        # StackOverflowError unless we limit the recursion.
        # But that don't help since we got a lot of 0's as the result....
        # This is fixed below...
        if i > 100
           return Dirac(0)
        end
        bi ~ biased_coin(i)
        bi1 ~ biased_coin(i+1)
        if bi != bi1
            Dirac(bi)
        else
            ci2 ~ coin(i+2)
            Dirac(ci2)
        end
    end

    biased_coin1 ~ biased_coin(1)
    coin1 ~ coin(1)
    # This is the fix. Not pretty though!
    true ~ Dirac(coin1 != 0)

end

model = fair_coin_from_a_biased_coin()

num_chains = 4

# chns = sample(model, Prior(), 10_000)
# chns = sample(model, MH(), 10_000)
chns = sample(model, PG(5), 10_000)
# chns = sample(model, SMC(), 10_000)
# chns = sample(model, IS(), 10_000)

display(chns)
# display(plot(chns))

show_var_dist_pct(chns, :biased_coin1)
show_var_dist_pct(chns, :coin1)


# ---
# jupyter:
#   jupytext:
#     formats: ipynb,jl:hydrogen
#     text_representation:
#       extension: .jl
#       format_name: hydrogen
#       format_version: '1.3'
#       jupytext_version: 1.11.2
#   kernelspec:
#     display_name: Julia 1.6.2
#     language: julia
#     name: julia-1.6
# ---

# %%
using Random
using Distributions
using StatsPlots
diststring(dist) = replace(string(dist), r"\{.*\}"=>"")

function sim(; dist_true=Gamma(20, 1/20), n = 100, L = 10^5)
    sample = Vector{Float64}(undef, n)
    AIC_normal0 = Vector{Float64}(undef, L)
    AIC_normal1 = similar(AIC_normal0)
    AIC_normal2 = similar(AIC_normal0)
    for i in 1:L
        rand!(dist_true, sample)
        normal0 = Normal(mean(dist_true), std(dist_true))
        normal1 = Normal(mean(sample), std(dist_true))
        normal2 = Normal(mean(sample), std(sample))
        AIC_normal0[i] = -2loglikelihood(normal0, sample)
        AIC_normal1[i] = -2loglikelihood(normal1, sample) + 1
        AIC_normal2[i] = -2loglikelihood(normal2, sample) + 2
    end
    AIC_normal0, AIC_normal1, AIC_normal2
end

function sim_and_plot(; dist_true=dist_true, n = 100, L = 10^5)
    AIC_normal0, AIC_normal1, AIC_normal2 = sim(; dist_true, n, L)

    Selected_from_012 = zeros(Int, 3)
    for i in 1:L
        Selected_from_012[argmin((AIC_normal0[i], AIC_normal1[i], AIC_normal2[i]))] += 1
    end
    group = ["normal0", "normal1", "normal2"]
    P1 = bar(group, Selected_from_012/L; label="", group, color=(1:3)')
    title!("$(diststring(dist_true)), n = $n"; titlefontsize=9)
    plot!(; ylim = (-0.01, 0.7))

    Selected_from_01 = zeros(Int, 2)
    for i in 1:L
        Selected_from_01[argmin((AIC_normal0[i], AIC_normal1[i]))] += 1
    end
    group = ["normal0", "normal1"]
    P2 = bar(group, Selected_from_01/L; label="", group, color=(1:2)')
    plot!(; ylim = (-0.01, 0.7))

    Selected_from_02 = zeros(Int, 2)
    for i in 1:L
        Selected_from_02[argmin((AIC_normal0[i], AIC_normal2[i]))] += 1
    end
    group = ["normal0", "normal2"]
    P3 = bar(group, Selected_from_02/L; label="", group, color=[1,3]')
    plot!(; ylim = (-0.01, 0.7))

    Selected_from_12 = zeros(Int, 2)
    for i in 1:L
        Selected_from_12[argmin((AIC_normal1[i], AIC_normal2[i]))] += 1
    end
    group = ["normal1", "normal2"]
    P4 = bar(group, Selected_from_12/L; label="", group, color=(2:3)')
    plot!(; ylim = (-0.01, 0.7))

    plot(P1, P2, P3, P4; size=(800, 300), layout=@layout[a{0.333w} b c d])
end

function sim_and_plot2(k; n = 100, L = 10^5)
    dist_true = Gamma(k, 1/k)
    normal0 = Normal(mean(dist_true), round(std(dist_true), digits=3))
    A = sim_and_plot(; n, L, dist_true)
    B = sim_and_plot(; n, L, dist_true = normal0)
    plot(A, B; size=(800, 600), layout=(2, 1))
end

# %%
sim_and_plot2(1)

# %%
sim_and_plot2(2)

# %%
sim_and_plot2(5)

# %%
sim_and_plot2(10)

# %%
sim_and_plot2(20)

# %% tags=[]
dist_true = Gamma(20, 1/20)
normal0 = Normal(mean(dist_true), std(dist_true))
plot(dist_true; label="$(diststring(dist_true))")
plot!(normal0; label="$(diststring(normal0))")
plot!(; legend=:outertop)

# %%

0021/normal dist estimation of gamma dist sample.jl


export proxL2!, normL2

"""
    proxL2!(x::Vector{T}, λ::Float64; kargs...) where T

proximal map for Tikhonov regularization.
"""
function proxL2!(x::T, λ::Float64; kargs...) where T<:AbstractArray
  x[:] .*= 1. / (1. + 2. *λ)#*x
end

"""
    normL2(x::Vector{T}, λ::Float64, kargs...)

returns the value of the L2-regularization term
"""
normL2(x::T, λ::Float64; kargs...) where T<:AbstractArray = λ*norm(x,2)^2


using KroneckerMatrices
using Base.Test

#square matrix nonsquare factors
Q,R = 3,4
S,T = 4,3
Aouter = rand(Q,R)
Ainner = ones(S,T)

#square matrix, square factors
Q,R = 3,3
S,T = 4,4
Bouter = rand(Q,R)
Binner = ones(S,T)

#tests for square Kronecker Products
for (outer,inner) in zip(Any[Aouter Bouter],Any[Ainner Binner])
    M = KronProdMat(outer,inner)
    Mfull = kronprod(outer,inner)
    @test size(M) == size(Mfull)
    K,L = size(M)
    for i = 1:K
        for j = 1:L
            @test M[i,j] == Mfull[i,j]
        end
    end
    @test full(M') == Mfull'
    @test rank(M) == rank(Mfull)
    @test det(M) == det(Mfull)

    fullsvdvals = sort(svdvals(Mfull),by = -)
    kronsvdvalstrunc = sort(full(svdvals(M)),by = -)
    kronsvdvals = zeros(eltype(M),length(fullsvdvals))
    kronsvdvals[1:length(kronsvdvalstrunc)] = kronsvdvalstrunc
    for (σ,τ) in zip(kronsvdvals,fullsvdvals)
        @test_approx_eq_eps(σ,τ,100*eps())
    end
end

#tests for square Kronecker products with square factors
B = KronProdMat(Bouter,Binner)
Bfull = kronprod(Bouter,Binner)

#U,Ufull = convert(Matrix{eltype(B)},eigvecs(B)),eigvecs(Bfull)
#Bsim,Bsimfull = U*U',Ufull*Ufull'
#N = length(B)
#for i = 1:N, j = 1:N
#    @test_approx_eq_eps(Bsim[i,j],Bsimfull[i,j],10*eps())
#end

#for now test sorted eigs, ordering is different for different situation in the full problem
#not sure what the right tolerance should be here (probably should try to compute)

#for (λ1,λ2) in zip(sort(vec(convert(Matrix{eltype(B)},eigvals(B))),by=real),sort(eigvals(Bfull),by=real))
#    @test_approx_eq_eps(λ1,λ2,10*eps())
#end


# Test tetrahedron delta function integration
using TetrahedronIntegration
using StaticArrays
using Test
using TetrahedronIntegration: _evaluate_piecewise_polynomial

@testset "delta" begin
    e0 = 0.4
    v0 = [0.0, 0.0, 1.0]
    L = 0.3
    de = L * v0[3] / 2
    @test delta_parallelepiped(e0, e0, v0, L) ≈ 1.0 / L
    @test delta_parallelepiped(e0 + de - 1e-3, e0, v0, L) ≈ 1.0 / L
    @test delta_parallelepiped(e0 - de + 1e-3, e0, v0, L) ≈ 1.0 / L
    @test delta_parallelepiped(e0 + de + 1e-3, e0, v0, L) ≈ 0
    @test delta_parallelepiped(e0 - de - 1e-3, e0, v0, L) ≈ 0

    # Test small v0
    @test delta_parallelepiped(e0, e0, [0.0, 0.0, 0.0], L) ≈ 0
end

@testset "delta_polynomial" begin
    e1234 = SVector(1.5, -0.2, 0.7, 0.3)
    polys = delta_tetrahedron_polynomial(e1234)
    for e in -0.3:0.1:1.6
        @test delta_tetrahedron(e, e1234) ≈ _evaluate_piecewise_polynomial(e, polys) atol=10*eps(eltype(e))
    end

    e1234 = SVector(1.0, -0.5, 1.0, -0.5)
    polys = delta_tetrahedron_polynomial(e1234)
    for e in -1.0:0.5:1.0
        @test delta_tetrahedron(e, e1234) ≈ _evaluate_piecewise_polynomial(e, polys)
    end
end